from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
from io import BytesIO
from PIL import Image
import base64
from models.model_loader import get_model, MODELS

app = FastAPI(title="Text-to-Image API", description="Generate images from text prompts using Stable Diffusion")

# Mount static files
app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
async def read_index():
    return FileResponse('static/index.html')

class GenerationRequest(BaseModel):
    prompt: str
    width: int = 512
    height: int = 512
    model_key: str = "sd15"

class GenerationResponse(BaseModel):
    image_base64 : str

@app.get("/models")
def list_models():
    return{
        "available_models":[
            {"key": key, "name": info["name"], "description": info["description"]}
            for key, info in MODELS.items()
        ]
    }
@app.post("/generate", response_model = GenerationResponse)
def generate_image(request: GenerationRequest):
    try:
        model_key = request.model_key if request.model_key in MODELS else "sd15"
        model = get_model(model_key) # type: ignore
        
        # Check if model is a function (fallback) or actual model
        if callable(model) and hasattr(model, '__name__') and model.__name__ == 'create_placeholder_image':
            # Use fallback placeholder generation
            result = model(request.prompt, request.height, request.width)
        else:
            # Use actual model
            result = model(request.prompt, height=request.height, width=request.width).images[0]
        
        buffered = BytesIO()
        result.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        
        return GenerationResponse(image_base64=img_str)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Image generation failed: {e}")