"""Single test clip — corrected WanImageToVideo workflow."""
import json, time, urllib.request, os, random

COMFY = "http://localhost:8188"
IMAGE_PATH = "assets/images/hero/hero-carpet-cleaning.webp"
OUT_DIR = "assets/videos/clips"
os.makedirs(OUT_DIR, exist_ok=True)


def upload_image(image_path):
    fname = os.path.basename(image_path)
    with open(image_path, "rb") as f:
        img_data = f.read()
    boundary = "----FormBoundary123456"
    body = (
        f"--{boundary}\r\n"
        f'Content-Disposition: form-data; name="image"; filename="{fname}"\r\n'
        f"Content-Type: image/webp\r\n\r\n"
    ).encode() + img_data + f"\r\n--{boundary}--\r\n".encode()
    req = urllib.request.Request(
        f"{COMFY}/upload/image", data=body,
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
    )
    with urllib.request.urlopen(req) as resp:
        result = json.loads(resp.read())
    print(f"  uploaded: {result['name']}")
    return result["name"]


def build_workflow(image_name, prompt, frames=25):
    # WanImageToVideo is a conditioning node, NOT a sampler.
    # outputs: [0]=positive CONDITIONING, [1]=negative CONDITIONING, [2]=latent LATENT
    # start_image is optional IMAGE — anchors first frame.
    return {
        "1": {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": "Wan2.2-TI2V-5B-Q4_K_M.gguf"}},
        "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors", "type": "wan"}},
        "3": {"class_type": "VAELoader", "inputs": {"vae_name": "wan_2.1_vae.safetensors"}},
        "4": {"class_type": "LoadImage", "inputs": {"image": image_name}},
        "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt}},
        "6": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": "blur, low quality, distortion, text, watermark, people, jitter"}},
        "7": {
            "class_type": "WanImageToVideo",
            "inputs": {
                "positive":    ["5", 0],
                "negative":    ["6", 0],
                "vae":         ["3", 0],
                "start_image": ["4", 0],
                "width": 832, "height": 480, "length": frames, "batch_size": 1,
            },
        },
        "8": {
            "class_type": "KSampler",
            "inputs": {
                "model":        ["1", 0],
                "positive":     ["7", 0],
                "negative":     ["7", 1],
                "latent_image": ["7", 2],
                "seed": 42, "steps": 20, "cfg": 6.0,
                "sampler_name": "uni_pc", "scheduler": "simple", "denoise": 1.0,
            },
        },
        # VAEDecodeTiled handles video (5D) latents — VAEDecode only handles images (4D)
        "9": {"class_type": "VAEDecodeTiled", "inputs": {"samples": ["8", 0], "vae": ["3", 0], "tile_size": 512, "overlap": 64, "temporal_size": 64, "temporal_overlap": 8}},
        "10": {
            "class_type": "SaveAnimatedWEBP",
            "inputs": {"images": ["9", 0], "filename_prefix": "wan_test", "fps": 12, "lossless": False, "quality": 85, "method": "default"},
        },
    }


def queue_prompt(workflow):
    data = json.dumps({"prompt": workflow}).encode()
    req = urllib.request.Request(f"{COMFY}/prompt", data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req) as resp:
        return json.loads(resp.read())["prompt_id"]


def wait_for_result(prompt_id, timeout=1800):
    start = time.time()
    while time.time() - start < timeout:
        with urllib.request.urlopen(f"{COMFY}/history/{prompt_id}") as resp:
            hist = json.loads(resp.read())
        if prompt_id in hist:
            entry = hist[prompt_id]
            if entry.get("status", {}).get("status_str") == "error":
                print(f"  ERROR: {entry['status'].get('messages', '')}")
                return None
            for node_out in entry.get("outputs", {}).values():
                if "gifs" in node_out:
                    return node_out["gifs"]
                if "images" in node_out:
                    return node_out["images"]
        elapsed = int(time.time() - start)
        print(f"  waiting... {elapsed}s", flush=True)
        time.sleep(15)
    return None


def download_output(vid_info, out_path):
    fname = vid_info["filename"]
    subfolder = vid_info.get("subfolder", "")
    img_type = vid_info.get("type", "output")
    url = f"{COMFY}/view?filename={fname}&subfolder={subfolder}&type={img_type}"
    with urllib.request.urlopen(url) as resp:
        data = resp.read()
    with open(out_path, "wb") as f:
        f.write(data)
    print(f"  saved: {out_path} ({len(data)//1024}KB)")


print("[TEST v2] WanImageToVideo → KSampler → VAEDecode → SaveAnimatedWEBP")
image_name = upload_image(IMAGE_PATH)
workflow = build_workflow(
    image_name,
    "slow dolly forward across clean plush cream carpet, gentle camera push toward the far wall, warm afternoon light, cinematic smooth motion",
    frames=9,
)
prompt_id = queue_prompt(workflow)
print(f"  queued: {prompt_id}")
results = wait_for_result(prompt_id)
if results:
    download_output(results[0], f"{OUT_DIR}/test-clip-01.webp")
    print("SUCCESS")
else:
    print("FAILED")