"""Request-scoped H3 stages for one Space, with no startup model residency.

The app owns one GPU decorator and serializes requests. Inside that allocation,
enter StageStore.stage, run the matching export, then exit only after the export
returns and releases its model. Cold downloads consume allocation time. Each
export loads fresh weights and keeps no pipeline.
"""
from __future__ import annotations

from contextlib import contextmanager
import faulthandler
import os
from pathlib import Path
import shutil
import sys
import tempfile
import threading
from threading import Lock
import time

MODEL_REPO = "MiniMaxAI/MiniMax-H3"
MODEL_REVISION = "42ed227ee7df40d41602854ae760620d6eb651fe"
# Pinned inventory: conditioner ~62.166 GiB, generator ~71.993 GiB, and the
# Ref2VA generator (transformer_ref partition) the same size as the generator.
# Round upward, then reserve extra space for downloads, adapter and outputs.
STAGE_BYTES = {"conditioner": 63 * 1024**3, "generator": 72 * 1024**3, "generator_ref": 72 * 1024**3}
STORAGE_MARGIN_BYTES = 8 * 1024**3
STAGE_PATTERNS = {
    "conditioner": ("modular_model_index.json", "model_index.json", "text_encoder/*", "processor/*", "tokenizer/*"),
    "generator": ("modular_model_index.json", "model_index.json", "transformer/*", "vae/*", "audio_vae/*", "scheduler/*", "audio_scheduler/*"),
    "generator_ref": ("modular_model_index.json", "model_index.json", "transformer_ref/*", "vae/*", "audio_vae/*", "scheduler/*", "audio_scheduler/*"),
}
# huggingface_hub 1.x shares one httpx client created with timeout=None, so a
# stuck Hub request never returns. A second staged run in one Space process sat
# silent after its stage marker until ZeroGPU aborted the worker (2026-10-22).
# Every download therefore starts a fresh session with finite timeouts, retries
# once after a transport failure, and reports progress. Every thread's stack is
# dumped if the stage's filesystem prelude and download together are still
# running after DOWNLOAD_STALL_DUMP_S.
HUB_HTTP_TIMEOUT_S = 60
HUB_CONNECT_TIMEOUT_S = 20
DOWNLOAD_ATTEMPTS = 2
DOWNLOAD_HEARTBEAT_S = 30
DOWNLOAD_STALL_DUMP_S = 180
# Loader kinds: Ref2VA reuses the conditioner partition with its own blocks.
_BLOCKS = {
    "conditioner": ("MiniMaxH3ConditionerBlocks", ("text_encoder",)),
    "generator": ("MiniMaxH3GeneratorBlocks", ("transformer", "vae", "audio_vae")),
    "conditioner_ref": ("MiniMaxH3Ref2VAConditionerBlocks", ("text_encoder",)),
    "generator_ref": ("MiniMaxH3Ref2VAGeneratorBlocks", ("transformer_ref", "vae", "audio_vae")),
}


def _bounded_hub_client():
    import httpx
    from huggingface_hub.utils._http import default_client_factory

    client = default_client_factory()
    client.timeout = httpx.Timeout(HUB_HTTP_TIMEOUT_S, connect=HUB_CONNECT_TIMEOUT_S)
    return client


def _transient(error):
    """True when a network transport failure caused the download error."""
    import httpx

    while error is not None:
        if isinstance(error, httpx.TransportError):
            return True
        error = error.__cause__
    return False


def _tree_bytes(path):
    total = 0
    for folder, _, files in os.walk(path):
        for name in files:
            try:
                total += os.lstat(os.path.join(folder, name)).st_size
            except OSError:
                pass
    return total


def _note(line):
    os.write(1, (line + "\n").encode())  # bypasses sys.stdout and its lock


@contextmanager
def _stall_dump(seconds):
    """Dump every thread's stack if the block is still running after `seconds`."""
    faulthandler.dump_traceback_later(seconds, file=sys.__stderr__)
    try:
        yield
    finally:
        faulthandler.cancel_dump_traceback_later()


