#!/usr/bin/env python3
"""
Cloud I2V Batch Processor

Uploads keyframe images to a cloud GPU (Vast.ai),
runs Wan 2.1 I2V generation, and downloads results.

This is 10-50x faster than local M4 Max for video generation.

Requires: pip install vastai httpx
"""

import argparse
import asyncio
import json
import os
import re
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import httpx


@dataclass
class CloudInstance:
    id: str
    ip: str
    port: int
    ssh_port: int
    hourly_cost: float
    gpu: str


def run_vastai(args: list[str], capture_json: bool = False) -> dict | str:
    """Run a vastai CLI command."""
    cmd = ["vastai"] + args
    if capture_json:
        cmd.append("--raw")

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        raise RuntimeError(f"vastai command failed: {result.stderr}")

    if capture_json:
        return json.loads(result.stdout)
    return result.stdout


def search_offers(max_price: float = 0.50, min_gpu_ram: int = 24) -> list[dict]:
    """Search for available GPU instances using vastai CLI."""
    query = f"gpu_ram >= {min_gpu_ram} num_gpus = 1 rentable = true dph < {max_price} reliability > 0.95"

    offers = run_vastai(["search", "offers", query, "-o", "dph+"], capture_json=True)

    # Filter for good GPUs (both underscore and space formats)
    good_gpus = ["RTX 4090", "RTX 5090", "RTX A6000", "A100", "H100", "A40",
                 "RTX_4090", "RTX_5090", "RTX_A6000"]
    filtered = [o for o in offers if any(g in o.get("gpu_name", "") for g in good_gpus)]

    if not filtered:
        # If no exact matches, return all offers with 24GB+ VRAM
        return offers[:10]

    return filtered[:10]  # Top 10 cheapest suitable options


def create_instance(offer_id: int, image: str = "runpod/pytorch:2.1.0-py3.10-cuda11.8.0-devel-ubuntu22.04") -> int:
    """Create a new instance from an offer."""
    # Use onstart script to install ComfyUI
    onstart = """#!/bin/bash
set -e
cd /root

# Ensure pip is available
apt-get update && apt-get install -y python3-pip || true
pip3 install comfy-cli httpx || pip install comfy-cli httpx
comfy --skip-prompt install
comfy node install ComfyUI-GGUF ComfyUI-VideoHelperSuite

# Download models
mkdir -p /root/comfy/ComfyUI/models/diffusion_models
mkdir -p /root/comfy/ComfyUI/models/text_encoders
mkdir -p /root/comfy/ComfyUI/models/vae

# Wan I2V model (only if not exists)
if [ ! -f /root/comfy/ComfyUI/models/diffusion_models/wan2.1-i2v-14b-480p-Q5_K_M.gguf ]; then
    wget -q -O /root/comfy/ComfyUI/models/diffusion_models/wan2.1-i2v-14b-480p-Q5_K_M.gguf \
        "https://huggingface.co/city96/Wan2.1-I2V-14B-480P-GGUF/resolve/main/wan2.1-i2v-14b-480p-Q5_K_M.gguf"
fi

# Text encoder
if [ ! -f /root/comfy/ComfyUI/models/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors ]; then
    wget -q -O /root/comfy/ComfyUI/models/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors \
        "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors"
fi

# VAE
if [ ! -f /root/comfy/ComfyUI/models/vae/wan_2.1_vae.safetensors ]; then
    wget -q -O /root/comfy/ComfyUI/models/vae/wan_2.1_vae.safetensors \
        "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/vae/wan_2.1_vae.safetensors"
fi

# Start ComfyUI
cd /root/comfy/ComfyUI
python main.py --listen 0.0.0.0 --port 8188 &
"""

    # Write onstart script to temp file
    onstart_file = Path("/tmp/vastai_onstart.sh")
    onstart_file.write_text(onstart)

    result = run_vastai([
        "create", "instance", str(offer_id),
        "--image", image,
        "--disk", "80",
        "--onstart-cmd", onstart,
        "--direct",
    ], capture_json=False)  # CLI doesn't return JSON for create

    # Parse instance ID from success message like "Started. {'new_contract': 12345, ...}"
    # or "success: created instance 12345"
    match = re.search(r"new_contract['\"]?\s*:\s*(\d+)", result)
    if match:
        return int(match.group(1))

    match = re.search(r"instance\s+(\d+)", result, re.IGNORECASE)
    if match:
        return int(match.group(1))

    # Try to find any number that looks like an instance ID
    match = re.search(r"(\d{7,})", result)
    if match:
        return int(match.group(1))

    raise RuntimeError(f"Failed to parse instance ID from: {result}")


def get_instance_info(instance_id: int) -> Optional[dict]:
    """Get instance connection info."""
    instances = run_vastai(["show", "instances"], capture_json=True)

    for inst in instances:
        if inst.get("id") == instance_id:
            return inst
    return None


