app.py

"""
Production FastAPI Backend for Image2Code
Wraps the VLM agent into a REST API for HuggingFace Spaces deployment
"""
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn
import tempfile
import os
from pathlib import Path
from typing import Optional
from uuid import uuid4
import sys
import shutil

# Import agent components
from agent_azure_vlm_tools import ModelManager, azure_vlm_app, CodeRefineState

app = FastAPI(
    title="Image2Code VLM API",
    description="Backend API for Image2Code using Qwen2.5-VL-7B-Instruct",
    version="1.0.0"
)

# CORS configuration for production
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:5173",  # Local dev
        "http://localhost:3000",
        "https://*.vercel.app",   # Vercel deployments
        # Add your specific production domain here
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global model manager (singleton)
models = None


@app.on_event("startup")
async def startup_event():
    """Load models on server startup"""
    global models
    print("🚀 Starting Image2Code VLM API...")
    print("📦 Loading Qwen2.5-VL-7B-Instruct model... This may take 2-3 minutes.")
    models = ModelManager()
    print("✅ Models loaded successfully!")


@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "service": "Image2Code VLM API",
        "model": "Qwen2.5-VL-7B-Instruct",
        "version": "1.0.0",
        "ready": models is not None
    }


@app.get("/health")
async def health_check():
    """Detailed health check"""
    return {
        "status": "healthy" if models is not None else "starting",
        "models_loaded": models is not None,
        "gpu_available": True,  # Update based on actual check
    }


@app.post("/api/generate")
async def generate_ui(
    prompt: str = Form(..., description="Text prompt describing the UI to generate"),
    image: UploadFile = File(..., description="Wireframe/screenshot image"),
):
    """
    Generate UI code from a prompt and wireframe image using Qwen VLM
    
    Args:
        prompt: Description of the desired UI
        image: Uploaded image file (PNG, JPG, etc.)
    
    Returns:
        JSON response with generated code, plan, and reasoning
    """
    if not models:
        raise HTTPException(
            status_code=503,
            detail="Models are still loading. Please wait a moment and try again."
        )
    
    # Validate image type
    if not image.content_type.startswith("image/"):
        raise HTTPException(
            status_code=400,
            detail=f"Invalid file type: {image.content_type}. Please upload an image."
        )
    
    # Create temporary directory for this request
    temp_dir = tempfile.mkdtemp()
    
    try:
        # Save uploaded image
        image_path = os.path.join(temp_dir, "input_wireframe.png")
        contents = await image.read()
        with open(image_path, "wb") as f:
            f.write(contents)
        
        # Define output paths
        out_html = os.path.join(temp_dir, "generated_output.html")
        out_brief = os.path.join(temp_dir, "generated_brief.txt")
        out_reldesc = os.path.join(temp_dir, "generated_reldesc.txt")
        
        # Build state for agent
        state = CodeRefineState(
            image_path=image_path,
            out_rel_desc=out_reldesc,
            out_brief=out_brief,
            out_html=out_html,
            vision_deployment="gpt-4.1-mini",
            text_deployment="gpt-4.1-mini",
            reldesc_tokens=700,
            brief_tokens=1100,
            code_tokens=2200,
            judge_tokens=900,
            temp=0.12,
            refine_max_iters=3,
            refine_threshold=8,
            shot_width=1536,
            shot_height=900
        )
        
        # Run the agent
        run_id = f"api-{uuid4()}"
        config = {"configurable": {"thread_id": run_id}}
        
        print(f"🎨 Processing wireframe with run_id: {run_id}")
        agent_result = azure_vlm_app.invoke(state, config=config)
        print(f"✅ Agent completed for run_id: {run_id}")
        
        # Read generated code
        generated_code = ""
        if os.path.exists(out_html):
            with open(out_html, 'r', encoding='utf-8') as f:
                generated_code = f.read()
        
        # Build response messages
        messages = []
        
        messages.append({
            "id": f"msg-{uuid4()}",
            "role": "assistant",
            "variant": "accent",
            "content": f"**Plan:** Generated UI from your wireframe using Qwen2.5-VL-7B-Instruct"
        })
        
        messages.append({
            "id": f"msg-{uuid4()}",
            "role": "assistant",
            "variant": "subtle",
            "content": """**Process:**
• Loaded and analyzed wireframe structure
• Identified UI components and layout
• Generated semantic HTML/CSS code
• Applied multi-stage refinement pipeline"""
        })
        
        if generated_code:
            messages.append({
                "id": f"msg-{uuid4()}",
                "role": "assistant",
                "variant": "accent",
                "content": f"```html\n{generated_code}\n```"
            })
        else:
            messages.append({
                "id": f"msg-{uuid4()}",
                "role": "assistant",
                "variant": "subtle",
                "content": "⚠️ Code generation completed but output file was empty."
            })
        
        return JSONResponse(content={
            "messages": messages,
            "status": {
                "kind": "success",
                "text": "UI code generated successfully",
                "detail": f"Run ID: {run_id}"
            },
            "usedFallback": False
        })
        
    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        print(f"❌ Error processing request: {error_details}")
        
        raise HTTPException(
            status_code=500,
            detail=f"Error generating UI: {str(e)}"
        )
    
    finally:
        # Cleanup temporary files
        try:
            shutil.rmtree(temp_dir)
        except Exception as e:
            print(f"⚠️ Failed to cleanup temp dir: {e}")


@app.post("/api/chat")
async def chat_with_vlm(
    prompt: str = Form(..., description="Chat message/question"),
    image: Optional[UploadFile] = File(None, description="Optional image for vision tasks"),
):
    """
    Simple chat endpoint for VLM queries
    """
    if not models:
        raise HTTPException(
            status_code=503,
            detail="Models not loaded yet"
        )
    
    return JSONResponse(content={
        "response": "Chat endpoint is available but not fully implemented yet.",
        "status": "success"
    })


if __name__ == "__main__":
    # For local testing
    port = int(os.environ.get("PORT", 7860))  # HF Spaces uses 7860
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=port,
        timeout_keep_alive=300  # 5 minutes for long-running requests
    )