@contextmanager
def _heartbeat(kind, path):
    """Log the stage directory's size every DOWNLOAD_HEARTBEAT_S while downloading."""
    stop = threading.Event()
    started = time.monotonic()

    def beat():
        while not stop.wait(DOWNLOAD_HEARTBEAT_S):
            _note(f"stage_download kind={kind} elapsed_s={time.monotonic() - started:.0f} "
                  f"bytes={_tree_bytes(path)}")

    threading.Thread(target=beat, name=f"stage-heartbeat-{kind}", daemon=True).start()
    try:
        yield
    finally:
        stop.set()


class StageStore:
    """Download one pinned partition into an owned temporary directory.

    Construction does no I/O. Reuse one store per app; simultaneous/nested stages
    on that store fail closed. The root and unrelated contents are never removed.
    Parent code may download a pinned adapter under the yielded generator path.
    """

    def __init__(self, root, model_repo, model_revision):
        if (model_repo, model_revision) != (MODEL_REPO, MODEL_REVISION):
            raise ValueError("unsupported pinned model; storage inventory must match the exact revision")
        self.root = Path(root)
        self.model_repo = model_repo
        self.model_revision = model_revision
        self._lock = Lock()

    @contextmanager
    def stage(self, kind):
        if kind not in STAGE_PATTERNS:
            raise ValueError("stage kind must be conditioner, generator or generator_ref")
        if not self._lock.acquire(blocking=False):
            raise RuntimeError("stage_already_active")
        owned = identity = None
        try:
            # The watchdog covers the filesystem prelude as well as the download.
            with _stall_dump(DOWNLOAD_STALL_DUMP_S):
                self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
                required = STAGE_BYTES[kind] + STORAGE_MARGIN_BYTES
                free = shutil.disk_usage(self.root).free
                if free < required:
                    raise RuntimeError(f"insufficient_stage_storage: {kind} requires {required} free bytes")
                owned = Path(tempfile.mkdtemp(prefix=f"h3-{kind}-", dir=self.root))
                identity = owned.lstat()
                _note(f"stage_prelude kind={kind} free_gib={free / 1024**3:.0f}")
                import huggingface_hub

                for attempt in range(1, DOWNLOAD_ATTEMPTS + 1):
                    # Replacing the factory also closes the shared client's pooled connections.
                    huggingface_hub.set_client_factory(_bounded_hub_client)
                    try:
                        with _heartbeat(kind, owned):
                            huggingface_hub.snapshot_download(
                                self.model_repo, revision=self.model_revision, token=False,
                                allow_patterns=list(STAGE_PATTERNS[kind]), local_dir=str(owned), max_workers=2,
                            )
                        break
                    except Exception as error:
                        if attempt == DOWNLOAD_ATTEMPTS or not _transient(error):
                            raise
                        print(f"stage_download_retry kind={kind} error={type(error).__name__}", flush=True)
            yield str(owned)
        finally:
            try:
                if owned is not None:
                    current = owned.lstat()
                    if identity is None or (current.st_dev, current.st_ino) != (identity.st_dev, identity.st_ino):
                        raise RuntimeError("owned_stage_replaced; refusing cleanup")
                    shutil.rmtree(owned)
            finally:
                self._lock.release()


def _open_image(path):
    from PIL import Image, ImageOps

    if not path:
        return None
    with Image.open(path) as source:
        return ImageOps.exif_transpose(source).convert("RGB")


def _load_pipeline(kind, model_path):
    if not Path(model_path).is_dir():
        raise ValueError("local_model_path_required")
    model_path = Path(model_path).resolve()
    import torch
    from diffusers import ComponentsManager
    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
    import h3_split_blocks

    MiniMaxH3ModularPipeline.min_duration = property(lambda self: 2.0)
    blocks_name, models = _BLOCKS[kind]
    pipe = getattr(h3_split_blocks, blocks_name)().init_pipeline(
        str(model_path), components_manager=ComponentsManager(), collection="workflow-studio",
    )
    # The H3 loader keeps both VAEs float32 via _keep_in_fp32_modules. Do not
    # blanket-cast the loaded pipeline when placing it on CUDA.
    # ModularPipeline interprets this outer dict as per-component kwargs. Each
    # model receives device_map="cuda", avoiding a complete CPU-resident stage
    # before transfer. This runs only inside the app's real GPU allocation.
    # The model's modular index contains Hub repository references. Override
    # their loading root explicitly; init_pipeline(local_path) does not rebase it.
    pipe.load_components(dtype=torch.bfloat16, local_files_only=True, token=False,
                         device_map={name: "cuda" for name in models}, low_cpu_mem_usage=True,
                         pretrained_model_name_or_path=str(model_path))
    return pipe


