#!/usr/bin/env python3
"""SCAIL-2 video generation script.

Directly calls ComfyUI nodes — no server, no API, no GUI.
Uses cuda:0 only.

Usage:
    CUDA_VISIBLE_DEVICES=0 python3 generate.py --pose pose.mp4 --ref ref.png -o output.mp4
"""

import argparse
import json
import os
import sys

# ── cuda:0 only ──
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import cv2
import numpy as np
import torch

# ── ComfyUI path ──
COMFY_DIR = "/home/ubuntu/ComfyUI"
sys.path.insert(0, COMFY_DIR)
import folder_paths
import nodes
from comfy_extras.nodes_custom_sampler import KSamplerSelect, BasicScheduler, SamplerCustom
from comfy_extras.nodes_scail import WanSCAILToVideo
from comfy_extras.nodes_post_processing import ColorTransfer


# ────────────────────────── helpers ──────────────────────────


def load_video_frames(path, max_frames=None, target_fps=None):
    """Load video → ComfyUI IMAGE tensor [B, H, W, C] float32 0-1."""
    cap = cv2.VideoCapture(path)
    if not cap.isOpened():
        raise FileNotFoundError(f"Cannot open video: {path}")
    src_fps = cap.get(cv2.CAP_PROP_FPS)
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(frame)
    cap.release()
    if not frames:
        raise ValueError(f"No frames in {path}")

    # fps conversion
    if target_fps and src_fps > 0 and src_fps != target_fps:
        step = src_fps / target_fps
        indices = np.arange(0, len(frames), step).astype(int)
        indices = indices[indices < len(frames)]
        frames = [frames[i] for i in indices]

    if max_frames:
        frames = frames[:max_frames]

    tensor = torch.from_numpy(np.stack(frames)).float() / 255.0  # B,H,W,C
    print(f"Loaded video: {tensor.shape[0]} frames, {tensor.shape[2]}x{tensor.shape[1]}, "
          f"src_fps={src_fps:.1f}")
    return tensor


def load_image(path):
    """Load image → ComfyUI IMAGE tensor [1, H, W, C] float32 0-1."""
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    if img is None:
        raise FileNotFoundError(f"Cannot open image: {path}")
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) if img.ndim == 3 else img
    tensor = torch.from_numpy(img).float() / 255.0
    if tensor.ndim == 2:  # grayscale
        tensor = tensor.unsqueeze(-1)
    tensor = tensor.unsqueeze(0)  # 1,H,W,C
    print(f"Loaded image: {tensor.shape[2]}x{tensor.shape[1]}")
    return tensor


def load_reference_images(path, target_h, target_w):
    """Load one image OR a directory of images → batched tensor [N, H, W, C].

    SCAIL-2 multi-reference: first image = primary ref, rest = additional views.
    All images are resized+center-cropped to (target_w, target_h).
    """
    if os.path.isdir(path):
        files = sorted(
            os.path.join(path, f) for f in os.listdir(path)
            if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
        )
        if not files:
            raise FileNotFoundError(f"No images in directory: {path}")
    else:
        files = [path]

    tensors = []
    for f in files:
        img = cv2.imread(f, cv2.IMREAD_COLOR)
        if img is None:
            raise FileNotFoundError(f"Cannot open image: {f}")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        h, w = img.shape[:2]
        # center-crop to target aspect ratio, then resize
        target_ar = target_w / target_h
        src_ar = w / h
        if src_ar > target_ar:  # too wide → crop sides
            new_w = int(h * target_ar)
            x0 = (w - new_w) // 2
            img = img[:, x0:x0 + new_w]
        else:  # too tall → crop top/bottom
            new_h = int(w / target_ar)
            y0 = (h - new_h) // 2
            img = img[y0:y0 + new_h, :]
        img = cv2.resize(img, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)
        t = torch.from_numpy(img).float() / 255.0
        tensors.append(t)

    batch = torch.stack(tensors)  # N,H,W,C
    print(f"Loaded {batch.shape[0]} reference image(s) from {path} "
          f"(resized to {target_w}x{target_h})")
    return batch


