"""Compare Wan VAE and TAEHV decode on saved FastWan latents."""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import numpy as np

from examples.inference.basic.mlx_wan_prompt_to_video import DEFAULT_MODEL_ROOT, decode_latents_to_video


def _torch_mps_memory() -> dict[str, int | None]:
    """
    Report current and recommended memory usage for the MPS backend.

    Returns:
        dict[str, int | None]: Memory metrics in bytes, or `None` values when
            PyTorch or MPS is unavailable.
    """
    try:
        import torch
    except ImportError:
        return {
            "current_allocated_bytes": None,
            "driver_allocated_bytes": None,
            "recommended_max_bytes": None,
        }
    if not torch.backends.mps.is_available():
        return {
            "current_allocated_bytes": None,
            "driver_allocated_bytes": None,
            "recommended_max_bytes": None,
        }
    return {
        "current_allocated_bytes": int(torch.mps.current_allocated_memory()),
        "driver_allocated_bytes": int(torch.mps.driver_allocated_memory()),
        "recommended_max_bytes": int(torch.mps.recommended_max_memory()),
    }


def _parse_backends(raw: str) -> list[str]:
    """
    Parse and validate a comma-separated list of decoding backends.

    Parameters:
        raw (str): Comma-separated backend names.

    Returns:
        list[str]: Trimmed, supported backend names in input order.

    Raises:
        ValueError: If the input contains an unsupported backend.
    """
    backends = [backend.strip() for backend in raw.split(",") if backend.strip()]
    allowed = {"wan-vae", "taehv"}
    unknown = sorted(set(backends) - allowed)
    if unknown:
        raise ValueError(f"Unsupported decode backends: {unknown}")
    return backends


def main() -> None:
    """
    Benchmark selected Wan latent decoding backends and record their performance metrics.

    Loads the specified latent array, decodes it with each selected backend, exports the
    results as MP4 files, and writes per-backend timing and Torch MPS memory metrics to
    `metrics.json`.
    """
    parser = argparse.ArgumentParser(description="Benchmark decode backends on saved Wan/FastWan latents.")
    parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT)
    parser.add_argument("--latents-path", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_decode_benchmark"))
    parser.add_argument("--backends", default="wan-vae,taehv")
    parser.add_argument("--fps", type=int, default=16)
    parser.add_argument("--torch-device", default="auto")
    parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16")
    parser.add_argument("--taehv-source-path", type=Path, default=None)
    parser.add_argument("--taehv-checkpoint-path", type=Path, default=None)
    parser.add_argument("--taehv-parallel", action="store_true")
    args = parser.parse_args()

    latents = np.load(args.latents_path)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    rows = []

    for backend in _parse_backends(args.backends):
        print(f"=== Decode backend: {backend} ===")
        before = _torch_mps_memory()
        start = time.perf_counter()
        output_path = args.output_dir / f"{args.latents_path.stem}_{backend}.mp4"
        decode_latents_to_video(
            model_root=args.model_root,
            latents_np=latents,
            output_path=output_path,
            fps=args.fps,
            device_arg=args.torch_device,
            dtype_arg=args.torch_dtype,
            backend=backend,
            taehv_source_path=args.taehv_source_path,
            taehv_checkpoint_path=args.taehv_checkpoint_path,
            taehv_parallel=args.taehv_parallel,
        )
        elapsed = time.perf_counter() - start
        after = _torch_mps_memory()
        metrics = {
            "backend": backend,
            "latents_path": str(args.latents_path),
            "latents_shape": list(latents.shape),
            "decode_export_s": elapsed,
            "torch_mps_current_before_bytes": before["current_allocated_bytes"],
            "torch_mps_current_after_bytes": after["current_allocated_bytes"],
            "torch_mps_driver_before_bytes": before["driver_allocated_bytes"],
            "torch_mps_driver_after_bytes": after["driver_allocated_bytes"],
            "torch_mps_recommended_max_bytes": after["recommended_max_bytes"],
            "output_path": str(output_path),
        }
        rows.append(metrics)
        print(json.dumps(metrics, indent=2))

    metrics_path = args.output_dir / "metrics.json"
    metrics_path.write_text(json.dumps(rows, indent=2))
    print(f"Wrote decode metrics to: {metrics_path}")


if __name__ == "__main__":
    main()