"""In-process FastH3 OmniRef engine — FP8 variant (benchmark against NVFP4).

Same Ref2VA pipeline as the NVFP4 Space, but the DiT linears are FP8 W8A8
(FastH3-OmniRef-FP8) instead of packed NVFP4. The sparse attention runs on
FastVideo's Triton tile-64 block-sparse kernel instead of the profile's bf16
tile-128 emulator.

FP8 specifics:
  - FastVideo's `rtx-4090` profile carries the FP8 numerics (FP8 linears, bf16
    VSA emulator, dense light VAE), but its offload policy streams the DiT
    layerwise for a 24 GB card. On the 96 GB ZeroGPU slice we override the
    offload to keep everything resident.
  - The FP8 repo ships `adaln_table.pt`; using it drops 24 GB of AdaLN
    projections from the load.
  - The text encoder is still the NVFP4 Qwen3-VL encoder (shared across the
    consumer releases), so the one-time flashinfer FP4 JIT still runs for it.
"""
from __future__ import annotations

import contextlib
import os
import tempfile
import threading
import time
from pathlib import Path

MODEL_REPO = os.getenv("FASTH3_MODEL_REPO", "FastVideo/FastH3-OmniRef-FP8")
PROFILE_NAME = "rtx-4090"  # FP8 numerics; offload and fine attention overridden below
ADALN_SUBPATH = os.getenv("FASTH3_ADALN_SUBPATH", "adaln_table.pt")
FINE_ATTENTION = os.getenv("FASTH3_FINE_ATTENTION", "triton64")
TILE64_CONFIG = (5, 4)
KERNEL_CACHE_DIR = os.getenv("FASTH3_KERNEL_CACHE_DIR", "/tmp/fasth3-kernels")
KERNEL_CACHE_REPO = os.getenv("FASTH3_KERNEL_CACHE_REPO", "multimodalart/fasth3-omni-kernel-cache")


def kernel_cache_env() -> None:
    os.environ.setdefault("TRITON_CACHE_DIR", f"{KERNEL_CACHE_DIR}/triton")
    os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", f"{KERNEL_CACHE_DIR}/inductor")
    os.environ.setdefault("FLASHINFER_WORKSPACE_BASE", KERNEL_CACHE_DIR)
    os.environ.setdefault("CUDA_CACHE_PATH", f"{KERNEL_CACHE_DIR}/nv")
    os.environ.setdefault("CUDA_CACHE_MAXSIZE", str(4 << 30))
    os.environ.setdefault("PYTORCH_KERNEL_CACHE_PATH", f"{KERNEL_CACHE_DIR}/torch")
    os.makedirs(os.environ["PYTORCH_KERNEL_CACHE_PATH"], exist_ok=True)


def kernel_cache_name() -> str:
    from importlib.metadata import version

    packages = ("torch", "triton", "flashinfer-python", "fastvideo", "fastvideo-kernel")
    return "-".join(version(p) for p in packages).replace("+", "_") + f"-{FINE_ATTENTION}.tar"


def fetch_kernel_cache() -> bool:
    import tarfile

    from huggingface_hub import hf_hub_download

    t0 = time.time()
    try:
        archive = hf_hub_download(KERNEL_CACHE_REPO, kernel_cache_name(), repo_type="dataset", token=False)
        with tarfile.open(archive) as tar:
            tar.extractall(KERNEL_CACHE_DIR, filter="tar")
    except Exception as error:
        print(f"[fasth3-fp8] no kernel cache ({type(error).__name__}); compiling on first use", flush=True)
        return False
    print(f"[fasth3-fp8] kernel cache {kernel_cache_name()} restored in {time.time() - t0:.1f}s", flush=True)
    return True


kernel_cache_env()


def _hf_token() -> str | None:
    return os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN") or None


def _patch_mig_device_ids() -> None:
    """ZeroGPU exposes a MIG slice via a UUID in CUDA_VISIBLE_DEVICES; FastVideo
    does int() on it and crashes. Fall back to the logical index for UUIDs."""
    import fastvideo.platforms.cuda as cuda

    if getattr(cuda, "_mig_patched", False):
        return
    orig = cuda.device_id_to_physical_device_id

    def safe(device_id: int = 0) -> int:
        try:
            return orig(device_id)
        except (ValueError, TypeError):
            return device_id

    cuda.device_id_to_physical_device_id = safe
    cuda._mig_patched = True
    print("[fasth3-fp8] patched device_id_to_physical_device_id for MIG UUID", flush=True)


