app/engine/app_planner.py

"""
Analyze user prompts, select the appropriate SDK, and plan app architecture.
"""

import re
from typing import Optional


# Keywords that strongly indicate a particular SDK
GRADIO_KEYWORDS = [
    "gradio", "demo", "interface", "chatbot", "chat", "image classifier",
    "image classification", "object detection", "text generation", "summarizer",
    "summarization", "sentiment", "translation", "translate", "speech",
    "audio", "video", "upload", "webcam", "microphone", "slider", "dropdown",
    "interactive", "ml model", "machine learning", "deep learning", "neural",
    "predict", "inference", "classify", "segment", "detect", "recognize",
    "generate image", "stable diffusion", "text-to-image", "image-to-text",
    "question answering", "qa", "ner", "named entity", "token classification",
    "zero-shot", "few-shot", "embedding", "similarity",
]

DOCKER_KEYWORDS = [
    "docker", "api", "rest", "fastapi", "flask", "endpoint", "websocket",
    "agent", "autonomous", "pipeline", "workflow", "backend", "server",
    "database", "redis", "celery", "queue", "cron", "schedule", "scrape",
    "crawl", "multi-step", "orchestrat", "microservice", "webhook",
    "authentication", "auth", "oauth", "custom server", "complex",
    "real-time", "streaming api",
]

STATIC_KEYWORDS = [
    "static", "landing page", "portfolio", "documentation", "docs",
    "blog", "website", "html", "css only", "showcase", "gallery",
    "resume", "cv", "personal site", "no backend", "frontend only",
]

# Templates for common app types
APP_TEMPLATES = {
    "chatbot": {
        "components": ["chat_interface", "system_prompt_config", "clear_button", "history"],
        "model_task": "text-generation",
        "description": "Conversational AI chatbot",
    },
    "image_classifier": {
        "components": ["image_upload", "label_output", "confidence_bars", "examples"],
        "model_task": "image-classification",
        "description": "Image classification with confidence scores",
    },
    "text_summarizer": {
        "components": ["text_input", "file_upload", "length_selector", "text_output"],
        "model_task": "summarization",
        "description": "Text summarization tool",
    },
    "sentiment_analyzer": {
        "components": ["text_input", "label_output", "chart_output"],
        "model_task": "text-classification",
        "description": "Sentiment analysis with visualization",
    },
    "text_generator": {
        "components": ["text_input", "parameter_controls", "text_output"],
        "model_task": "text-generation",
        "description": "Text generation with configurable parameters",
    },
    "translator": {
        "components": ["text_input", "language_selector", "text_output"],
        "model_task": "translation",
        "description": "Multi-language translation tool",
    },
    "object_detector": {
        "components": ["image_upload", "annotated_image_output", "json_output"],
        "model_task": "object-detection",
        "description": "Object detection with bounding boxes",
    },
    "speech_to_text": {
        "components": ["audio_input", "text_output", "language_selector"],
        "model_task": "automatic-speech-recognition",
        "description": "Speech recognition / transcription",
    },
    "image_generator": {
        "components": ["text_input", "parameter_controls", "image_output", "gallery"],
        "model_task": "text-to-image",
        "description": "AI image generation from text prompts",
    },
    "question_answering": {
        "components": ["context_input", "question_input", "answer_output", "highlight"],
        "model_task": "question-answering",
        "description": "Extractive question answering",
    },
    "rest_api": {
        "components": ["fastapi_app", "model_endpoint", "docs", "health_check"],
        "model_task": "text-generation",
        "description": "REST API serving ML models",
    },
    "portfolio": {
        "components": ["hero_section", "projects_section", "skills_section", "contact_form"],
        "model_task": None,
        "description": "Personal portfolio website",
    },
}


