app/codegen/readme_generator.py

"""
Generate beautiful README.md files for generated HF Spaces.
"""

from typing import Optional


EMOJI_MAP = {
    "chatbot": "💬",
    "image_classifier": "🖼️",
    "text_summarizer": "📝",
    "sentiment_analyzer": "😊",
    "text_generator": "✍️",
    "translator": "🌐",
    "object_detector": "🔍",
    "speech_to_text": "🎤",
    "image_generator": "🎨",
    "question_answering": "❓",
    "rest_api": "🔌",
    "portfolio": "💼",
}

COLOR_MAP = {
    "chatbot": ("blue", "purple"),
    "image_classifier": ("green", "yellow"),
    "text_summarizer": ("indigo", "blue"),
    "sentiment_analyzer": ("pink", "red"),
    "text_generator": ("purple", "indigo"),
    "translator": ("cyan", "blue"),
    "object_detector": ("orange", "red"),
    "speech_to_text": ("green", "cyan"),
    "image_generator": ("purple", "pink"),
    "question_answering": ("yellow", "orange"),
    "rest_api": ("gray", "blue"),
    "portfolio": ("blue", "cyan"),
}


class ReadmeGenerator:
    """Generate HF Space README.md with proper frontmatter."""

    def generate(self, plan: dict, sdk: str) -> str:
        """Generate a complete README.md for the Space."""
        app_type = plan.get("app_type", "custom")
        title = plan.get("title", "My App")
        description = plan.get("description", "A Hugging Face Space")
        app_name = plan.get("app_name", "my-app")
        models = plan.get("recommended_models", [])
        components = plan.get("components", [])

        emoji = EMOJI_MAP.get(app_type, "🚀")
        color_from, color_to = COLOR_MAP.get(app_type, ("blue", "cyan"))

        # Build frontmatter
        frontmatter = self._build_frontmatter(
            title=title,
            emoji=emoji,
            color_from=color_from,
            color_to=color_to,
            sdk=sdk,
            models=models,
        )

        # Build body
        body = self._build_body(
            title=title,
            description=description,
            app_type=app_type,
            sdk=sdk,
            models=models,
            components=components,
            plan=plan,
        )

        return frontmatter + "\n" + body

    def _build_frontmatter(
        self,
        title: str,
        emoji: str,
        color_from: str,
        color_to: str,
        sdk: str,
        models: list,
    ) -> str:
        lines = [
            "---",
            f'title: "{title}"',
            f"emoji: {emoji}",
            f"colorFrom: {color_from}",
            f"colorTo: {color_to}",
        ]

        if sdk == "gradio":
            lines.append("sdk: gradio")
            lines.append("sdk_version: 5.9.1")
        elif sdk == "docker":
            lines.append("sdk: docker")
            lines.append("app_port: 7860")
        else:
            lines.append("sdk: static")

        if models:
            lines.append("models:")
            for m in models[:3]:
                lines.append(f'  - "{m["id"]}"')

        lines.append("pinned: false")
        lines.append("license: mit")
        lines.append("---")

        return "\n".join(lines)

    def _build_body(
        self,
        title: str,
        description: str,
        app_type: str,
        sdk: str,
        models: list,
        components: list,
        plan: dict,
    ) -> str:
        sections = []

        # Title and description
        sections.append(f"# {title}\n")
        sections.append(f"{description}\n")

        # Features section
        sections.append("## Features\n")
        feature_descriptions = {
            "chat_interface": "Interactive chat interface with streaming responses",
            "system_prompt_config": "Configurable system prompt",
            "clear_button": "Clear conversation history",
            "history": "Persistent chat history",
            "image_upload": "Image upload support",
            "label_output": "Classification labels with confidence scores",
            "confidence_bars": "Visual confidence bars",
            "examples": "Pre-loaded example inputs",
            "text_input": "Text input area",
            "file_upload": "File upload support",
            "length_selector": "Adjustable output length",
            "text_output": "Text output display",
            "chart_output": "Data visualization charts",
            "parameter_controls": "Adjustable model parameters",
            "language_selector": "Multi-language selection",
            "audio_input": "Audio/microphone input",
            "annotated_image_output": "Annotated image output with bounding boxes",
            "json_output": "Structured JSON output",
            "gallery": "Image gallery view",
            "image_output": "Image output display",
            "fastapi_app": "FastAPI web application",
            "model_endpoint": "Model inference endpoint",
            "docs": "Interactive API documentation",
            "health_check": "Health check endpoint",
            "html_page": "Responsive HTML pages",
            "css_styles": "Custom CSS styling",
            "js_scripts": "Interactive JavaScript",
            "context_input": "Context paragraph input",
            "question_input": "Question input field",
            "answer_output": "Answer display with confidence",
            "highlight": "Answer highlighting in context",
            "file_component": "File handling component",
            "video_input": "Video input support",
            "hero_section": "Hero section with CTA",
            "projects_section": "Projects showcase",
            "skills_section": "Skills display",
            "contact_form": "Contact form",
        }
        for comp in components:
            desc = feature_descriptions.get(comp, comp.replace("_", " ").title())
            sections.append(f"- {desc}")
        sections.append("")

        # Models section
        if models:
            sections.append("## Models Used\n")
            for m in models:
                sections.append(f"- **[{m['id']}](https://huggingface.co/{m['id']})** - {m.get('desc', '')}")
            sections.append("")

        # Tech stack
        sections.append("## Tech Stack\n")
        if sdk == "gradio":
            sections.append("- [Gradio](https://gradio.app/) - UI framework")
            sections.append("- [Hugging Face Hub](https://huggingface.co/) - Model inference")
        elif sdk == "docker":
            sections.append("- [FastAPI](https://fastapi.tiangolo.com/) - Web framework")
            sections.append("- [Docker](https://docker.com/) - Containerization")
            sections.append("- [Hugging Face Hub](https://huggingface.co/) - Model inference")
        else:
            sections.append("- HTML5 / CSS3 / JavaScript")
        sections.append("")

        # Usage
        sections.append("## Usage\n")
        if sdk == "gradio":
            sections.append("1. Open the Space URL")
            sections.append("2. Interact with the interface")
            sections.append("3. Results will be displayed automatically")
        elif sdk == "docker":
            sections.append("### API Endpoints\n")
            sections.append("Visit `/docs` for interactive API documentation.\n")
            sections.append("```bash")
            sections.append('curl -X POST "/query" -H "Content-Type: application/json" -d \'{"query": "Hello"}\'')
            sections.append("```")
        else:
            sections.append("Simply visit the Space URL to view the site.")
        sections.append("")

        # Footer
        sections.append("---\n")
        sections.append("*Generated by [AutoApp Builder](https://huggingface.co/spaces/autoapp-builder)*")

        return "\n".join(sections)