"""Model lifecycle: load at startup, report status, install decode hooks.

`load_models()` is invoked once at startup by `app.py:status()`. It populates
the `PIPE`, `MANAGER`, `LOADED_IN`, `TURBO_AVAILABLE`, `_TURBO_LoRA` globals
in `state.py` (everything else reads from there).
"""
from __future__ import annotations

import time

import spaces  # noqa: F401 — patches torch.cuda so the 72 GiB load can happen at startup

from config import (
    ATTENTION,
    CONDITIONER_SPACE,
    CHECKPOINT,
    MODEL_REPO,
    PLACEMENT,
    TURBO_LORA_FILE,
    TURBO_LORA_REPO,
)
from pipeline import lower_duration_floor
import state as _state
from state import _log


def status() -> str:
    if _state.LOAD_ERROR:
        return _state.LOAD_ERROR
    if _state.PIPE is None:
        return f"Loading `{CHECKPOINT}` transformer and `{MODEL_REPO}` VAEs. Watch the Space logs."
    import h3_aoti

    return (
        f"Ready · `{CHECKPOINT}` · **bf16 transformer / fp32 VAEs** · placement `{PLACEMENT}` · attention `{ATTENTION}` · "
        f"{h3_aoti.status() if CHECKPOINT == 'base' else 'AoTI off (curve checkpoint)'} · "
        f"loaded in {_state.LOADED_IN:.0f}s · conditioner `{CONDITIONER_SPACE}`"
        + (f" · turbo ready" if _state.TURBO_AVAILABLE else "")
    )