class AppPlanner:
    """Analyzes user prompts and produces a structured app plan."""

    def analyze(self, prompt: str, sdk_preference: str = "auto") -> dict:
        """
        Analyze the user prompt and return a structured plan.

        Returns dict with keys:
            sdk, app_name, app_type, description, components,
            model_task, template_key, title
        """
        prompt_lower = prompt.lower()

        # Detect the best matching template
        template_key = self._match_template(prompt_lower)
        template = APP_TEMPLATES.get(template_key, {})

        # Determine SDK
        if sdk_preference != "auto":
            sdk = sdk_preference
        else:
            sdk = self._select_sdk(prompt_lower, template_key)

        # Generate an app name from the prompt
        app_name = self._generate_app_name(prompt_lower)

        # Determine a human-readable title
        title = self._generate_title(prompt, template_key)

        plan = {
            "sdk": sdk,
            "app_name": app_name,
            "app_type": template_key or "custom",
            "title": title,
            "description": template.get("description", self._extract_description(prompt)),
            "components": template.get("components", self._infer_components(prompt_lower, sdk)),
            "model_task": template.get("model_task", self._infer_model_task(prompt_lower)),
            "template_key": template_key,
            "original_prompt": prompt,
        }

        return plan

    def _select_sdk(self, prompt_lower: str, template_key: Optional[str]) -> str:
        """Auto-select the best SDK based on prompt analysis."""
        # Check template-based hints
        if template_key in ("rest_api",):
            return "docker"
        if template_key in ("portfolio",):
            return "static"

        gradio_score = sum(1 for kw in GRADIO_KEYWORDS if kw in prompt_lower)
        docker_score = sum(1 for kw in DOCKER_KEYWORDS if kw in prompt_lower)
        static_score = sum(1 for kw in STATIC_KEYWORDS if kw in prompt_lower)

        scores = {"gradio": gradio_score, "docker": docker_score, "static": static_score}
        best = max(scores, key=scores.get)

        # Default to gradio if no strong signal
        if scores[best] == 0:
            return "gradio"

        return best

    def _match_template(self, prompt_lower: str) -> Optional[str]:
        """Find the best matching template for the prompt."""
        scores = {}
        keyword_map = {
            "chatbot": ["chat", "chatbot", "conversation", "dialogue", "talk"],
            "image_classifier": ["classify image", "image classif", "image recogni", "predict image"],
            "text_summarizer": ["summar", "summarize", "summarization", "tldr", "shorten text"],
            "sentiment_analyzer": ["sentiment", "emotion", "mood", "opinion"],
            "text_generator": ["text generat", "write text", "generate text", "story generat", "complete text"],
            "translator": ["translat", "language translation"],
            "object_detector": ["object detect", "detect object", "bounding box", "find objects"],
            "speech_to_text": ["speech", "transcri", "voice to text", "audio to text", "asr"],
            "image_generator": ["generate image", "text to image", "stable diffusion", "create image", "dall-e", "image generat"],
            "question_answering": ["question answer", "qa ", "answer question", "reading comprehension"],
            "rest_api": ["rest api", "api service", "endpoint", "api server", "fastapi"],
            "portfolio": ["portfolio", "personal site", "resume", "cv site", "showcase"],
        }

        for key, keywords in keyword_map.items():
            score = sum(1 for kw in keywords if kw in prompt_lower)
            if score > 0:
                scores[key] = score

        if not scores:
            return None

        return max(scores, key=scores.get)

    def _generate_app_name(self, prompt_lower: str) -> str:
        """Generate a slug-style app name."""
        # Try to extract key nouns
        words = re.sub(r"[^a-z0-9\s]", "", prompt_lower).split()
        stop_words = {
            "a", "an", "the", "is", "are", "was", "were", "be", "been",
            "being", "have", "has", "had", "do", "does", "did", "will",
            "would", "could", "should", "may", "might", "can", "shall",
            "build", "create", "make", "develop", "design", "generate",
            "that", "which", "who", "whom", "this", "these", "those",
            "i", "me", "my", "we", "our", "you", "your", "it", "its",
            "they", "them", "their", "what", "where", "when", "how",
            "and", "or", "but", "if", "then", "else", "so", "for",
            "with", "from", "to", "of", "in", "on", "at", "by", "about",
            "into", "through", "during", "before", "after", "above",
            "below", "between", "under", "using", "use", "uses",
            "let", "users", "user", "app", "application",
        }
        meaningful = [w for w in words if w not in stop_words and len(w) > 2][:4]
        if not meaningful:
            meaningful = ["my", "app"]
        return "-".join(meaningful)

    def _generate_title(self, prompt: str, template_key: Optional[str]) -> str:
        """Generate a human-friendly title."""
        if template_key and template_key in APP_TEMPLATES:
            return APP_TEMPLATES[template_key]["description"]
        # Capitalize first ~6 words of prompt
        words = prompt.split()[:6]
        title = " ".join(words)
        if len(prompt.split()) > 6:
            title += "..."
        return title

    def _extract_description(self, prompt: str) -> str:
        """Extract a short description from the prompt."""
        sentences = prompt.split(".")
        if sentences:
            return sentences[0].strip()[:200]
        return prompt[:200]

    def _infer_components(self, prompt_lower: str, sdk: str) -> list:
        """Infer UI components from the prompt when no template matches."""
        components = []
        if sdk == "gradio":
            if any(w in prompt_lower for w in ["image", "photo", "picture", "upload image"]):
                components.append("image_upload")
            if any(w in prompt_lower for w in ["text", "input", "enter", "paste", "type"]):
                components.append("text_input")
            if any(w in prompt_lower for w in ["audio", "sound", "voice", "microphone"]):
                components.append("audio_input")
            if any(w in prompt_lower for w in ["video", "webcam"]):
                components.append("video_input")
            if any(w in prompt_lower for w in ["chart", "plot", "graph", "visual"]):
                components.append("chart_output")
            if any(w in prompt_lower for w in ["file", "upload", "download"]):
                components.append("file_component")
            if not components:
                components = ["text_input", "text_output"]
        elif sdk == "docker":
            components = ["fastapi_app", "model_endpoint", "health_check"]
        else:
            components = ["html_page", "css_styles", "js_scripts"]
        return components

    def _infer_model_task(self, prompt_lower: str) -> Optional[str]:
        """Infer the HF model task from prompt keywords."""
        task_keywords = {
            "text-generation": ["generat", "write", "complete", "chat", "story"],
            "text-classification": ["classif", "sentiment", "emotion", "categor"],
            "summarization": ["summar", "tldr", "shorten"],
            "translation": ["translat"],
            "image-classification": ["image classif", "recognize image", "identify image"],
            "object-detection": ["object detect", "bounding box", "detect object"],
            "text-to-image": ["generate image", "text to image", "create image"],
            "automatic-speech-recognition": ["speech", "transcri", "asr", "voice to text"],
            "question-answering": ["question answer", "qa"],
            "token-classification": ["ner", "named entity", "token classif"],
            "fill-mask": ["fill mask", "mask", "cloze"],
        }

        for task, keywords in task_keywords.items():
            if any(kw in prompt_lower for kw in keywords):
                return task

        return "text-generation"  # sensible default