"""Load aokit AOT-compiled block graphs onto a BriaFibo2Transformer2DModel, on ZeroGPU.

The transformer's `layers` ModuleList is heterogeneous: injection layers (ids in the model config)
carry gated cross-attention to the text encoder, the rest do not. So one graph cannot cover the whole
ModuleList (as in the Flux/klein case). Instead the package holds three weight-less graphs --
`plain`, `injection`, `refiner` -- and we patch each block with the matching one. Each patched block
keeps its own eager `forward` as a fallback, so a shape/version the graph was not built for still runs.

Package (built by job_fibo2_aoti.py, see the AOTI repo): config.json + package/submodules/<name>/package.pt2.
The compiled graphs are weight-less, so the SAME package serves dev and turbo (identical architecture).
"""

import json
import os
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download, snapshot_download

AOTI_REPO = os.environ.get("FIBO2_AOTI_REPO", "diffusers-internal-dev/Fibo-2.0-dev-aoti-sm120")
AOTI_ENABLED = os.environ.get("AOTI", "1") != "0"
HF_TOKEN = os.environ.get("HF_TOKEN")


def _with_eager_fallback(compiled, eager, label):
    warned = []

    def forward(*args, **kwargs):
        try:
            return compiled(*args, **kwargs)
        except Exception as e:  # shape/dtype the graph was not built for -> run eager
            if not warned:
                warned.append(True)
                print(f"AOTI {label} fell back to eager: {e!r}", flush=True)
            return eager(*args, **kwargs)

    return forward


def load_aoti(transformer):
    """Patch `transformer` in place with the compiled graphs. Raises if the package is
    unavailable or built for a different torch; the caller should fall back to eager."""
    from aokit import LazyAOTIModel, patch

    cfg = json.loads(Path(hf_hub_download(AOTI_REPO, "config.json", token=HF_TOKEN)).read_text())
    if cfg.get("torch") != torch.__version__:
        raise RuntimeError(f"AOTI built for torch {cfg.get('torch')}, runtime has {torch.__version__}")

    pkg = Path(snapshot_download(AOTI_REPO, allow_patterns="package/*", token=HF_TOKEN)) / "package"
    inj = set(cfg["injection_layer_ids"])
    models = {
        name: LazyAOTIModel(pkg / "submodules" / name / "package.pt2")
        for name in cfg["stacks"]
        if (pkg / "submodules" / name / "package.pt2").exists()
    }

    def bind(block, model, label):
        eager = block.forward
        patch(block, model)  # sets block.forward to the compiled graph (bound to block's own weights)
        block.forward = _with_eager_fallback(block.forward, eager, label)

    def bind_refiner(block, m_t2i, m_edit, label):
        # The refiner is called with a different arg set on t2i vs edit; try the t2i graph,
        # then the edit graph, then eager. Each graph is bound to this block's own weights.
        eager = block.forward
        patch(block, m_t2i)
        f_t2i = block.forward
        f_edit = m_edit.with_weights(block.state_dict())
        warned = []

        def forward(*args, **kwargs):
            try:
                return f_t2i(*args, **kwargs)
            except Exception:
                try:
                    return f_edit(*args, **kwargs)
                except Exception as e:
                    if not warned:
                        warned.append(True)
                        print(f"AOTI {label} fell back to eager: {e!r}", flush=True)
                    return eager(*args, **kwargs)

        block.forward = forward

    if "plain" in models or "injection" in models:
        for i, block in enumerate(transformer.layers):
            model = models.get("injection") if i in inj else models.get("plain")
            if model is not None:
                bind(block, model, f"layers.{i}")
    if "refiner" in models:
        redit = models.get("refiner_edit")
        for block in transformer.noise_refiner:
            if redit is not None:
                bind_refiner(block, models["refiner"], redit, "noise_refiner")
            else:
                bind(block, models["refiner"], "noise_refiner")

    print(f"AOTI loaded: {list(models)} from {AOTI_REPO}", flush=True)