app/main.py

import os
import io
import json
import uuid
import zipfile
import traceback
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, Request, Form, HTTPException
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

from app.engine.app_planner import AppPlanner
from app.engine.model_recommender import ModelRecommender
from app.codegen.repo_generator import RepoGenerator
from app.validators.code_checker import CodeChecker

app = FastAPI(title="AutoApp Builder", version="1.0.0")

BASE_DIR = Path(__file__).resolve().parent
GENERATED_DIR = Path("/tmp/autoapp_generated")
GENERATED_DIR.mkdir(parents=True, exist_ok=True)

app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static")), name="static")
templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))

# In-memory store for generated projects (keyed by session id)
projects: dict = {}

planner = AppPlanner()
recommender = ModelRecommender()
generator = RepoGenerator()
checker = CodeChecker()


@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    examples = [
        {
            "title": "Image Classifier",
            "prompt": "Build a Gradio app that classifies images using a pretrained ResNet model. Users upload an image and get top-5 predictions with confidence bars.",
        },
        {
            "title": "AI Chatbot",
            "prompt": "Create a chatbot that uses a Hugging Face language model to have conversations. Include chat history, system prompt configuration, and a clear button.",
        },
        {
            "title": "Text Summarizer",
            "prompt": "Build an app that summarizes long text documents. Let users paste text or upload a .txt file, choose summary length (short/medium/long), and display the result.",
        },
        {
            "title": "Sentiment Dashboard",
            "prompt": "Create an interactive sentiment analysis tool. Users enter text and see sentiment scores (positive/negative/neutral) visualized with charts.",
        },
        {
            "title": "REST API Service",
            "prompt": "Build a Docker-based REST API that serves a text generation model with endpoints for completion, summarization, and translation. Include API docs.",
        },
        {
            "title": "Portfolio Site",
            "prompt": "Create a beautiful static portfolio website for a data scientist. Include sections for projects, skills, publications, and contact info with a dark theme.",
        },
    ]
    return templates.TemplateResponse(
        "home.html", {"request": request, "examples": examples}
    )


@app.post("/generate", response_class=HTMLResponse)
async def generate(
    request: Request,
    prompt: str = Form(...),
    sdk_preference: str = Form("auto"),
    model_size: str = Form("medium"),
    gpu_needed: bool = Form(False),
    features: str = Form(""),
):
    try:
        # Step 1: Plan the app
        plan = planner.analyze(prompt, sdk_preference)

        # Step 2: Recommend models
        recommended_models = recommender.recommend(plan, model_size, gpu_needed)
        plan["recommended_models"] = recommended_models

        # Step 3: Parse additional features
        feature_list = [f.strip() for f in features.split(",") if f.strip()]
        plan["extra_features"] = feature_list

        # Step 4: Generate repository files
        repo_files = generator.generate(plan, prompt)

        # Step 5: Validate the generated code
        validation = checker.check(repo_files, plan["sdk"])

        # Step 6: Store the project
        project_id = str(uuid.uuid4())[:8]
        projects[project_id] = {
            "plan": plan,
            "files": repo_files,
            "validation": validation,
            "prompt": prompt,
        }

        # Build file tree structure
        file_tree = _build_file_tree(repo_files)

        # Generate architecture diagram
        arch_diagram = _generate_arch_diagram(plan)

        return templates.TemplateResponse(
            "result.html",
            {
                "request": request,
                "project_id": project_id,
                "plan": plan,
                "files": repo_files,
                "file_tree": file_tree,
                "validation": validation,
                "arch_diagram": arch_diagram,
                "prompt": prompt,
            },
        )
    except Exception as e:
        traceback.print_exc()
        return templates.TemplateResponse(
            "home.html",
            {
                "request": request,
                "examples": [],
                "error": f"Generation failed: {str(e)}. Please try again with a different prompt.",
            },
        )


@app.post("/edit", response_class=HTMLResponse)
async def edit_project(
    request: Request,
    project_id: str = Form(...),
    edit_prompt: str = Form(...),
):
    if project_id not in projects:
        raise HTTPException(status_code=404, detail="Project not found")

    project = projects[project_id]
    original_plan = project["plan"]
    original_files = project["files"]

    try:
        # Re-generate with edit instructions
        updated_files = generator.edit(original_plan, original_files, edit_prompt)

        validation = checker.check(updated_files, original_plan["sdk"])

        project["files"] = updated_files
        project["validation"] = validation

        file_tree = _build_file_tree(updated_files)
        arch_diagram = _generate_arch_diagram(original_plan)

        return templates.TemplateResponse(
            "result.html",
            {
                "request": request,
                "project_id": project_id,
                "plan": original_plan,
                "files": updated_files,
                "file_tree": file_tree,
                "validation": validation,
                "arch_diagram": arch_diagram,
                "prompt": project["prompt"],
                "edit_prompt": edit_prompt,
            },
        )
    except Exception as e:
        traceback.print_exc()
        # Return original project with error
        file_tree = _build_file_tree(original_files)
        arch_diagram = _generate_arch_diagram(original_plan)
        return templates.TemplateResponse(
            "result.html",
            {
                "request": request,
                "project_id": project_id,
                "plan": original_plan,
                "files": original_files,
                "file_tree": file_tree,
                "validation": project["validation"],
                "arch_diagram": arch_diagram,
                "prompt": project["prompt"],
                "edit_error": f"Edit failed: {str(e)}",
            },
        )