def save_video(tensor, path, fps=24):
    """Save ComfyUI IMAGE tensor [B, H, W, C] float32 0-1 → mp4."""
    arr = (tensor.clone().cpu().numpy() * 255).astype(np.uint8)  # B,H,W,C
    h, w = arr.shape[1], arr.shape[2]
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(path, fourcc, fps, (w, h))
    for frame in arr:
        writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
    writer.release()
    print(f"Saved: {path} ({arr.shape[0]} frames, {w}x{h}, {fps} fps)")


# ────────────────────────── main ──────────────────────────


def main():
    p = argparse.ArgumentParser(description="SCAIL-2 video generation")
    p.add_argument("--pose", required=True, help="Pose driving video (mp4)")
    p.add_argument("--ref", required=True,
                   help="Reference image OR directory of reference images "
                        "(first = primary, rest = additional views)")
    p.add_argument("-o", "--output", default="output.mp4")
    p.add_argument("--positive", default="masterpiece, best quality, high quality, detailed")
    p.add_argument("--negative", default="")
    p.add_argument("--width", type=int, default=512)
    p.add_argument("--height", type=int, default=896)
    p.add_argument("--chunk-length", type=int, default=81)
    p.add_argument("--overlap", type=int, default=5)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--cfg", type=float, default=1.0)
    p.add_argument("--steps", type=int, default=6)
    p.add_argument("--sampler", default="euler")
    p.add_argument("--scheduler", default="simple")
    p.add_argument("--shift", type=float, default=5.0,
                   help="ModelSamplingSD3 shift for flow model (0 to disable)")
    p.add_argument("--fps", type=int, default=24)
    p.add_argument("--max-frames", type=int, default=None)
    p.add_argument("--replacement-mode", action="store_true")
    p.add_argument("--lora", default=None)
    p.add_argument("--lora-strength", type=float, default=0.8)
    p.add_argument("--mask", action="store_true",
                   help="Generate SAM3 pose/reference masks (cached on disk)")
    p.add_argument("--mask-cache", default=None,
                   help="Mask cache directory (default: ./mask_cache)")
    p.add_argument("--config", default=None, help="JSON config file (overrides CLI)")

    args = p.parse_args()
    if args.config:
        with open(args.config) as f:
            cfg = json.load(f)
        for k, v in cfg.items():
            setattr(args, k, v)

    device = torch.device("cuda:0")
    print(f"Using device: {device}")

    # ── 1. Load inputs ──
    print("\n=== Loading inputs ===")
    pose_video = load_video_frames(args.pose, max_frames=args.max_frames)
    ref_image = load_reference_images(args.ref, args.height, args.width)

    # ── 2. SAM3 masks FIRST (before the 14B model eats the GPU) ──
    pose_mask = ref_mask = None
    if args.mask:
        print("\n=== SAM3 masks ===")
        import masks as mask_mod
        kwargs = {}
        if args.mask_cache:
            kwargs["cache_dir"] = args.mask_cache
        pose_mask, ref_mask = mask_mod.get_masks(
            pose_path=args.pose, ref_path=args.ref,
            pose_video=pose_video, ref_images=ref_image,
            width=args.width, height=args.height,
            replacement_mode=args.replacement_mode, **kwargs)

    # ── 3. Load models ──
    print("\n=== Loading models ===")

    # Diffusion model (SCAIL-2)
    import importlib.util
    _kjnodes_spec = importlib.util.spec_from_file_location(
        'kjnodes_model_loader',
        os.path.join(COMFY_DIR, 'custom_nodes/comfyui-kjnodes/nodes/model_optimization_nodes.py')
    )
    _kjnodes_mod = importlib.util.module_from_spec(_kjnodes_spec)
    _kjnodes_spec.loader.exec_module(_kjnodes_mod)
    model_loader = _kjnodes_mod.DiffusionModelLoaderKJ()
    model = model_loader.patch_and_load(
        model_name="wan2.1_14B_SCAIL_2_fp8_scaled.safetensors",
        weight_dtype="default",
        compute_dtype="default",
        patch_cublaslinear=False,
        sage_attention="disabled",
        enable_fp16_accumulation=False,
    )[0]
    print(f"  Model loaded")

    # LoRA
    if args.lora:
        from nodes import LoraLoaderModelOnly
        lora_loader = LoraLoaderModelOnly()
        model = lora_loader.load_lora_model_only(
            model=model, lora_name=args.lora, strength_model=args.lora_strength,
        )[0]
        print(f"  LoRA loaded: {args.lora} @ {args.lora_strength}")

    # ModelSamplingSD3 shift (required for Wan/flow models)
    if args.shift > 0:
        from comfy_extras.nodes_model_advanced import ModelSamplingSD3
        model = ModelSamplingSD3().patch(model=model, shift=args.shift)[0]
        print(f"  ModelSamplingSD3 shift={args.shift}")

    # VAE
    from nodes import VAELoader
    vae_loader = VAELoader()
    vae = vae_loader.load_vae("wan_2.1_vae.safetensors")[0]
    print(f"  VAE loaded")

    # CLIP (text encoder)
    from nodes import CLIPLoader
    clip_loader = CLIPLoader()
    clip = clip_loader.load_clip(
        clip_name="umt5_xxl_fp8_e4m3fn_scaled.safetensors",
        type="wan",
    )[0]
    print(f"  CLIP loaded")

    # CLIP Vision
    from nodes import CLIPVisionLoader, CLIPVisionEncode
    clip_vision_loader = CLIPVisionLoader()
    clip_vision = clip_vision_loader.load_clip("clip_vision_h_fp16.safetensors")[0]
    print(f"  CLIP Vision loaded")

    # ── 4. Text encoding ──
    print("\n=== Encoding prompts ===")
    from nodes import CLIPTextEncode
    text_encoder = CLIPTextEncode()
    positive = text_encoder.encode(clip=clip, text=args.positive)[0]
    negative = text_encoder.encode(clip=clip, text=args.negative)[0]
    print(f"  Positive: {args.positive[:60]}...")
    print(f"  Negative: {args.negative[:60]}...")

    # CLIP Vision encode
    clip_vision_encode = CLIPVisionEncode()
    clip_vision_output = clip_vision_encode.encode(
        clip_vision=clip_vision, image=ref_image, crop="center"
    )[0]
    print(f"  CLIP Vision encoded")

    # ── 4. Sampler & sigmas ──
    sampler = KSamplerSelect.execute(sampler_name=args.sampler).args[0]
    sigmas = BasicScheduler.execute(
        model=model, scheduler=args.scheduler, steps=args.steps, denoise=1.0
    ).args[0]

    # ── 5. Generate (chunked loop) ──
    print(f"\n=== Generating ({pose_video.shape[0]} frames, "
          f"chunk={args.chunk_length}, overlap={args.overlap}) ===")

    _scail_spec = importlib.util.spec_from_file_location(
        'scail_auto_extend',
        os.path.join(COMFY_DIR, 'custom_nodes/scail-auto-extend/__init__.py')
    )
    _scail_mod = importlib.util.module_from_spec(_scail_spec)
    _scail_spec.loader.exec_module(_scail_mod)
    gen = _scail_mod.SCAILAutoExtend()
    with torch.inference_mode():
        output_frames, frame_count = gen.generate(
            model=model,
            positive=positive,
            negative=negative,
            vae=vae,
            sampler=sampler,
            sigmas=sigmas,
            pose_video=pose_video,
            width=args.width,
            height=args.height,
            noise_seed=args.seed,
            cfg=args.cfg,
            chunk_length=args.chunk_length,
            overlap=args.overlap,
            seed_mode="increment",
            color_transfer=True,
            reference_image=ref_image,
            reference_image_mask=ref_mask,
            pose_video_mask=pose_mask,
            clip_vision_output=clip_vision_output,
            replacement_mode=args.replacement_mode,
        )

    # ── 6. Save ──
    print(f"\n=== Saving {frame_count} frames ===")
    save_video(output_frames, args.output, fps=args.fps)
    print("Done!")


if __name__ == "__main__":
    main()