def wait_for_instance(instance_id: int, timeout: int = 900) -> dict:
    """Wait for instance to be ready."""
    print(f"Waiting for instance {instance_id} to be ready...")
    start = time.time()
    last_status = None

    while time.time() - start < timeout:
        info = get_instance_info(instance_id)
        if info:
            status = info.get("actual_status", "")
            if status != last_status:
                print(f"  Status: {status}")
                last_status = status

            if status == "running":
                return info  # Don't wait here - we'll check ComfyUI separately

        time.sleep(10)

    raise TimeoutError(f"Instance {instance_id} did not become ready in {timeout}s")


def destroy_instance(instance_id: int):
    """Destroy an instance."""
    try:
        run_vastai(["destroy", "instance", str(instance_id)])
        print(f"Destroyed instance {instance_id}")
    except Exception as e:
        print(f"Warning: Failed to destroy instance: {e}")


def build_i2v_workflow(image_name: str, motion_prompt: str, num_frames: int = 33, steps: int = 6) -> dict:
    """Build ComfyUI workflow for Wan I2V."""
    import random
    seed = random.randint(0, 2**32 - 1)

    return {
        "1": {"class_type": "LoadImage", "inputs": {"image": image_name}},
        "2": {"class_type": "ImageScale", "inputs": {
            "image": ["1", 0], "width": 832, "height": 480,
            "upscale_method": "lanczos", "crop": "center"
        }},
        "3": {"class_type": "UnetLoaderGGUF", "inputs": {
            "unet_name": "wan2.1-i2v-14b-480p-Q5_K_M.gguf"
        }},
        "4": {"class_type": "CLIPLoader", "inputs": {
            "clip_name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors", "type": "wan"
        }},
        "5": {"class_type": "VAELoader", "inputs": {
            "vae_name": "wan_2.1_vae.safetensors"
        }},
        "6": {"class_type": "CLIPTextEncode", "inputs": {
            "text": motion_prompt, "clip": ["4", 0]
        }},
        "7": {"class_type": "CLIPTextEncode", "inputs": {
            "text": "blurry, distorted, watermark, static", "clip": ["4", 0]
        }},
        "8": {"class_type": "WanImageToVideo", "inputs": {
            "positive": ["6", 0], "negative": ["7", 0], "vae": ["5", 0],
            "width": 832, "height": 480, "length": num_frames,
            "batch_size": 1, "start_image": ["2", 0],
        }},
        "9": {"class_type": "ModelSamplingSD3", "inputs": {
            "model": ["3", 0], "shift": 8.0
        }},
        "10": {"class_type": "KSampler", "inputs": {
            "model": ["9", 0], "positive": ["8", 0], "negative": ["8", 1],
            "latent_image": ["8", 2], "seed": seed, "steps": steps,
            "cfg": 5.0, "sampler_name": "uni_pc", "scheduler": "normal",
            "denoise": 1.0,
        }},
        "11": {"class_type": "VAEDecode", "inputs": {
            "samples": ["10", 0], "vae": ["5", 0]
        }},
        "12": {"class_type": "VHS_VideoCombine", "inputs": {
            "frame_rate": 16, "loop_count": 0,
            "filename_prefix": f"i2v_{Path(image_name).stem}",
            "format": "video/h264-mp4", "pingpong": False,
            "save_output": True, "images": ["11", 0],
        }},
    }


async def submit_workflow(base_url: str, workflow: dict) -> str:
    """Submit workflow to ComfyUI and return prompt_id."""
    async with httpx.AsyncClient(timeout=60) as client:
        resp = await client.post(f"{base_url}/prompt", json={"prompt": workflow})
        resp.raise_for_status()
        return resp.json()["prompt_id"]


async def wait_for_completion(base_url: str, prompt_id: str, timeout: int = 600) -> Optional[str]:
    """Wait for workflow to complete."""
    async with httpx.AsyncClient(timeout=30) as client:
        start = time.time()
        while time.time() - start < timeout:
            await asyncio.sleep(10)

            resp = await client.get(f"{base_url}/history/{prompt_id}")
            if resp.status_code == 200:
                history = resp.json()
                if prompt_id in history:
                    entry = history[prompt_id]
                    status = entry.get("status", {}).get("status_str", "")

                    if status == "success":
                        outputs = entry.get("outputs", {})
                        for node_id, output in outputs.items():
                            if "gifs" in output:
                                return output["gifs"][0]["filename"]
                            if "videos" in output:
                                return output["videos"][0]["filename"]
                    elif "error" in status.lower():
                        print(f"  Workflow error: {entry}")
                        return None

        print(f"  Timeout waiting for {prompt_id}")
        return None


