from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
import uvicorn
import sys
from liblibai_client import LibLibAIClient

app = FastAPI(title="LibLib ControlNet API")

class GenerateRequest(BaseModel):
    image: str
    prompt: str
    negative_prompt: str = "lowres, bad anatomy, text, error"
    width: int = 512
    height: int = 512
    steps: int = 20
    cfg_scale: float = 7
    img_count: int = 1
    control_weight: float = 1.0
    preprocessor: int = 1
    canny_low: int = 100
    canny_high: int = 200

@app.post("/generate")
async def generate(req: GenerateRequest):
    try:
        client = LibLibAIClient()
        result = client.generate_image(
            image=req.image,
            prompt=req.prompt,
            negative_prompt=req.negative_prompt,
            width=req.width,
            height=req.height,
            steps=req.steps,
            cfg_scale=req.cfg_scale,
            img_count=req.img_count,
            control_weight=req.control_weight,
            preprocessor=req.preprocessor,
            canny_low=req.canny_low,
            canny_high=req.canny_high
        )
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health():
    return {"status": "ok"}

if __name__ == "__main__":
    port = int(sys.argv[sys.argv.index("--port") + 1]) if "--port" in sys.argv else 8001
    uvicorn.run(app, host="0.0.0.0", port=port)