def _clear_failure_frames(error):
    """Release models held by chained loader errors, without losing traceback locations."""
    import traceback

    pending, seen = [error], set()
    while pending:
        current = pending.pop()
        if id(current) in seen:
            continue
        seen.add(id(current))
        traceback.clear_frames(current.__traceback__)
        pending.extend(link for link in (current.__cause__, current.__context__) if link is not None)


def encode_local(plan, image_path, last_image_path, model_path):
    """GPU-call body returning (CPU embeds, CPU tags, height, width, frames).

    The app calls this inside its whole-workflow allocation. No Hub access or pipeline retention occurs
    here. FL2VA conditioning resolves canvas geometry but does not expose a frame
    output, so its concrete, already-aligned input frame count is retained.
    """
    import gc
    import torch

    pipe = state = None
    try:
        pipe = _load_pipeline("conditioner", model_path)
        pipe.to("cuda")
        with torch.inference_mode():
            state = pipe(
                prompt=plan["prompt"], image=_open_image(image_path), last_image=_open_image(last_image_path),
                height=plan["height"], width=plan["width"], num_frames=plan["num_frames"],
            )
            height, width = state.get("height"), state.get("width")
            frames = state.get("num_frames")
            if frames is None:
                frames = plan["num_frames"]
            if (height, width, frames) != (plan["height"], plan["width"], plan["num_frames"]):
                raise ValueError("conditioning_geometry_mismatch")
            return (state.get("prompt_embeds").detach().cpu().contiguous(),
                    state.get("text_token_tags").detach().cpu().contiguous(),
                    int(height), int(width), int(frames))
    except BaseException as error:
        # A retained provider traceback otherwise keeps failed loader/call frames
        # and their model references alive across the next stage.
        _clear_failure_frames(error)
        raise
    finally:
        pipe = state = None
        gc.collect()
        torch.cuda.empty_cache()


def generate_local(plan, embeds, tags, image_path, last_image_path, model_path, adapter_path):
    """GPU-call body returning (frames, CPU contiguous audio, sampling_rate).

    adapter_path is a pinned safetensors file staged before this export, inside
    the whole-workflow allocation. The base checkpoint is freshly loaded on every invocation. LoRA hooks
    are scoped to the call and never fused into the base weights.
    """
    from copy import deepcopy
    import gc
    import torch
    from adapter_runtime import adapter_context, adapter_entries

    pipe = state = entries = adapter_state = None
    try:
        pipe = _load_pipeline("generator", model_path)
        pipe.vae.set_attention_backend("native")
        pipe.audio_vae.set_attention_backend("native")
        pipe.transformer.set_attention_backend("_native_cudnn")
        pipe.scheduler = type(pipe.scheduler).from_config(deepcopy(dict(pipe.scheduler.config)))
        pipe.audio_scheduler = type(pipe.audio_scheduler).from_config(deepcopy(dict(pipe.audio_scheduler.config)))
        pipe.to("cuda")
        entries, scale = [], 1.0
        if plan["adapter"] != "off":
            from h3_adapter_spec import PROFILES
            from safetensors.torch import load_file

            profile = PROFILES[plan["adapter"]]
            adapter_state = load_file(str(adapter_path), device="cpu")
            entries, scale = adapter_entries(profile, pipe.transformer, adapter_state)
            pipe.scheduler.set_shift(profile.video_shift)
            pipe.audio_scheduler.set_shift(profile.audio_shift)
        with torch.inference_mode(), adapter_context(pipe.transformer, entries, scale):
            state = pipe(
                prompt_embeds=embeds.to("cuda"), text_token_tags=tags,
                image=_open_image(image_path), last_image=_open_image(last_image_path),
                height=plan["height"], width=plan["width"], num_frames=plan["num_frames"],
                num_inference_steps=plan["steps"],
                generator=torch.Generator("cpu").manual_seed(plan["seed"]),
            )
            frames = state.get("videos")[0]
            if isinstance(frames, torch.Tensor):
                frames = frames.detach().cpu().contiguous()
            return frames, state.get("audio")[0].detach().cpu().contiguous(), state.get("sampling_rate")
    except BaseException as error:
        _clear_failure_frames(error)
        raise
    finally:
        pipe = state = entries = adapter_state = None
        gc.collect()
        torch.cuda.empty_cache()


