from fastapi import FastAPI
from pydantic import BaseModel
from typing import Optional
import httpx, json, base64, asyncio, time, random, os

app = FastAPI(title="ComfyUI OpenAI Adapter", version="0.1.0")

COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://10.0.30.97:8188")
WORKFLOW_PREFIX = os.environ.get("WORKFLOW_PREFIX", "oai_adapter")

# FLUX.2 workflow using Qwen3VL, EmptyFlux2LatentImage, and TAEF2
# Adapt model names to your local ComfyUI installation.
FLUX2_WORKFLOW = {
    "1": {"inputs": {"unet_name": "flux-2-klein-9b-Q4_K_S.gguf"}, "class_type": "UnetLoaderGGUF"},
    "2": {"inputs": {"clip_name": "Qwen3VL-8B-Instruct-Q4_K_M.gguf", "type": "flux2"}, "class_type": "CLIPLoaderGGUF"},
    "3": {"inputs": {"vae_name": "taef2"}, "class_type": "VAELoader"},
    "4": {"inputs": {"width": 1024, "height": 1024, "batch_size": 1}, "class_type": "EmptyFlux2LatentImage"},
    "5": {"inputs": {"text": "prompt", "clip": ["2", 0]}, "class_type": "CLIPTextEncode"},
    "6": {"inputs": {"text": "", "clip": ["2", 0]}, "class_type": "CLIPTextEncode"},
    "7": {"inputs": {"seed": 42, "steps": 4, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple",
                      "denoise": 1.0, "model": ["1", 0], "positive": ["5", 0], "negative": ["6", 0],
                      "latent_image": ["4", 0]}, "class_type": "KSampler"},
    "8": {"inputs": {"samples": ["7", 0], "vae": ["3", 0]}, "class_type": "VAEDecode"},
    "9": {"inputs": {"filename_prefix": WORKFLOW_PREFIX, "images": ["8", 0]}, "class_type": "SaveImage"}
}

class ImageRequest(BaseModel):
    prompt: str
    n: Optional[int] = 1
    size: Optional[str] = "1024x1024"
    response_format: Optional[str] = "b64_json"

async def queue_workflow(wf: dict) -> str:
    async with httpx.AsyncClient() as client:
        r = await client.post(f"{COMFYUI_URL}/api/prompt", json={"prompt": wf})
        r.raise_for_status()
        return r.json()["prompt_id"]

async def wait_for_image(pid: str, timeout: int = 300) -> list:
    async with httpx.AsyncClient() as client:
        start = time.time()
        while time.time() - start < timeout:
            r = await client.get(f"{COMFYUI_URL}/history/{pid}")
            if r.status_code == 200:
                data = r.json()
                if pid in data and "outputs" in data[pid]:
                    imgs = []
                    for node_out in data[pid]["outputs"].values():
                        if "images" in node_out:
                            imgs.extend(node_out["images"])
                    if imgs:
                        return imgs
            await asyncio.sleep(1)
        raise TimeoutError(f"Timeout waiting for {pid}")

async def fetch_image(filename: str, subfolder: str = "", folder_type: str = "output") -> bytes:
    async with httpx.AsyncClient() as client:
        params = {"filename": filename, "subfolder": subfolder, "type": folder_type}
        r = await client.get(f"{COMFYUI_URL}/view", params=params)
        r.raise_for_status()
        return r.content

@app.post("/v1/images/generations")
async def generate(req: ImageRequest):
    imgs = []
    for i in range(req.n):
        wf = json.loads(json.dumps(FLUX2_WORKFLOW))
        wf["5"]["inputs"]["text"] = req.prompt
        w, h = (1024, 1024)
        if "x" in req.size:
            try:
                w, h = map(int, req.size.split("x"))
            except ValueError:
                pass
        wf["4"]["inputs"]["width"] = w
        wf["4"]["inputs"]["height"] = h
        wf["7"]["inputs"]["seed"] = random.randint(1, 2**32)

        pid = await queue_workflow(wf)
        files = await wait_for_image(pid)
        if not files:
            raise RuntimeError("No images returned")
        data = await fetch_image(files[0]["filename"], files[0].get("subfolder", ""), files[0].get("type", "output"))
        imgs.append({"b64_json": base64.b64encode(data).decode()})

    return {"created": int(time.time()), "data": imgs}

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

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