_CUDA_ORIGINALS: dict = {}
_DEFERRED_VALIDATIONS: list = []


def _is_cpu(device) -> bool:
    return device is None or str(device) == "cpu"


class _CapturedSlice:
    def __init__(self, inner, device) -> None:
        self._inner = inner
        self._device = device

    def __getitem__(self, key):
        tensor = self._inner[key]
        return tensor if _is_cpu(self._device) else tensor.to(self._device)

    def __getattr__(self, name):
        return getattr(self._inner, name)


def _captured_safe_open(original):
    class CapturedSafeOpen:
        def __init__(self, filename, framework="pt", device="cpu", **kwargs) -> None:
            self._inner = original(filename, framework=framework, device="cpu", **kwargs)
            self._device = device

        def __enter__(self):
            self._inner.__enter__()
            return self

        def __exit__(self, *exc):
            return self._inner.__exit__(*exc)

        def get_tensor(self, name):
            tensor = self._inner.get_tensor(name)
            return tensor if _is_cpu(self._device) else tensor.to(self._device)

        def get_tensors(self):
            tensors = self._inner.get_tensors()
            return tensors if _is_cpu(self._device) else {k: v.to(self._device) for k, v in tensors.items()}

        def get_slice(self, name):
            return _CapturedSlice(self._inner.get_slice(name), self._device)

        def __getattr__(self, name):
            return getattr(self._inner, name)

    return CapturedSafeOpen


def zerogpu_main_process_shims() -> None:
    """Let FastVideo's in-process worker initialize under ZeroGPU's CUDA emulation.

    `spaces` emulates CUDA in the main process but leaves `set_device` and
    `mem_get_info` pointing at the real runtime, which has no device there, and
    NVML sees no GPU, so FastVideo's platform probe would fall back to MPS.
    safetensors' `safe_open(device="cuda")` allocates on the real runtime, so files
    are read on CPU and moved with `.to()`, which ZeroGPU captures and packs.
    FastVideo probes CU_DEVICE_ATTRIBUTE_INTEGRATED with a ctypes cuInit; a failed
    cuInit in the GPU-less parent is inherited by every forked GPU worker, so the
    answer (a discrete card) is pinned instead. The NVFP4 text encoder's runtime
    check imports flashinfer, whose CuTe DSL loads the CUTLASS runtime against the
    driver, so that check (and its FP4 op registration) runs in the first GPU call.
    NCCL refuses to start without a GPU; with one rank every FastVideo collective
    short-circuits, so a gloo group is equivalent, and a one-rank barrier (which
    torch would otherwise route through CUDA device selection) is skipped.
    """
    import torch

    if _CUDA_ORIGINALS:
        return
    import fastvideo.platforms as platforms
    from fastvideo.platforms.cuda import NonNvmlCudaPlatform

    platforms._current_platform = NonNvmlCudaPlatform()

    import safetensors
    import safetensors.torch

    captured = _captured_safe_open(safetensors.safe_open)
    safetensors.safe_open = captured
    safetensors.torch.safe_open = captured
    import importlib

    for module_name in (
        "fastvideo.models.loader.weight_utils",
        "fastvideo.models.loader.component_loader",
        "fastvideo.models.loader.lora_patch",
        "fastvideo.models.loader.text_encoder_quantization",
    ):
        importlib.import_module(module_name).safe_open = captured
    _CUDA_ORIGINALS["set_device"] = torch.cuda.set_device
    _CUDA_ORIGINALS["mem_get_info"] = torch.cuda.mem_get_info
    total = torch.cuda.get_device_properties(0).total_memory
    torch.cuda.set_device = lambda *args, **kwargs: None
    torch.cuda.mem_get_info = lambda *args, **kwargs: (total, total)

    import torch.distributed as dist

    def gloo_instead_of_nccl(function):
        def wrapped(*args, backend=None, **kwargs):
            return function(*args, backend="gloo" if backend in (None, "nccl") else backend, **kwargs)

        return wrapped

    dist.init_process_group = gloo_instead_of_nccl(dist.init_process_group)
    dist.new_group = gloo_instead_of_nccl(dist.new_group)
    single_rank_barrier = dist.barrier

    def barrier(*args, **kwargs):
        if dist.is_initialized() and dist.get_world_size() == 1:
            return None
        return single_rank_barrier(*args, **kwargs)

    dist.barrier = barrier

    import fastvideo.platforms.cuda as fastvideo_cuda

    fastvideo_cuda._cuda_driver_device_is_integrated = lambda device_id: False

    from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import MiniMaxH3SerializedNVFP4Config

    validate_runtime = MiniMaxH3SerializedNVFP4Config.validate_runtime

    def deferred_validate_runtime(self, device):
        _DEFERRED_VALIDATIONS.append((validate_runtime, self, device))

    MiniMaxH3SerializedNVFP4Config.validate_runtime = deferred_validate_runtime


