#!/usr/bin/env python3
"""End-to-end Text-to-Video (and Text-to-Image) inference for Lance-3B-Video-MLX.

Self-contained: works from this repo directory after `huggingface-cli download`.
Auto-fetches the Wan 2.2 VAE companion repo (`RockTalk/Wan2.2-VAE-MLX`) on first
run if its weights aren't already present alongside this script.

Usage:
    # T2V (default — 9-frame video at 256x256, ~22 s on M4 Studio)
    python inference.py --prompt "a calm ocean wave rolling onto a sandy beach"

    # Tune frame count: T = (T_lat - 1) * 4 + 1
    #   T_lat=1 -> 1 frame (image), T_lat=3 -> 9, T_lat=8 -> 29, T_lat=31 -> 121
    python inference.py --prompt "..." --t-lat 8

    # T2I fast path
    python inference.py --prompt "..." --t-lat 1 --size 512 --steps 30

Outputs:
  - <out>.png            — horizontal strip of all frames
  - <out>_frame*.png     — each frame as a separate file
  - <out>.mp4            — MP4 (if --mp4 and `imageio[ffmpeg]` is installed)

Verified on M4 Studio (128 GB). Requires Apple Silicon with MLX >= 0.29.
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
from PIL import Image

_materialize = getattr(mx, "eval")

REPO_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(REPO_DIR))

from lance_mlx.lance import Lance, LanceConfig  # noqa: E402
from lance_mlx.vae_wan22 import Wan2_2_VAE  # noqa: E402

try:
    from mlx_vlm.models.qwen2_5_vl.config import (
        ModelConfig, TextConfig, VisionConfig,
    )
except ImportError as e:
    raise SystemExit("mlx-vlm is required. Install with: pip install 'mlx-vlm>=0.3'") from e

try:
    from transformers import AutoTokenizer
except ImportError as e:
    raise SystemExit("transformers is required. Install with: pip install transformers") from e


def build_lance_config(cfg_json: dict) -> LanceConfig:
    qwen = cfg_json["qwen2_5_vl_config"]
    vc = qwen["vision_config"]
    text_cfg = TextConfig(
        model_type="qwen2_5_vl",
        hidden_size=qwen["hidden_size"],
        intermediate_size=qwen["intermediate_size"],
        num_hidden_layers=qwen["num_hidden_layers"],
        num_attention_heads=qwen["num_attention_heads"],
        num_key_value_heads=qwen["num_key_value_heads"],
        vocab_size=qwen["vocab_size"],
        rms_norm_eps=qwen["rms_norm_eps"],
        rope_theta=qwen["rope_theta"],
        rope_scaling=qwen["rope_scaling"],
        tie_word_embeddings=qwen.get("tie_word_embeddings", True),
    )
    vision_cfg = VisionConfig(
        model_type="qwen2_5_vl",
        hidden_size=vc["hidden_size"], out_hidden_size=vc["out_hidden_size"],
        intermediate_size=vc["intermediate_size"], depth=vc["depth"],
        num_heads=vc["num_heads"], patch_size=vc["patch_size"],
        spatial_merge_size=vc["spatial_merge_size"], in_channels=vc["in_chans"],
        spatial_patch_size=vc["spatial_patch_size"],
        temporal_patch_size=vc["temporal_patch_size"],
        window_size=vc["window_size"],
        fullatt_block_indexes=vc["fullatt_block_indexes"],
        tokens_per_second=vc["tokens_per_second"],
    )
    mc = ModelConfig(
        text_config=text_cfg, vision_config=vision_cfg, model_type="qwen2_5_vl",
        image_token_id=qwen["image_token_id"],
        video_token_id=qwen["video_token_id"],
        vision_start_token_id=qwen["vision_start_token_id"],
        vision_end_token_id=qwen["vision_end_token_id"],
        vision_token_id=qwen["vision_token_id"],
    )
    return LanceConfig(
        qwen_config=mc,
        latent_patch_size=tuple(cfg_json["latent_patch_size"]),
        max_latent_size=cfg_json["max_latent_size"],
        max_num_frames=cfg_json["max_num_frames"],
        max_num_latent_frames_override=cfg_json.get("max_num_latent_frames"),
        latent_channel=cfg_json["latent_channel"],
        vae_downsample_spatial=cfg_json["vae_downsample_spatial"],
        vae_downsample_temporal=cfg_json["vae_downsample_temporal"],
        timestep_shift=cfg_json["timestep_shift"],
    )


def ensure_vae_weights(repo_dir: Path) -> Path:
    candidate = repo_dir / "wan22_vae.safetensors"
    if candidate.exists():
        return candidate
    try:
        from huggingface_hub import hf_hub_download
    except ImportError as e:
        raise SystemExit(
            "huggingface_hub is required to fetch the Wan VAE. "
            "Install with: pip install huggingface_hub"
        ) from e
    print("[setup] Fetching Wan 2.2 VAE from RockTalk/Wan2.2-VAE-MLX ...")
    downloaded = Path(hf_hub_download(
        repo_id="RockTalk/Wan2.2-VAE-MLX",
        filename="model.safetensors",
    ))
    target = repo_dir / "wan22_vae.safetensors"
    try:
        target.symlink_to(downloaded)
    except OSError:
        import shutil
        shutil.copy(downloaded, target)
    return target


def save_outputs(video_np: np.ndarray, out_path: Path, want_mp4: bool, fps: int) -> None:
    """video_np: (T, H, W, 3) in [-1, 1]. Saves strip + per-frame PNGs + optional MP4."""
    video_u8 = np.clip((video_np + 1.0) * 127.5, 0, 255).astype(np.uint8)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    strip = np.concatenate(list(video_u8), axis=1)  # (H, T*W, 3)
    Image.fromarray(strip).save(out_path)
    print(f"[ok] frame strip -> {out_path}")

    base = out_path.with_suffix("")
    for i, frame in enumerate(video_u8):
        Image.fromarray(frame).save(f"{base}_frame{i:02d}.png")
    print(f"[ok] per-frame PNGs -> {base}_frame*.png")

    if want_mp4:
        try:
            import imageio.v3 as iio
            mp4_path = out_path.with_suffix(".mp4")
            iio.imwrite(mp4_path, video_u8, fps=fps, codec="libx264")
            print(f"[ok] mp4 -> {mp4_path}")
        except Exception as exc:
            print(f"[warn] MP4 export failed: {exc}")
            print("       Install with: pip install 'imageio[ffmpeg]'")


def main(args: argparse.Namespace) -> None:
    repo = REPO_DIR
    n_frames = (args.t_lat - 1) * 4 + 1
    print("=== Lance-3B-Video-MLX ===")
    print(f"prompt: {args.prompt!r}")
    print(f"out:    {args.out}")
    print(f"size:   {args.size}x{args.size}  steps: {args.steps}  "
          f"T_lat={args.t_lat} → {n_frames} frames  cfg: {args.cfg}\n")

    t0 = time.time()
    cfg_json = json.loads((repo / "config.json").read_text())
    lance_cfg = build_lance_config(cfg_json)
    model = Lance(lance_cfg)
    print(f"[ok] Lance built ({time.time()-t0:.1f}s)")

    t0 = time.time()
    weights = mx.load(str(repo / "model.safetensors"))
    non_vit = {k: v for k, v in weights.items() if not k.startswith("vit_model.")}
    n_vit = len(weights) - len(non_vit)
    model.load_weights(list(non_vit.items()), strict=True)
    _materialize(model.parameters())
    print(f"[ok] strict load — {len(non_vit)} tensors ({time.time()-t0:.1f}s, "
          f"dropped {n_vit} ViT tensors not needed for generation)")

    vae_path = ensure_vae_weights(repo)
    t0 = time.time()
    vae = Wan2_2_VAE(
        z_dim=48, c_dim=160, dim_mult=(1, 2, 4, 4),
        temperal_downsample=(False, True, True),
    )
    vae.model.load_weights(list(mx.load(str(vae_path)).items()), strict=True)
    _materialize(vae.model.parameters())
    print(f"[ok] VAE strict load from {vae_path.name} ({time.time()-t0:.1f}s)")

    tok = AutoTokenizer.from_pretrained(str(repo))
    ids = tok(args.prompt, add_special_tokens=False, return_tensors="np").input_ids[0]
    text_ids = mx.array(ids, dtype=mx.int32)

    def tok_id(s: str) -> int:
        out = tok.convert_tokens_to_ids(s)
        if out is None or out == tok.unk_token_id:
            raise RuntimeError(f"special token {s!r} not found in tokenizer")
        return out

    special_token_ids = {
        "bos":              tok_id("<|im_start|>"),
        "eos":              tok_id("<|im_end|>"),
        "start_of_image":   tok_id("<|vision_start|>"),
        "end_of_image":     tok_id("<|vision_end|>"),
        "image_token_id":   cfg_json["qwen2_5_vl_config"]["image_token_id"],
    }
    print(f"[ok] tokenized: {len(ids)} prompt tokens")

    H_lat = args.size // lance_cfg.vae_downsample_spatial
    W_lat = args.size // lance_cfg.vae_downsample_spatial
    latent_shape = (args.t_lat, H_lat, W_lat)
    print(f"\nRunning {args.steps}-step denoising loop ...")
    t0 = time.time()
    final_latent = model.sample_t2i(
        prompt_token_ids=text_ids,
        latent_shape=latent_shape,
        special_token_ids=special_token_ids,
        num_steps=args.steps,
        timestep_shift=lance_cfg.timestep_shift,
        seed=args.seed,
        cfg_scale=args.cfg,
    )
    _materialize(final_latent)
    sample_dt = time.time() - t0
    print(f"[ok] sampled. latent {final_latent.shape}  "
          f"({sample_dt:.1f}s, {sample_dt/args.steps*1000:.0f} ms/step)")

    print("Decoding through VAE (streaming for T>1) ...")
    t0 = time.time()
    video = vae.decode(final_latent)
    _materialize(video)
    print(f"[ok] VAE decode ({time.time()-t0:.1f}s)  shape={video.shape}")

    video_np = np.asarray(video).squeeze(0)  # (T, H, W, 3)
    save_outputs(video_np, Path(args.out), want_mp4=args.mp4, fps=args.fps)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt", default="a calm ocean wave rolling onto a sandy beach")
    ap.add_argument("--out", default="output.png")
    ap.add_argument("--steps", type=int, default=24)
    ap.add_argument("--size", type=int, default=256,
                    help="square frame size (256 recommended for T2V)")
    ap.add_argument("--t-lat", type=int, default=3,
                    help="latent frame count; output frames = (t_lat-1)*4 + 1")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--cfg", type=float, default=4.0)
    ap.add_argument("--mp4", action="store_true",
                    help="also export an MP4 (needs `pip install 'imageio[ffmpeg]'`)")
    ap.add_argument("--fps", type=int, default=8,
                    help="MP4 frame rate")
    main(ap.parse_args())