from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import io
from model_utils import predict, download_model_if_not_exists

app = FastAPI(
    title="Face Shape Classification API",
    description="API for classifying face shapes using the fahd9999/face_shape_classification model.",
    version="1.0.0"
)

# Enable CORS for Vercel frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins for simplicity in this context, or specify Vercel domain
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.on_event("startup")
async def startup_event():
    """Ensure model is downloaded on startup."""
    try:
        download_model_if_not_exists()
    except Exception as e:
        print(f"Startup warning: Could not download model: {e}")

@app.get("/")
async def root():
    return {"message": "Face Shape Classification API is running"}

@app.post("/predict")
async def predict_endpoint(file: UploadFile = File(...)):
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image")
    
    try:
        contents = await file.read()
        image_stream = io.BytesIO(contents)
        result = predict(image_stream)
        return JSONResponse(content=result)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

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