"""Wan 2.2 TI2V-5B text/image→video via mlx-video (native MLX, Apple Silicon).

Replaces the previous ComfyUI/PyTorch-MPS path, which produced structurally-valid
but visually garbage MP4s on every job from 2026-07-17 onward and still reported
success. Benchmarked 2026-08-02 at identical settings (832x480, 49f, 20 steps,
cfg 5.0, seed 1234): MLX q8 132s and clean, ComfyUI fp16 216s and garbage.

Same operator id, same params, same output contract (outdir/wan_{seed}.mp4) — the
queue and any downstream consumer see no change. Weights are the q8 MLX conversion
in vendor/mlx-video-models; the generator holds them in unified memory itself, so
there is no resident-server handshake and no PYTORCH_MPS_HIGH_WATERMARK_RATIO
tuning (MLX manages its own allocation).
"""
import argparse
import json
import os
import random
import subprocess
import sys
import time
from pathlib import Path

# job env has a minimal PATH; we shell out to ffprobe for output validation
os.environ["PATH"] = "/opt/homebrew/bin:/usr/local/bin:" + os.environ.get("PATH", "")

ROOT = Path(__file__).resolve().parents[3]
GEN = ROOT / "venvs" / "mlxvideo" / "bin" / "mlx_video.wan_2.generate"
MODEL_DIR = ROOT / "vendor" / "mlx-video-models" / "Wan2.2-TI2V-5B-mlx-q8"
FPS = 24  # Wan 2.2

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)
outdir = Path(a.outdir)

prompt = (p.get("prompt") or "").strip()
if not prompt:
    print("ERROR: prompt is required")
    sys.exit(1)
if not GEN.exists():
    print(f"ERROR: mlx-video not installed at {GEN}. Run scripts/install_mlx_video.sh")
    sys.exit(1)
if not (MODEL_DIR / "model.safetensors").exists():
    print(f"ERROR: MLX weights missing at {MODEL_DIR} — see manifest description")
    sys.exit(1)

length = int(p.get("length", 49))
if length % 4 != 1:
    length = (length // 4) * 4 + 1  # model requires 4n+1 frames
width, height = int(p.get("width", 832)), int(p.get("height", 480))
seed = int(p.get("seed", -1))
if seed < 0:
    seed = random.randint(0, 2**31 - 1)

start_image = a.input[0] if a.input else None
if start_image and not Path(start_image).exists():
    print(f"ERROR: input image not found: {start_image}")
    sys.exit(1)

out_mp4 = outdir / f"wan_{seed}.mp4"
cmd = [
    str(GEN),
    "--model-dir", str(MODEL_DIR),
    "--prompt", prompt,
    "--negative-prompt", p.get(
        "negative", "blurry, distorted, low quality, static image, watermark, text"),
    "--width", str(width),
    "--height", str(height),
    "--num-frames", str(length),
    "--steps", str(int(p.get("steps", 20))),
    "--guide-scale", str(float(p.get("cfg", 5.0))),
    "--seed", str(seed),
    "--output-path", str(out_mp4),
]
if start_image:
    cmd += ["--image", start_image]

print(f"seed={seed} mode={'i2v' if start_image else 't2v'} "
      f"{width}x{height}x{length}f backend=mlx-q8", flush=True)

t0 = time.time()
r = subprocess.run(cmd, capture_output=True, text=True)

# NOTE: do not gate on the exit code alone. Under memory pressure the generator
# has been observed finishing the render, writing a complete and valid MP4, and
# only then being SIGKILLed during interpreter teardown (rc=-9, "leaked
# semaphore" warning). Judge the artefact on disk first; the exit code is only
# authoritative when there is nothing usable to inspect.

# Echo the generator's own stage timings (T5 / denoise / VAE) into the job log.
import re
ANSI = re.compile(r"\x1b\[[0-9;]*m")  # generator colourises; job logs are plain text
for line in (r.stdout or "").splitlines():
    s = ANSI.sub("", line).strip()
    if any(k in s for k in ("T5 encoding", "Models loaded", "Denoising:", "VAE decode", "Total time")):
        if "it/s" not in s and "s/it" not in s:  # skip the tqdm bar
            print(f"  {s}", flush=True)

# --- output validation ---------------------------------------------------
# The old ComfyUI path marked six garbage jobs "done" because nothing here ever
# checked the result. These are structural checks only (a real content-quality
# check would need a perceptual model and would risk false failures on
# legitimately flat footage) — but they do catch the silent-empty-output case.
if not out_mp4.exists() or out_mp4.stat().st_size < 50_000:
    print(f"ERROR: no usable video produced at {out_mp4}")
    if r.returncode != 0:
        print(f"       mlx-video exited {r.returncode}")
        print((r.stderr or r.stdout)[-1200:])
    sys.exit(1)

probe = subprocess.run(
    ["ffprobe", "-v", "error", "-select_streams", "v:0", "-count_packets",
     "-show_entries", "stream=width,height,nb_read_packets", "-of", "csv=p=0", str(out_mp4)],
    capture_output=True, text=True)
if probe.returncode != 0:
    print(f"ERROR: output is not a decodable video: {probe.stderr[-400:]}")
    sys.exit(1)

try:
    gw, gh, gframes = (int(x) for x in probe.stdout.strip().split(",")[:3])
except ValueError:
    print(f"ERROR: could not parse ffprobe output: {probe.stdout!r}")
    sys.exit(1)

if (gw, gh) != (width, height):
    print(f"ERROR: dimension mismatch — asked {width}x{height}, got {gw}x{gh}")
    sys.exit(1)
if gframes < length:
    print(f"ERROR: truncated output — asked {length} frames, got {gframes}")
    sys.exit(1)

if r.returncode != 0:
    print(f"note: output validated OK despite generator exit {r.returncode} "
          f"(killed at teardown, render was complete)", flush=True)

print(f"done: {gframes} frames @ {gw}x{gh} → {out_mp4.name} "
      f"in {time.time() - t0:.1f}s", flush=True)