async def run_batch(
    instance_info: dict,
    images_dir: Path,
    output_dir: Path,
    motion_prompt: str,
    steps: int = 6,
    frames: int = 33,
) -> list[Path]:
    """Run batch I2V on cloud instance."""

    # Get connection info
    ssh_host = instance_info.get("ssh_host", "").split(":")[0]
    ssh_port = int(instance_info.get("ssh_port", 22))
    public_ip = instance_info.get("public_ipaddr", ssh_host)

    # ComfyUI port (usually 8188, mapped to a high port)
    ports = instance_info.get("ports", {})
    comfy_port = None
    for port_map in ports.values() if isinstance(ports, dict) else []:
        if "8188" in str(port_map):
            comfy_port = port_map.get("HostPort", 8188)
            break

    if not comfy_port:
        comfy_port = 8188  # Default

    base_url = f"http://{public_ip}:{comfy_port}"
    print(f"ComfyUI URL: {base_url}")

    # Find images
    images = list(images_dir.glob("*.png")) + list(images_dir.glob("*.jpg"))
    print(f"Found {len(images)} images to process")

    # Upload images via SSH
    print("Uploading images...")

    # Find SSH key for Vast.ai (check for dedicated key first, then default)
    ssh_key = Path.home() / ".ssh" / "id_ed25519_vastai"
    if not ssh_key.exists():
        ssh_key = Path.home() / ".ssh" / "id_ed25519"
    if not ssh_key.exists():
        ssh_key = Path.home() / ".ssh" / "id_rsa"

    ssh_opts = f"-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -i {ssh_key}"
    scp_port = f"-P {ssh_port}"  # SCP uses capital P for port

    for img in images:
        result = subprocess.run(
            f"scp {ssh_opts} {scp_port} {img} root@{ssh_host}:/root/comfy/ComfyUI/input/",
            shell=True, capture_output=True, text=True
        )
        if result.returncode != 0:
            print(f"  Warning: Failed to upload {img.name}: {result.stderr}")
            raise RuntimeError(f"SCP failed: {result.stderr}")

    # Process each image
    output_videos = []
    for i, img in enumerate(images):
        print(f"Processing {i+1}/{len(images)}: {img.name}")

        workflow = build_i2v_workflow(img.name, motion_prompt, frames, steps)

        try:
            prompt_id = await submit_workflow(base_url, workflow)
            print(f"  Submitted: {prompt_id}")

            video_file = await wait_for_completion(base_url, prompt_id, timeout=600)
            if video_file:
                output_videos.append(video_file)
                print(f"  Complete: {video_file}")
            else:
                print(f"  Failed or timed out")
        except Exception as e:
            print(f"  Error: {e}")

    # Download results
    output_dir.mkdir(parents=True, exist_ok=True)
    print(f"\nDownloading {len(output_videos)} videos...")

    for video in output_videos:
        result = subprocess.run(
            f"scp {ssh_opts} {scp_port} root@{ssh_host}:/root/comfy/ComfyUI/output/{video} {output_dir}/",
            shell=True, capture_output=True, text=True
        )
        if result.returncode != 0:
            print(f"  Warning: Failed to download {video}: {result.stderr}")

    return list(output_dir.glob("*.mp4"))