def load_models() -> str | None:
    """Load the denoising half at startup.

    `MiniMaxH3GeneratorBlocks` declares `transformer`, `vae`, `audio_vae`, the two schedulers and `video_processor`,
    so `load_components` fetches exactly those subfolders — `text_encoder/` and `transformer_ref/` are never touched.
    Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a bfloat16 audio VAE decodes
    the soundtrack roughly 20 dB too quiet.
    """
    if _state.PIPE is not None or _state.LOAD_ERROR is not None:
        return _state.LOAD_ERROR

    started = time.time()
    try:
        import torch
        from diffusers import ComponentsManager

        from h3_split_blocks import MiniMaxH3GeneratorBlocks

        lower_duration_floor()
        manager = ComponentsManager()
        blocks = MiniMaxH3GeneratorBlocks()
        _log(f"loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...")
        pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
        if CHECKPOINT != "base":
            from h3_checkpoint import load_transformer
            # Register first so the original transformer is never downloaded.
            pipe.update_components(transformer=load_transformer(
                MODEL_REPO, _log, merged_turbo=CHECKPOINT == "10eros-turbo-beta5"
            ))
        pipe.load_components(dtype=torch.bfloat16)
        pipe.transformer.set_attention_backend(ATTENTION)
        _log(f"components loaded in {time.time() - started:.0f}s")

        if TURBO_LORA_REPO and CHECKPOINT == "base":
            # Turbo weights (~850 MB) exceed the user custom-LoRA quota; load them
            # under a higher system limit. Failure must not abort the main model load.
            try:
                from safetensors.torch import load_file as _load_lora
                from lora.fetch import _fetch_lora_file

                _t0 = time.time()
                _log(f"loading turbo LoRA from {TURBO_LORA_REPO} ...")
                # System turbo file is ~851 MB; custom-LoRA cap is 512 MB by default.
                _lora_path = _fetch_lora_file(
                    TURBO_LORA_REPO, TURBO_LORA_FILE, max_bytes=2 * 1024 * 1024 * 1024
                )

                _lora_sd = _load_lora(_lora_path)

                _lora_pairs: dict[str, dict[str, object]] = {}
                for _k, _v in _lora_sd.items():
                    if ".lora_A." in _k:
                        _base = _k.split(".lora_A.")[0]
                        _lora_pairs.setdefault(_base, {})["A"] = _v
                    elif ".lora_B." in _k:
                        _base = _k.split(".lora_B.")[0]
                        _lora_pairs.setdefault(_base, {})["B"] = _v

                _state._TURBO_LoRA = []
                for _base_name, _pair in _lora_pairs.items():
                    if "A" not in _pair or "B" not in _pair:
                        continue
                    _base_name = _base_name.removeprefix("transformer.")
                    _parts = _base_name.split(".")
                    _obj = pipe.transformer
                    for _p in _parts:
                        _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
                    _A = _pair["A"].to(_obj.weight.dtype)
                    _B = _pair["B"].to(_obj.weight.dtype)
                    _state._TURBO_LoRA.append((_obj, _A, _B))

                del _lora_sd, _lora_pairs
                import gc as _gc; _gc.collect()
                from lora.apply import _validate_lora_shapes
                _validate_lora_shapes(_state._TURBO_LoRA)
                _state.TURBO_AVAILABLE = bool(_state._TURBO_LoRA)
                _log(f"turbo LoRA parsed ({len(_state._TURBO_LoRA)} modules in {time.time() - _t0:.0f}s, merge pending)")
                if _state._TURBO_LoRA:
                    _m, _A, _B = _state._TURBO_LoRA[0]
                    _log(f"  sample: module={type(_m).__name__}, "
                         f"A.shape={_A.shape}, B.shape={_B.shape}, "
                         f"A.device={_A.device}, dtype={_A.dtype}")
            except Exception as _turbo_e:
                _state.TURBO_AVAILABLE = False
                _state._TURBO_LoRA = None
                _log(f"turbo LoRA unavailable ({type(_turbo_e).__name__}: {_turbo_e}); continuing without it", "WARN")

        # Still startup, still free: an AoTI package carries no weights and opens its archive lazily inside the GPU
        # worker. Off unless `H3_AOTI=1`.
        import h3_aoti

        if CHECKPOINT == "base":
            h3_aoti.maybe_load(pipe.transformer)
        else:
            _log("AoTI disabled: the original package does not support curve-compressed AdaLN")

        if PLACEMENT == "pack":
            # Scoped to the transformer. `spaces` packs every startup-resident CUDA tensor into a second on-disk copy,
            # and packing all 77.3 GB busts the 150 GB storage quota; the 61.7 GB transformer alone fits. The ~10 GB of
            # fp32 VAEs move on the first GPU call instead.
            pipe.transformer.to("cuda")
            import gc as _gc
            _gc.collect()
            torch.cuda.empty_cache()

        if PLACEMENT == "offload":
            manager.enable_auto_cpu_offload(device="cuda")
            _arm_decode_hooks(pipe)

        _state.PIPE, _state.MANAGER = pipe, manager
        _state.LOADED_IN = time.time() - started
        _log(f"ready in {_state.LOADED_IN:.0f}s")
        from lora.fetch import _lora_cache_path
        if _lora_cache_path("probe", "probe.safetensors"):
            from state import _LORA_CACHE_DIR
            _log(f"LoRA bucket cache ready: {_LORA_CACHE_DIR}")
    except Exception as _e:
        _state.LOAD_ERROR = f"Loading {CHECKPOINT} with {MODEL_REPO} components failed: {type(_e).__name__}: {_e}"
        _log(f"load_models failed: {_state.LOAD_ERROR}", "ERROR")
    return _state.LOAD_ERROR


def _arm_decode_hooks(pipe):
    """Force every VAE decode through the .to("cpu") path even when placement=offload.

    `ComponentsManager.enable_auto_cpu_offload` toggles device on `__call__` boundaries only;
    this hook patches `decode` directly so the last `state.get("videos")` survives a worker recycle.
    """
    armed = False
    for module in (pipe.vae, getattr(pipe, "audio_vae", None)):
        if module is None or not hasattr(module, "decode"):
            continue
        orig_decode = module.decode
        def _patched_decode(self, *a, _orig=orig_decode, **k):
            _orig(*a, **k)
            try:
                self.to("cpu")
            except Exception:
                pass
        module.decode = _patched_decode.__get__(module, type(module))
        armed = True
    if armed:
        module.decode = armed