from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
from comfy_launcher import launch_comfy_server

app = FastAPI(title="RunComfy Launcher API")


class LaunchRequest(BaseModel):
    version_id: Optional[str] = "90f77137-ba75-400d-870f-204c614ae8a3"
    server_type: Optional[str] = "medium"
    duration: Optional[int] = 3600


class LaunchResponse(BaseModel):
    server_id: str
    comfy_url: str
    status: str
    usage_instruction: str


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


@app.post("/launch", response_model=LaunchResponse)
def launch(req: LaunchRequest):
    try:
        result = launch_comfy_server(
            version_id=req.version_id,
            server_type=req.server_type,
            duration=req.duration
        )
        return result
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except TimeoutError as e:
        raise HTTPException(status_code=504, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"启动失败: {str(e)}")


if __name__ == "__main__":
    import argparse
    import uvicorn
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8002)
    args = parser.parse_args()
    uvicorn.run(app, host="0.0.0.0", port=args.port)
