# worker_layers.py — stage 9 part 1 depth-layer decomposition. FastAPI :8007.
# wraps the comfyui layered workflow (comfyui runs as its own external process, not session-mgr).
# panel in -> N layer pngs uploaded to minio, layer uris out.
import os, json, time, uuid
from fastapi import FastAPI
from pydantic import BaseModel
import requests
import transport

app = FastAPI()
transport.install_logging(app, "layers")
SHM = "/dev/shm"
COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://127.0.0.1:8188")
LAYERED_WORKFLOW = os.path.join(os.path.dirname(__file__), "legacy", "qwen_layered_workflow.json")


class LayerInput(BaseModel):
    panel_uri: str
    panel_id: str = ""
    num_layers: int = 4
    prompt: str = ""
    manga_id: str = ""
    chapter_id: str = ""


def _comfy_upload(path: str) -> str:
    with open(path, "rb") as f:
        r = requests.post(f"{COMFYUI_URL}/upload/image",
                          files={"image": (os.path.basename(path), f)},
                          data={"overwrite": "true"}, timeout=30)
    r.raise_for_status()
    return r.json()["name"]


def _comfy_run(panel_path: str, num_layers: int, prompt_text: str) -> list[str]:
    """returns comfyui /view urls for the produced layers."""
    workflow = json.load(open(LAYERED_WORKFLOW))
    workflow["10"]["inputs"]["image"] = _comfy_upload(panel_path)
    workflow["6"]["inputs"]["text"] = prompt_text
    workflow["83"]["inputs"]["layers"] = num_layers - 1  # node returns incl. background
    pid = requests.post(f"{COMFYUI_URL}/prompt",
                        json={"prompt": workflow, "client_id": str(uuid.uuid4())},
                        timeout=30).json()["prompt_id"]
    deadline = time.time() + 600
    while time.time() < deadline:
        h = requests.get(f"{COMFYUI_URL}/history/{pid}", timeout=30).json()
        if pid in h:
            imgs = h[pid]["outputs"]["9"]["images"]
            return [f"{COMFYUI_URL}/view?filename={i['filename']}&subfolder={i['subfolder']}&type={i['type']}"
                    for i in imgs]
        time.sleep(2)
    raise TimeoutError(f"comfyui prompt {pid} timed out")


def _comfy_up() -> bool:
    try:
        return requests.get(f"{COMFYUI_URL}/system_stats", timeout=2).ok
    except requests.RequestException:
        return False


@app.post("/layers")
async def layers(data: LayerInput):
    if not _comfy_up():
        return {"layer_uris": [], "skipped": "comfyui down"}  # ponytail: skip stage if ComfyUI not running
    local = transport.get(data.panel_uri, f"{SHM}/layer_{uuid.uuid4().hex[:8]}.png")
    # orchestrator doesn't pass manga/chapter to layers; derive from the panel uri prefix.
    parts = data.panel_uri.replace("s3://", "").split("/")
    manga_id, chapter_id = parts[1], parts[2]
    view_urls = _comfy_run(local, data.num_layers, data.prompt)
    layer_uris = []
    for idx, url in enumerate(view_urls):
        png = f"{SHM}/layer_{uuid.uuid4().hex[:6]}.png"
        with open(png, "wb") as f:
            f.write(requests.get(url, timeout=60).content)
        uri = f"s3://manga/{manga_id}/{chapter_id}/layers/{data.panel_id or 'p'}/{idx}.png"
        transport.put(png, uri)
        os.remove(png)
        layer_uris.append(uri)
    os.remove(local)
    return {"layer_uris": layer_uris}


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


if __name__ == "__main__":
    # self-check: mock comfyui history -> view-url extraction picks the SaveImage node images.
    fake_history = {"pid": {"outputs": {"9": {"images": [
        {"filename": "a.png", "subfolder": "", "type": "output"},
        {"filename": "b.png", "subfolder": "sub", "type": "output"},
    ]}}}}
    imgs = fake_history["pid"]["outputs"]["9"]["images"]
    urls = [f"x?filename={i['filename']}&subfolder={i['subfolder']}" for i in imgs]
    assert len(urls) == 2 and "a.png" in urls[0]
    assert os.path.exists(LAYERED_WORKFLOW), LAYERED_WORKFLOW
    print("worker_layers self-check ok")