async def main():
    parser = argparse.ArgumentParser(description="Cloud I2V Batch Processor")
    parser.add_argument("--images", type=Path, required=True, help="Input images directory")
    parser.add_argument("--output", type=Path, default=Path("i2v_output"), help="Output directory")
    parser.add_argument("--motion", default="subtle organic motion, gentle breathing, cinematic")
    parser.add_argument("--steps", type=int, default=6, help="Sampling steps (4-12)")
    parser.add_argument("--frames", type=int, default=33, help="Output frames (33=~2s at 16fps)")
    parser.add_argument("--max-price", type=float, default=0.50, help="Max $/hr")
    parser.add_argument("--dry-run", action="store_true", help="Show offers only")
    parser.add_argument("--yes", "-y", action="store_true", help="Skip confirmations")
    parser.add_argument("--instance", type=int, help="Use existing instance ID")

    args = parser.parse_args()

    # Count images
    images = list(args.images.glob("*.png")) + list(args.images.glob("*.jpg"))
    if not images:
        print(f"No images found in {args.images}")
        sys.exit(1)

    print(f"Found {len(images)} images to process")

    # Search for offers
    print(f"\nSearching for GPU instances under ${args.max_price}/hr...")
    offers = search_offers(args.max_price)

    if not offers:
        print("No suitable GPU instances found. Try increasing --max-price")
        sys.exit(1)

    print(f"\nTop offers:")
    print(f"{'ID':<12} {'GPU':<20} {'RAM':<8} {'$/hr':<10} {'Location':<20}")
    print("-" * 75)
    for o in offers[:5]:
        print(f"{o['id']:<12} {o.get('gpu_name', 'N/A'):<20} {o.get('gpu_ram', 0):<8.0f} ${o.get('dph_total', 0):<9.4f} {o.get('geolocation', 'N/A'):<20}")

    # Cost estimate
    best = offers[0]
    est_time_min = len(images) * 3  # ~3 min per clip on fast GPU
    est_cost = (est_time_min / 60 + 0.5) * best['dph_total']  # +30 min for setup

    print(f"\nEstimated: {est_time_min} min processing + 30 min setup = ${est_cost:.2f}")

    if args.dry_run:
        print("\n(Dry run - no instance created)")
        return

    # Confirm
    if not args.yes:
        response = input("\nProceed with cheapest option? [y/N] ")
        if response.lower() != "y":
            print("Cancelled")
            return
    else:
        print("\n--yes flag: proceeding automatically")

    instance_id = args.instance
    instance_info = None

    try:
        if not instance_id:
            # Create instance
            print(f"\nCreating instance from offer {best['id']}...")
            instance_id = create_instance(best['id'])
            print(f"Created instance: {instance_id}")

            # Wait for ready
            instance_info = wait_for_instance(instance_id)
        else:
            instance_info = get_instance_info(instance_id)
            if not instance_info:
                print(f"Instance {instance_id} not found")
                sys.exit(1)

        ssh_host = instance_info.get("ssh_host", "").split(":")[0]
        ssh_port = int(instance_info.get("ssh_port", 22))
        public_ip = instance_info.get("public_ipaddr", ssh_host)

        print(f"\nInstance ready!")
        print(f"  SSH: ssh -p {ssh_port} root@{ssh_host}")

        # Find SSH key
        ssh_key = Path.home() / ".ssh" / "id_ed25519_vastai"
        if not ssh_key.exists():
            ssh_key = Path.home() / ".ssh" / "id_ed25519"
        if not ssh_key.exists():
            ssh_key = Path.home() / ".ssh" / "id_rsa"
        ssh_opts = f"-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o ConnectTimeout=10 -i {ssh_key}"

        # Get ComfyUI port from port mappings
        ports = instance_info.get("ports", {})
        comfy_port = 8188
        for port_key, port_info in (ports.items() if isinstance(ports, dict) else []):
            if "8188" in str(port_key):
                if isinstance(port_info, list) and port_info:
                    comfy_port = port_info[0].get("HostPort", 8188)
                break

        comfy_url = f"http://{public_ip}:{comfy_port}"

        # Wait for ComfyUI setup to complete (downloads ~18GB of models)
        # Check both SSH connectivity AND ComfyUI input directory exists
        print(f"  Waiting for setup to complete (this can take 5-10 minutes for model downloads)...")

        comfy_ready = False
        for attempt in range(90):  # Wait up to 15 minutes
            elapsed = attempt * 10

            # First check if SSH works and input directory exists
            try:
                ssh_check = subprocess.run(
                    f"ssh {ssh_opts} -p {ssh_port} root@{ssh_host} 'test -d /root/comfy/ComfyUI/input && echo OK'",
                    shell=True, capture_output=True, text=True, timeout=20
                )
            except subprocess.TimeoutExpired:
                if attempt % 6 == 0:
                    print(f"    SSH connection timed out, retrying... ({elapsed}s)")
                await asyncio.sleep(10)
                continue

            if ssh_check.returncode == 0 and "OK" in ssh_check.stdout:
                # SSH works and directory exists, now check ComfyUI API
                try:
                    async with httpx.AsyncClient(timeout=10) as client:
                        resp = await client.get(f"{comfy_url}/system_stats")
                        if resp.status_code == 200:
                            print(f"  ComfyUI ready! (took {elapsed}s)")
                            comfy_ready = True
                            break
                except Exception:
                    pass

                if attempt % 3 == 0:
                    print(f"    ComfyUI directory exists, waiting for API... ({elapsed}s)")
            else:
                if attempt % 6 == 0:
                    print(f"    Setup in progress... ({elapsed}s)")

            await asyncio.sleep(10)

        if not comfy_ready:
            print("  Warning: ComfyUI setup may not be complete, but proceeding...")

        # Run batch
        downloaded = await run_batch(
            instance_info,
            args.images,
            args.output,
            args.motion,
            args.steps,
            args.frames,
        )

        print(f"\nComplete! Downloaded {len(downloaded)} videos to {args.output}")

    finally:
        if instance_id and not args.instance:
            if not args.yes:
                response = input("\nDestroy instance? [Y/n] ")
                if response.lower() != "n":
                    destroy_instance(instance_id)
            else:
                # Auto-destroy with --yes flag
                destroy_instance(instance_id)


if __name__ == "__main__":
    asyncio.run(main())