import os
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from typing import Optional
from comfy_launcher import launch_comfy_server
from dotenv import load_dotenv

env_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), ".env")
if os.path.exists(env_path):
    load_dotenv(env_path)
else:
    load_dotenv()

app = FastAPI()

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

@app.post("/launch")
def launch(req: LaunchRequest):
    return launch_comfy_server(
        version_id=req.version_id,
        server_type=req.server_type,
        duration=req.duration
    )

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