def _typed_references(references):
    """(kind, path) pairs as the model's reference objects, in reading order."""
    from diffusers.modular_pipelines.minimax_h3 import (
        MiniMaxH3AudioReference, MiniMaxH3ImageReference, MiniMaxH3VideoReference,
    )

    types = {"image": MiniMaxH3ImageReference, "audio": MiniMaxH3AudioReference,
             "video": MiniMaxH3VideoReference}
    return [types[kind].from_file(path) for kind, path in references]


def encode_local_ref2va(plan, references, model_path, image_short_edge):
    """Ref2VA conditioning: a vision block per reference ahead of the prompt.

    Both Ref2VA halves resize reference images in a shared setup step (default
    2048 px short edge), so this must match the generator's edge.
    Returns (CPU embeds, CPU tags, height, width, frames) like encode_local.
    """
    import gc
    import torch

    pipe = state = None
    try:
        pipe = _load_pipeline("conditioner_ref", model_path)
        pipe.register_to_config(reference_image_short_edge=int(image_short_edge))
        pipe.to("cuda")
        with torch.inference_mode():
            state = pipe(prompt=plan["prompt"], references=_typed_references(references),
                         height=plan["height"], width=plan["width"], num_frames=plan["num_frames"])
            height, width = state.get("height"), state.get("width")
            frames = state.get("num_frames")
            if frames is None:
                frames = plan["num_frames"]
            if (height, width, frames) != (plan["height"], plan["width"], plan["num_frames"]):
                raise ValueError("conditioning_geometry_mismatch")
            return (state.get("prompt_embeds").detach().cpu().contiguous(),
                    state.get("text_token_tags").detach().cpu().contiguous(),
                    int(height), int(width), int(frames))
    except BaseException as error:
        _clear_failure_frames(error)
        raise
    finally:
        pipe = state = None
        gc.collect()
        torch.cuda.empty_cache()


def generate_local_ref2va(plan, embeds, tags, references, model_path, image_short_edge):
    """Ref2VA generation on the transformer_ref partition; references are decoded here.

    Returns (frames, CPU contiguous audio, sampling_rate) like generate_local.
    """
    from copy import deepcopy
    import gc
    import torch

    pipe = state = None
    try:
        pipe = _load_pipeline("generator_ref", model_path)
        pipe.register_to_config(reference_image_short_edge=int(image_short_edge))
        pipe.vae.set_attention_backend("native")
        pipe.audio_vae.set_attention_backend("native")
        pipe.transformer_ref.set_attention_backend("_native_cudnn")
        pipe.scheduler = type(pipe.scheduler).from_config(deepcopy(dict(pipe.scheduler.config)))
        pipe.audio_scheduler = type(pipe.audio_scheduler).from_config(deepcopy(dict(pipe.audio_scheduler.config)))
        pipe.to("cuda")
        with torch.inference_mode():
            state = pipe(
                prompt_embeds=embeds.to("cuda"), text_token_tags=tags,
                references=_typed_references(references),
                height=plan["height"], width=plan["width"], num_frames=plan["num_frames"],
                num_inference_steps=plan["steps"],
                generator=torch.Generator("cpu").manual_seed(plan["seed"]),
            )
            frames = state.get("videos")[0]
            if isinstance(frames, torch.Tensor):
                frames = frames.detach().cpu().contiguous()
            return frames, state.get("audio")[0].detach().cpu().contiguous(), state.get("sampling_rate")
    except BaseException as error:
        _clear_failure_frames(error)
        raise
    finally:
        pipe = state = None
        gc.collect()
        torch.cuda.empty_cache()