#!/usr/bin/env python3
"""Wan 2.2 Optimized for Quality + Smoothness"""

import json
import urllib.request
import urllib.error
import uuid

COMFY_URL = "http://localhost:8188"
INPUT_IMAGE = "pool_party.png"

# Optimized prompt - explicit hand quality and smooth motion
PROMPT = """Static camera shot. Two men stand affectionately beside a swimming pool at sunset. The larger man on the left gently lifts the white shirt off the smaller man, pulling it slowly upward and over his head in one smooth motion. He then playfully tickles the smaller man's chest with his fingertips in a teasing, loving way. Both men smile and laugh together. Detailed well-formed hands with proper fingers. Natural smooth body movement. The pool, patio, and trees remain perfectly still in the background. Camera remains perfectly static. No zooming, no panning, no camera shake."""

NEGATIVE = "blurry, low quality, distorted hands, malformed hands, extra fingers, missing fingers, fused fingers, claw hands, jerky motion, choppy movement, fast motion, glitchy, artifacts, watermark, text, subtitle"

workflow = {
    "1": {
        "class_type": "LoadImage",
        "inputs": {"image": INPUT_IMAGE, "upload": "image"}
    },
    "2": {
        "class_type": "ImageResizeKJv2",
        "inputs": {
            "image": ["1", 0],
            "width": 848,
            "height": 480,
            "upscale_method": "lanczos",
            "keep_proportion": "crop",
            "pad_color": "0, 0, 0",
            "crop_position": "center",
            "divisible_by": 32
        }
    },
    "8": {
        "class_type": "WanVideoBlockSwap",
        "inputs": {"blocks_to_swap": 20, "offload_img_emb": False, "offload_txt_emb": False}
    },
    # LightX2V LoRA for better quality
    "20": {
        "class_type": "WanVideoLoraSelect",
        "inputs": {
            "lora_name": "WanVideo/Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank64_bf16.safetensors",
            "strength": 1.5,
            "prev_lora": None
        }
    },
    "3": {
        "class_type": "WanVideoModelLoader",
        "inputs": {
            "model": "WanVideo/2_2/Wan2_2-I2V-A14B-HIGH_fp8_e4m3fn_scaled_KJ.safetensors",
            "base_precision": "fp16_fast",
            "quantization": "fp8_e4m3fn_scaled",
            "load_device": "offload_device",
            "attention_mode": "sdpa",
            "block_swap_args": ["8", 0],
            "lora": ["20", 0]
        }
    },
    "21": {
        "class_type": "WanVideoLoraSelect",
        "inputs": {
            "lora_name": "WanVideo/Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank64_bf16.safetensors",
            "strength": 1.0,
            "prev_lora": None
        }
    },
    "4": {
        "class_type": "WanVideoModelLoader",
        "inputs": {
            "model": "WanVideo/2_2/Wan2_2-I2V-A14B-LOW_fp8_e4m3fn_scaled_KJ.safetensors",
            "base_precision": "fp16_fast",
            "quantization": "fp8_e4m3fn_scaled",
            "load_device": "offload_device",
            "attention_mode": "sdpa",
            "block_swap_args": ["8", 0],
            "lora": ["21", 0]
        }
    },
    "5": {
        "class_type": "WanVideoVAELoader",
        "inputs": {"model_name": "wanvideo/Wan2_1_VAE_bf16.safetensors", "precision": "bf16"}
    },
    "6": {
        "class_type": "LoadWanVideoT5TextEncoder",
        "inputs": {
            "model_name": "umt5-xxl-enc-fp8_e4m3fn.safetensors",
            "precision": "bf16",
            "load_device": "offload_device",
            "quantization": "fp8_e4m3fn"
        }
    },
    "7": {
        "class_type": "WanVideoTextEncode",
        "inputs": {
            "positive_prompt": PROMPT,
            "negative_prompt": NEGATIVE,
            "t5": ["6", 0],
            "force_offload": True,
            "use_disk_cache": False
        }
    },
    "9": {
        "class_type": "WanVideoImageToVideoEncode",
        "inputs": {
            "width": ["2", 1],
            "height": ["2", 2],
            "num_frames": 81,
            "noise_aug_strength": 0,
            "start_latent_strength": 1,
            "end_latent_strength": 1,
            "force_offload": True,
            "vae": ["5", 0],
            "start_image": ["2", 0]
        }
    },
    # Pass 1: HIGH with unipc (smoother), low CFG
    "10": {
        "class_type": "WanVideoSampler",
        "inputs": {
            "model": ["3", 0],
            "image_embeds": ["9", 0],
            "steps": 8,
            "cfg": 1.2,
            "shift": 8,
            "seed": 42,
            "force_offload": True,
            "scheduler": "unipc",
            "riflex_freq_index": 0,
            "text_embeds": ["7", 0],
            "end_step": 4
        }
    },
    # Pass 2: LOW with unipc, even lower CFG
    "11": {
        "class_type": "WanVideoSampler",
        "inputs": {
            "model": ["4", 0],
            "image_embeds": ["9", 0],
            "steps": 8,
            "cfg": 1.0,
            "shift": 8,
            "seed": 42,
            "force_offload": True,
            "scheduler": "unipc",
            "riflex_freq_index": 0,
            "text_embeds": ["7", 0],
            "samples": ["10", 0],
            "start_step": 4
        }
    },
    "12": {
        "class_type": "WanVideoDecode",
        "inputs": {
            "vae": ["5", 0],
            "samples": ["11", 0],
            "enable_vae_tiling": False,
            "enable_tiling": False,
            "tile_x": 272,
            "tile_y": 272,
            "tile_stride_x": 144,
            "tile_stride_y": 128
        }
    },
    "13": {
        "class_type": "VHS_VideoCombine",
        "inputs": {
            "images": ["12", 0],
            "frame_rate": 24,
            "loop_count": 0,
            "filename_prefix": "Wan22_Optimized",
            "format": "video/h264-mp4",
            "pix_fmt": "yuv420p",
            "crf": 16,
            "save_metadata": True,
            "trim_to_audio": False,
            "pingpong": False,
            "save_output": True
        }
    }
}

def submit():
    prompt_id = str(uuid.uuid4())
    data = {"prompt": workflow, "client_id": "selena-agent"}
    req = urllib.request.Request(
        f"{COMFY_URL}/prompt",
        data=json.dumps(data).encode('utf-8'),
        headers={'Content-Type': 'application/json'}
    )
    try:
        with urllib.request.urlopen(req) as response:
            result = json.loads(response.read().decode('utf-8'))
            print(f"✅ Submitted! ID: {result.get('prompt_id', prompt_id)}")
            return result.get('prompt_id')
    except urllib.error.HTTPError as e:
        err = e.read().decode('utf-8')
        print(f"❌ Error: {err[:800]}")
        return None

if __name__ == "__main__":
    print("🎬 Wan 2.2 Optimized:")
    print("   - LightX2V LoRA enabled")
    print("   - unipc scheduler (smoother)")
    print("   - Lower CFG (1.2/1.0)")
    print("   - 24fps output")
    print()
    submit()