"""
FastAPI server for Qwen Image Layered model.
Compatible with Hugging Face Inference Endpoints custom container format.
"""
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional, List, Dict, Any
import uvicorn
import base64
import io
from PIL import Image

# Import our handler
from handler import EndpointHandler

app = FastAPI()

# Initialize handler on startup
handler = None

@app.on_event("startup")
async def startup_event():
    global handler
    print("Initializing model...")
    handler = EndpointHandler()
    print("Model ready!")

class InferenceRequest(BaseModel):
    inputs: Dict[str, Any]
    parameters: Optional[Dict[str, Any]] = None

class HealthResponse(BaseModel):
    status: str

@app.get("/health")
async def health() -> HealthResponse:
    return HealthResponse(status="ok")

@app.get("/")
async def root():
    return {"status": "Qwen Image Layered Endpoint Ready"}

@app.post("/")
async def predict(request: InferenceRequest):
    if handler is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    data = {
        "inputs": request.inputs,
        "parameters": request.parameters or {}
    }
    
    try:
        result = handler(data)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8080)