"""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 128-token VSA tiles run on the
bf16 Triton emulator in both variants (the FP4 sparse-attention kernel is absent
from the published fastvideo-kernel wheel), so this isolates FP8 vs NVFP4 linears
+ decoder on ZeroGPU's RTX PRO 6000.

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 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 + bf16 VSA emulator; offload overridden below
ADALN_SUBPATH = os.getenv("FASTH3_ADALN_SUBPATH", "adaln_table.pt")


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)


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 + bf16 VSA emulator 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=True),
            sequential_load=False,
        )
        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)

    def generate(self, prompt, references, height, width, num_frames, seed) -> str:
        if self.generator is None:
            raise RuntimeError("engine not initialized")
        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()
        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