@app.get("/download/{project_id}")
async def download_zip(project_id: str):
    if project_id not in projects:
        raise HTTPException(status_code=404, detail="Project not found")

    project = projects[project_id]
    files = project["files"]
    plan = project["plan"]
    app_name = plan.get("app_name", "my-hf-space")

    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        for filepath, content in files.items():
            zf.writestr(f"{app_name}/{filepath}", content)

    buf.seek(0)
    return StreamingResponse(
        buf,
        media_type="application/zip",
        headers={"Content-Disposition": f'attachment; filename="{app_name}.zip"'},
    )


@app.get("/api/file/{project_id}/{filepath:path}")
async def get_file(project_id: str, filepath: str):
    if project_id not in projects:
        raise HTTPException(status_code=404, detail="Project not found")

    files = projects[project_id]["files"]
    if filepath not in files:
        raise HTTPException(status_code=404, detail="File not found")

    return JSONResponse({"filename": filepath, "content": files[filepath]})


def _build_file_tree(files: dict) -> list:
    """Build a nested file tree structure from flat file dict."""
    tree = []
    dirs_seen = set()
    sorted_files = sorted(files.keys())

    for filepath in sorted_files:
        parts = filepath.split("/")
        # Add directory entries
        for i in range(len(parts) - 1):
            dir_path = "/".join(parts[: i + 1])
            if dir_path not in dirs_seen:
                dirs_seen.add(dir_path)
                tree.append(
                    {
                        "path": dir_path,
                        "name": parts[i],
                        "type": "dir",
                        "depth": i,
                    }
                )
        # Add file entry
        tree.append(
            {
                "path": filepath,
                "name": parts[-1],
                "type": "file",
                "depth": len(parts) - 1,
            }
        )

    return tree


def _generate_arch_diagram(plan: dict) -> str:
    """Generate ASCII architecture diagram."""
    sdk = plan.get("sdk", "gradio")
    app_name = plan.get("app_name", "App")
    components = plan.get("components", [])

    if sdk == "gradio":
        diagram = f"""
+------------------------------------------------------+
|                  Hugging Face Spaces                  |
|  +------------------------------------------------+  |
|  |              {app_name:^30s}          |  |
|  |  +------------------------------------------+  |  |
|  |  |            Gradio Interface               |  |  |
|  |  |  +------------+    +------------------+   |  |  |
|  |  |  |   Inputs   |--->|   Processing     |   |  |  |
|  |  |  +------------+    |  +------------+  |   |  |  |
|  |  |                    |  |  HF Model  |  |   |  |  |
|  |  |  +------------+    |  +------------+  |   |  |  |
|  |  |  |  Outputs   |<---|                  |   |  |  |
|  |  |  +------------+    +------------------+   |  |  |
|  |  +------------------------------------------+  |  |
|  +------------------------------------------------+  |
+------------------------------------------------------+"""
    elif sdk == "docker":
        diagram = f"""
+------------------------------------------------------+
|                  Hugging Face Spaces                  |
|  +------------------------------------------------+  |
|  |            Docker Container                     |  |
|  |  +------------------------------------------+  |  |
|  |  |         {app_name:^30s}      |  |  |
|  |  |  +----------+  +----------+  +---------+ |  |  |
|  |  |  |  FastAPI  |  |  Model   |  |  Utils  | |  |  |
|  |  |  |  Routes   |  |  Service |  |         | |  |  |
|  |  |  +-----+-----+  +----+-----+  +---------+ |  |  |
|  |  |        |              |                    |  |  |
|  |  |  +-----v--------------v-----------------+ |  |  |
|  |  |  |        API Endpoints                  | |  |  |
|  |  |  +---------------------------------------+ |  |  |
|  |  +------------------------------------------+  |  |
|  +------------------------------------------------+  |
+------------------------------------------------------+"""
    else:
        diagram = f"""
+------------------------------------------------------+
|                  Hugging Face Spaces                  |
|  +------------------------------------------------+  |
|  |            Static Site                          |  |
|  |  +------------------------------------------+  |  |
|  |  |         {app_name:^30s}      |  |  |
|  |  |  +----------+  +----------+  +---------+ |  |  |
|  |  |  |  HTML     |  |  CSS     |  |   JS   | |  |  |
|  |  |  |  Pages    |  |  Styles  |  | Scripts | |  |  |
|  |  |  +----------+  +----------+  +---------+ |  |  |
|  |  +------------------------------------------+  |  |
|  +------------------------------------------------+  |
+------------------------------------------------------+"""

    return diagram