"""Wan 2.2 TI2V-5B text/image→video via the resident ComfyUI (same pattern as
comfyui_sd: pure stdlib, ComfyUI's venv does the heavy lifting, model stays
cached in RAM between jobs). Frames come back as PNGs; ffmpeg muxes the MP4."""
import argparse
import json
import mimetypes
import os
import random
import subprocess
import sys
import time
import urllib.parse
import urllib.request
import uuid
from pathlib import Path

# job env has a minimal PATH; we shell out to ffmpeg for the mux
os.environ["PATH"] = "/opt/homebrew/bin:/usr/local/bin:" + os.environ.get("PATH", "")

ROOT = Path(__file__).resolve().parents[3]
COMFY = ROOT / "vendor" / "comfyui"
API = "http://127.0.0.1:8188"
FPS = 24  # Wan 2.2

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)
outdir = Path(a.outdir)

prompt = (p.get("prompt") or "").strip()
if not prompt:
    print("ERROR: prompt is required")
    sys.exit(1)
if not (COMFY / "main.py").exists():
    print(f"ERROR: ComfyUI not installed at {COMFY}. Run scripts/install_comfyui.sh")
    sys.exit(1)
model_file = COMFY / "models" / "diffusion_models" / "wan2.2_ti2v_5B_fp16.safetensors"
if not model_file.exists():
    print("ERROR: Wan weights missing — see manifest description for the three files")
    sys.exit(1)


def alive():
    try:
        urllib.request.urlopen(f"{API}/", timeout=3)
        return True
    except Exception:
        return False


def ensure_server():
    if alive():
        print("comfyui: already resident", flush=True)
        return
    print("comfyui: starting resident server ...", flush=True)
    log = open("/tmp/comfyui.log", "ab")
    subprocess.Popen(
        [str(COMFY / ".venv" / "bin" / "python"), "main.py", "--port", "8188"],
        cwd=str(COMFY), stdout=log, stderr=log, start_new_session=True)
    for _ in range(90):
        if alive():
            print("comfyui: up", flush=True)
            return
        time.sleep(2)
    print("ERROR: ComfyUI did not come up — see /tmp/comfyui.log")
    sys.exit(1)


def upload_image(path):
    """POST multipart to /upload/image, return server-side filename."""
    boundary = uuid.uuid4().hex
    fname = Path(path).name
    ctype = mimetypes.guess_type(fname)[0] or "application/octet-stream"
    body = (f"--{boundary}\r\nContent-Disposition: form-data; name=\"image\"; "
            f"filename=\"{fname}\"\r\nContent-Type: {ctype}\r\n\r\n").encode()
    body += Path(path).read_bytes()
    body += f"\r\n--{boundary}--\r\n".encode()
    req = urllib.request.Request(f"{API}/upload/image", data=body, headers={
        "Content-Type": f"multipart/form-data; boundary={boundary}"})
    return json.load(urllib.request.urlopen(req))["name"]


def build(seed, start_image):
    length = int(p.get("length", 49))
    if length % 4 != 1:
        length = (length // 4) * 4 + 1  # model requires 4n+1 frames
    lat = {"vae": ["v", 0], "width": int(p.get("width", 832)),
           "height": int(p.get("height", 480)), "length": length, "batch_size": 1}
    g = {
        "u": {"class_type": "UNETLoader", "inputs": {
            "unet_name": "wan2.2_ti2v_5B_fp16.safetensors", "weight_dtype": "default"}},
        "c": {"class_type": "CLIPLoader", "inputs": {
            "clip_name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors", "type": "wan",
            "device": "default"}},
        "v": {"class_type": "VAELoader", "inputs": {"vae_name": "wan2.2_vae.safetensors"}},
        "mo": {"class_type": "ModelSamplingSD3", "inputs": {"model": ["u", 0], "shift": 8.0}},
        "pos": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["c", 0]}},
        "neg": {"class_type": "CLIPTextEncode", "inputs": {
            "text": p.get("negative", "blurry, distorted, low quality, static image, watermark, text"),
            "clip": ["c", 0]}},
        "lat": {"class_type": "Wan22ImageToVideoLatent", "inputs": lat},
        "ks": {"class_type": "KSampler", "inputs": {
            "seed": seed, "steps": int(p.get("steps", 20)), "cfg": float(p.get("cfg", 5.0)),
            "sampler_name": "uni_pc", "scheduler": "simple", "denoise": 1.0,
            "model": ["mo", 0], "positive": ["pos", 0], "negative": ["neg", 0],
            "latent_image": ["lat", 0]}},
        "dec": {"class_type": "VAEDecode", "inputs": {"samples": ["ks", 0], "vae": ["v", 0]}},
        "sav": {"class_type": "SaveImage", "inputs": {
            "filename_prefix": "mb_wan", "images": ["dec", 0]}},
    }
    if start_image:
        g["img"] = {"class_type": "LoadImage", "inputs": {"image": start_image}}
        lat["start_image"] = ["img", 0]
    return g


ensure_server()
seed = int(p.get("seed", -1))
if seed < 0:
    seed = random.randint(0, 2**31 - 1)
start_image = upload_image(a.input[0]) if a.input else None
print(f"seed={seed} mode={'i2v' if start_image else 't2v'} "
      f"{p.get('width', 832)}x{p.get('height', 480)}x{p.get('length', 49)}f", flush=True)

body = json.dumps({"prompt": build(seed, start_image)}).encode()
req = urllib.request.Request(f"{API}/prompt", data=body,
                             headers={"Content-Type": "application/json"})
try:
    pid = json.load(urllib.request.urlopen(req))["prompt_id"]
except urllib.error.HTTPError as e:
    print(f"ERROR: ComfyUI rejected the workflow: {e.read().decode()[:600]}")
    sys.exit(1)

t0 = time.time()
images = []
while time.time() - t0 < 5400:
    h = json.load(urllib.request.urlopen(f"{API}/history/{pid}"))
    if pid in h:
        st = h[pid].get("status", {})
        if st.get("status_str") == "error":
            print("ERROR: generation failed — check /tmp/comfyui.log")
            print(json.dumps(st)[:400])
            sys.exit(1)
        if h[pid].get("outputs"):
            for node in h[pid]["outputs"].values():
                images += node.get("images", [])
            break
    time.sleep(5)

if not images:
    print("ERROR: no frames produced (timeout after 90min)")
    sys.exit(1)

import tempfile
frames_dir = Path(tempfile.mkdtemp(prefix="wan_frames_"))  # outside outdir so frames don't register as assets
for i, im in enumerate(images):
    q = urllib.parse.urlencode({"filename": im["filename"],
                                "subfolder": im.get("subfolder", ""),
                                "type": im.get("type", "output")})
    (frames_dir / f"f{i:05d}.png").write_bytes(
        urllib.request.urlopen(f"{API}/view?{q}").read())

out_mp4 = outdir / f"wan_{seed}.mp4"
r = subprocess.run(["ffmpeg", "-y", "-framerate", str(FPS),
                    "-i", str(frames_dir / "f%05d.png"),
                    "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "18",
                    str(out_mp4)], capture_output=True, text=True)
if r.returncode != 0 or not out_mp4.exists():
    print("ERROR: ffmpeg mux failed:", r.stderr[-800:])
    sys.exit(1)
print(f"done: {len(images)} frames → {out_mp4.name} in {time.time() - t0:.1f}s", flush=True)