def restore_cuda() -> None:
    """Inside a GPU call the device is real: put the runtime functions back."""
    import torch

    for name, function in _CUDA_ORIGINALS.items():
        setattr(torch.cuda, name, function)
    while _DEFERRED_VALIDATIONS:
        validate, config, device = _DEFERRED_VALIDATIONS.pop(0)
        validate(config, device)


def _find_pipeline():
    import gc

    from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3BasePipeline

    found = [o for o in gc.get_objects() if isinstance(o, MiniMaxH3BasePipeline)]
    return found[0] if len(found) == 1 else None


def _reuse_built_flashinfer_modules() -> None:
    from flashinfer.jit.core import JitSpecNvcc

    build = JitSpecNvcc.build
    if getattr(build, "reuses_built", False):
        return

    def build_unless_built(self, *args, **kwargs):
        if self.jit_library_path.exists():
            return None
        return build(self, *args, **kwargs)

    build_unless_built.reuses_built = True
    JitSpecNvcc.build = build_unless_built


def _ensure_tile64_kernel() -> None:
    if FINE_ATTENTION != "triton64":
        return
    import fastvideo.attention.backends.video_sparse_attn_h3 as vsa

    if vsa.block_sparse_attn_64_bhsd is None:
        from fastvideo_kernel.block_sparse_attn import block_sparse_attn

        vsa.block_sparse_attn_64_bhsd = block_sparse_attn
    from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import _attn_fwd_sparse

    pinned = [c for c in _attn_fwd_sparse.configs if (c.num_stages, c.num_warps) == TILE64_CONFIG]
    if pinned:
        _attn_fwd_sparse.configs = pinned


STAGE_PROGRESS = {
    "conditioning_stage": (0.02, "Encoding the prompt and references"),
    "latent_preparation_stage": (0.08, "Preparing latents"),
    "video_decoding_stage": (0.86, "Decoding video"),
    "audio_decoding_stage": (0.94, "Decoding audio"),
}


@contextlib.contextmanager
def _progress_hooks(pipeline, num_steps, on_progress):
    if on_progress is None or pipeline is None:
        yield
        return
    stages = getattr(pipeline, "_stage_name_mapping", {})
    transformer = getattr(stages.get("denoising_stage"), "transformer", None)
    step = 0
    patched = []

    def report(fraction, desc):
        try:
            on_progress(fraction, desc)
        except Exception as error:
            print(f"[fasth3-fp8] progress report failed: {error!r}", flush=True)

    def before_step(module, args, kwargs):
        nonlocal step
        step += 1
        report(0.1 + 0.75 * (step - 1) / num_steps, f"Denoising step {step} of {num_steps}")

    def wrap(name, stage):
        forward = stage.forward

        def staged(*args, **kwargs):
            if name in STAGE_PROGRESS:
                report(*STAGE_PROGRESS[name])
            out = forward(*args, **kwargs)
            if name == "audio_decoding_stage":
                report(0.97, "Muxing video and audio")
            return out

        stage.forward = staged
        patched.append(stage)

    for name, stage in stages.items():
        wrap(name, stage)
    handle = transformer.register_forward_pre_hook(before_step, with_kwargs=True) if transformer is not None else None
    try:
        yield
    finally:
        if handle is not None:
            handle.remove()
        for stage in patched:
            del stage.forward


def prefetch() -> str:
    from huggingface_hub import snapshot_download

    return snapshot_download(MODEL_REPO, token=_hf_token())


