"""Load every model once.

Warmup dominates a single run — roughly 71s of graph build and weight fault against
~17s of actual compute for the structure stage alone. So anything serving more than
one job (the MODELBEAST operator, a batch script) must build this ONCE and keep it,
never fork per job. The trellis-2 lane on this fleet shows the same shape: 47.9s cold
against 2.5s warm pipeline load.
"""

from __future__ import annotations

import json
from pathlib import Path

CKPTS = Path("/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts")
PIPELINE_JSON = CKPTS.parent / "pipeline.json"
REPO_WEIGHTS = Path(__file__).resolve().parents[1] / "weights"

FILES = {
    "ss_flow": "ss_flow_img_dit_1_3B_64_bf16",
    "ss_dec": "ss_dec_conv3d_16l8_fp16",
    "slat_512": "slat_flow_img2shape_dit_1_3B_512_bf16",
    "slat_1024": "slat_flow_img2shape_dit_1_3B_1024_bf16",
    "slat_tex": "slat_flow_imgshape2tex_dit_1_3B_1024_bf16",
    "shape_dec": "shape_dec_next_dc_f16c32_fp16",
    "tex_dec": "tex_dec_next_dc_f16c32_fp16",
}


def _weights(stem: str) -> Path:
    """Decoders were converted into the repo; flows pass through untouched."""
    local = REPO_WEIGHTS / f"{stem}.safetensors"
    return local if local.exists() else CKPTS / f"{stem}.safetensors"


def normalization(kind: str = "shape") -> dict:
    cfg = json.loads(PIPELINE_JSON.read_text())
    return cfg.get("args", cfg)[f"{kind}_slat_normalization"]


def load_all(device: str | None = None, with_texture: bool = False) -> dict:
    """Every model plus both conditioners. ~24GB of weights; MLX loads them lazily."""
    from . import slat_flow, ss_dec, ss_flow
    from .cond import ProjConditioner
    from .decoders import load as load_dec

    def cfg(stem):
        return CKPTS / f"{stem}.json"

    flows = [("ss_flow", ss_flow.load), ("ss_dec", ss_dec.load),
             ("slat_512", slat_flow.load), ("slat_1024", slat_flow.load)]
    if with_texture:
        flows.append(("slat_tex", slat_flow.load))

    models = {}
    for key, loader in flows:
        stem = FILES[key]
        model, rep = loader(_weights(stem), cfg(stem))
        if rep["missing"] or rep["unmapped"]:
            raise RuntimeError(f"{key}: {len(rep['missing'])} missing, "
                               f"{len(rep['unmapped'])} unmapped")
        models[key] = model

    for key in (["shape_dec", "tex_dec"] if with_texture else ["shape_dec"]):
        stem = FILES[key]
        model, rep = load_dec(_weights(stem), cfg(stem))
        if rep["missing"] or rep["unmapped"]:
            raise RuntimeError(f"{key}: incomplete weight mapping")
        models[key] = model

    models["cond_512"] = ProjConditioner("shape_512", device=device)
    models["cond_1024"] = ProjConditioner("shape_1024", device=device)
    if with_texture:
        # tex_1024 differs from shape_1024 only in naf_target_size (1024 vs 512), but
        # that changes the high-res branch it samples, so it needs its own conditioner
        models["cond_tex"] = ProjConditioner("tex_1024", device=device)
    return models