class FastH3OmniRefEngine:
    def __init__(self) -> None:
        self.generator = None
        self.num_steps = 8
        self._lock = threading.Lock()

    def initialize(self) -> None:
        if self.generator is not None:
            return
        with self._lock:
            if self.generator is not None:
                return
            self._build()

    def _build(self) -> None:
        import dataclasses

        import torch  # noqa: F401

        try:
            p = torch.cuda.get_device_properties(0)
            print(f"[fasth3-fp8] GPU: {p.name}, total VRAM {p.total_memory / 1e9:.1f} GB", flush=True)
        except Exception as e:  # pragma: no cover
            print(f"[fasth3-fp8] device props unavailable: {e!r}", flush=True)

        import fasth3_ref as ref

        _patch_mig_device_ids()

        snapshot = prefetch()
        adaln = os.path.join(snapshot, ADALN_SUBPATH)
        if not os.path.isfile(adaln):
            print(f"[fasth3-fp8] no adaln table at {adaln}; loading full AdaLN proj", flush=True)
            adaln = None

        composed = os.path.join(tempfile.gettempdir(), "fasth3_omniref_fp8_composed")
        model_dir = ref.compose_model_dir(Path(snapshot), Path(snapshot), Path(composed))
        contract = ref.load_contract(model_dir)
        self.num_steps = int(contract["num_inference_steps"])

        from fastvideo.api import OffloadConfig

        # FP8 numerics from rtx-4090, but resident on 96 GB.
        profile = ref.DEVICE_PROFILES[PROFILE_NAME]
        profile = dataclasses.replace(
            profile,
            offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=False,
                                  vae=False, pin_cpu_memory=False),
            sequential_load=False,
        )
        if FINE_ATTENTION == "triton64":
            env = {k: v for k, v in profile.env.items() if k != "FASTVIDEO_H3_VSA_FINE_ATTENTION"}
            profile = dataclasses.replace(profile, env=env)
            os.environ.pop("FASTVIDEO_H3_VSA_FINE_ATTENTION", None)
        ref.apply_profile(profile, None, adaln)  # dense decoder (no NVFP4 decoder on FP8)
        ref.apply_attention_prefix("default")
        print(f"[fasth3-fp8] profile: FP8 resident on {profile.label} numerics", flush=True)

        from fastvideo import VideoGenerator

        cfg = ref.build_generator_config(model_dir, num_gpus=1, lossless_accel=False, profile=profile)
        t0 = time.time()
        self.generator = VideoGenerator.from_config(cfg)
        print(f"[fasth3-fp8] generator ready in {time.time() - t0:.1f}s", flush=True)
        self.pipeline = _find_pipeline()
        print(f"[fasth3-fp8] fine attention: {FINE_ATTENTION}", flush=True)

    def generate(self, prompt, references, height, width, num_frames, seed, on_progress=None) -> str:
        if self.generator is None:
            raise RuntimeError("engine not initialized")
        _reuse_built_flashinfer_modules()
        restore_cuda()
        _ensure_tile64_kernel()
        if self.pipeline is not None and not self.pipeline.post_init_called:
            t0 = time.time()
            self.pipeline.post_init()
            print(f"[fasth3-fp8] pipeline stages ready in {time.time() - t0:.1f}s", flush=True)
        from fastvideo.api import GenerationRequest, InputConfig, OutputConfig, SamplingConfig
        from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference

        references = [MiniMaxH3Reference(source=path, media_type=kind) for kind, path in references]

        out_dir = tempfile.mkdtemp(prefix="fasth3_fp8_out_")
        out_path = os.path.join(out_dir, "fasth3_omniref.mp4")
        t0 = time.time()
        with _progress_hooks(self.pipeline, self.num_steps, on_progress):
            result = self.generator.generate(
                GenerationRequest(
                    prompt=prompt,
                    negative_prompt="",
                    inputs=InputConfig(references=references),
                    sampling=SamplingConfig(height=height, width=width, num_frames=num_frames, fps=24,
                                            num_inference_steps=self.num_steps, guidance_scale=1.0,
                                            batch_cfg=False, seed=seed),
                    output=OutputConfig(output_path=out_path, save_video=True, return_frames=False),
                )
            )
        dt = time.time() - t0
        path = getattr(result, "video_path", None) or out_path
        peak = getattr(result, "peak_memory_mb", None)
        print(f"[fasth3-fp8] generated {width}x{height} {num_frames}f in {dt:.1f}s"
              + (f" | peak {peak / 1024:.1f} GiB" if peak else ""), flush=True)
        return path