"""ZeroGPU AOTI for SenseNova-U1.5-8B-MoT.

One compiled `Qwen3DecoderLayer` gen-path package, reused by all 42 layers.

The per-denoising-step hot path is `language_model.model` (94.9% of wall time on
a 2048x2048 T2I), and inside it all the time is in the 42 identical
`Qwen3DecoderLayer`s. Their image-generation path reads the preallocated flash
KV buffers off the cache object as *attributes*, which `torch.export` cannot
trace, so `sensenova_u1.models.neo_unify.modeling_qwen3` adds a pure-tensor
restatement (`gen_block_core` / `Qwen3GenBlock` / `Qwen3DecoderLayer.gen_core`)
that takes the prefix K/V as explicit tensor arguments. Only `gen_core` is
compiled; the und/prefill paths stay eager.

A Space only calls `maybe_load(model)`; everything else is the build path.
"""

from __future__ import annotations

import json
import os
from pathlib import Path

AOTI = os.environ.get("SN_AOTI", "0") == "1"
AOTI_REPO = os.environ.get("SN_AOTI_REPO", "zerogpu-hacking/sensenova-u1-5-aoti")
AOTI_REPO_TYPE = os.environ.get("SN_AOTI_REPO_TYPE", "model")
# A package is valid for exactly one `<dtype>/torch<X.Y>/sm<cc>/<shape>`, and a
# mismatched one segfaults rather than raising, so `maybe_load` refuses anything
# that is not an exact match.
AOTI_KEY = os.environ.get("SN_AOTI_KEY", "bf16/torch2.8/sm120/dynamic")
AOTI_SHAPE = os.environ.get("SN_AOTI_SHAPE", "dynamic")
AOTI_DURATION = int(os.environ.get("SN_AOTI_DURATION", "1500"))

SUBMODULE = "layers"
METADATA_FILENAME = "metadata.json"
ALIASES_FILENAME = "constant_aliases.json"

# Image-token sequence length. The T2I aspect buckets run 1152x3456 .. 2880x1440
# at patch 16 x merge 2, i.e. 3888..4096 tokens; editing prefixes are longer, so
# leave generous headroom rather than recompiling per bucket.
S_MIN, S_MAX = int(os.environ.get("SN_AOTI_S_MIN", "256")), int(os.environ.get("SN_AOTI_S_MAX", "16384"))
# Text (or text+reference-image) prefix already in the KV cache.
P_MIN, P_MAX = int(os.environ.get("SN_AOTI_P_MIN", "8")), int(os.environ.get("SN_AOTI_P_MAX", "32768"))

_LOADED: set[int] = set()


def trunk_of(model):
    """The `Qwen3Model` whose `.layers` hold the 42 decoder layers."""
    return model.language_model.model


def layers_of(model):
    return trunk_of(model).layers


def artifact_key() -> str | None:
    """`<dtype>/torch<X.Y>/sm<cc>/<shape>` for this process, or None without CUDA."""
    try:
        import torch

        torch_version = ".".join(torch.__version__.split(".")[:2])
        major, minor = torch.cuda.get_device_capability()
    except Exception:
        return None
    return f"bf16/torch{torch_version}/sm{major}{minor}/{AOTI_SHAPE}"


def status() -> str:
    return (
        f"AOTI **on** · `{AOTI_REPO}` ({AOTI_REPO_TYPE}) · shape `{AOTI_SHAPE}`"
        if AOTI
        else "AOTI **off** (`SN_AOTI=1` to load compiled decoder layers)"
    )


# --------------------------------------------------------------------------- weights


def layer_weights(layer) -> dict:
    """Constant map for one layer.

    `state_dict()` is NOT enough: the two rotary `inv_freq` buffers are
    registered non-persistent, so they are absent from `state_dict()` but are
    real graph constants once `package_constants_in_so=False`. Binding them
    partially is a SIGSEGV, not an error, so include every parameter and buffer.
    """
    import torch
    from torch._functorch._aot_autograd.subclass_parametrization import (
        unwrap_tensor_subclass_parameters,
    )

    module = layer
    if any(type(p) is not torch.nn.Parameter for p in layer.parameters()):
        from spaces.zero.torch.aoti import _shallow_clone_module

        module = _shallow_clone_module(layer)
        unwrap_tensor_subclass_parameters(module)
    return {**dict(module.named_parameters()), **dict(module.named_buffers())}


def anonymous_constants(exported_program) -> list[str]:
    return [
        spec.target
        for spec in exported_program.graph_signature.input_specs
        if spec.kind.name == "CONSTANT_TENSOR" and str(spec.target).startswith("_tensor_constant")
    ]


def constant_aliases(exported_program) -> dict:
    """`{'_tensor_constant<N>': '<real fqn>'}` for anonymously lifted constants.

    Empty for this model as long as everything stays a registered parameter or
    buffer; written anyway so the loader can bind if that ever changes.
    """
    targets = [
        spec.target
        for spec in exported_program.graph_signature.input_specs
        if spec.kind.name == "CONSTANT_TENSOR"
    ]
    return {f"_tensor_constant{i}": t for i, t in enumerate(targets)}


# --------------------------------------------------------------------------- load side


def patch_layers(model, package_dir, verbose: bool = True) -> int:
    """Point every layer's `gen_core` at the one compiled package.

    Weights are read on the first call, not here: this can run at startup, and
    ZeroGPU rebinds `param.data` to fresh CUDA tensors when the GPU worker
    materializes the model, which would leave patch-time tensors stale.
    """
    from spaces.zero.torch.aoti import LazyAOTIModel

    package_dir = Path(package_dir)
    pt2 = package_dir / "submodules" / SUBMODULE / "package.pt2"
    if not pt2.is_file():
        raise FileNotFoundError(pt2)

    aliases_path = package_dir / "submodules" / SUBMODULE / ALIASES_FILENAME
    aliases = json.loads(aliases_path.read_text()) if aliases_path.is_file() else {}

    shared = LazyAOTIModel(pt2)
    layers = layers_of(model)

    for layer in layers:
        layer.gen_core = _make_gen_core(shared, pt2, aliases, layer)

    if verbose:
        print(f"[sn-aoti] {len(layers)} decoder layers patched (gen_core)", flush=True)
    return len(layers)


_WARNED: set[str] = set()


def _warn_once(message: str) -> None:
    if message not in _WARNED:
        _WARNED.add(message)
        print(f"[sn-aoti] {message}", flush=True)


def _eager_gen_core(layer, hidden_states, indexes, prefix_k, prefix_v):
    # sn_aoti is a top-level module, not part of the sensenova_u1 package,
    # so this has to be an absolute import.
    from sensenova_u1.models.neo_unify.modeling_qwen3 import gen_block_core

    return gen_block_core(layer, hidden_states, indexes, prefix_k, prefix_v)


def _in_range(hidden_states, prefix_k) -> bool:
    """Is this call inside the shapes the package was exported for?

    Batch is *specialized* to 1 (only dims 1 were made dynamic), and S/P have
    explicit bounds. Everything the demo does sits well inside this, but an
    out-of-range call must degrade to eager rather than rely on the compiled
    guards -- so check rather than find out.
    """
    batch, seq = hidden_states.shape[0], hidden_states.shape[1]
    prefix = prefix_k.shape[1]
    return batch == 1 and S_MIN <= seq <= S_MAX and P_MIN <= prefix <= P_MAX


def _make_gen_core(shared, pt2, aliases, layer):
    """One layer's replacement `gen_core`, binding its weights on first call."""
    bound: dict = {}
    state = {"disabled": False}

    def gen_core(hidden_states, indexes, prefix_k, prefix_v):
        if state["disabled"]:
            return _eager_gen_core(layer, hidden_states, indexes, prefix_k, prefix_v)
        if not _in_range(hidden_states, prefix_k):
            _warn_once(
                f"shape outside the compiled range (B={hidden_states.shape[0]}, "
                f"S={hidden_states.shape[1]}, P={prefix_k.shape[1]}); running eager"
            )
            return _eager_gen_core(layer, hidden_states, indexes, prefix_k, prefix_v)
        # Any failure here -- unbindable constants, a guard the compiled model
        # rejects, anything -- disables the compiled path for this layer and
        # runs eager instead. The invariant that matters is that a partially
        # bound package is never *executed*; declining to use it satisfies that
        # just as well as raising, and keeps the demo serving.
        try:
            first = not bound
            if first:
                bound["weights"] = _resolve_constants(shared, pt2, aliases, layer)
            return shared(bound["weights"], first, hidden_states, indexes, prefix_k, prefix_v)
        except Exception as error:
            _warn_once(f"compiled path unusable ({type(error).__name__}: {error}); "
                       f"this layer runs eager from now on")
            state["disabled"] = True
            return _eager_gen_core(layer, hidden_states, indexes, prefix_k, prefix_v)

    return gen_core


def _resolve_constants(shared, pt2, aliases, layer) -> dict:
    """Bind every compiled constant to one of the layer's own tensors, or refuse.

    Name first, then the compile-side alias sidecar. Partial binding is not an
    option: an unset constant is a SIGSEGV, not an exception.
    """
    weights = layer_weights(layer)
    fqns = _constant_fqns(shared, pt2)
    resolved = {name: weights[name] for name in fqns if name in weights}
    for name in [n for n in fqns if n not in resolved]:
        target = aliases.get(name)
        if target in weights:
            resolved[name] = weights[target]
    missing = [n for n in fqns if n not in resolved]
    if missing:
        raise RuntimeError(
            f"{len(missing)} of {len(fqns)} AOTI constants unbindable from the layer's "
            f"parameters/buffers: {missing[:8]}. Refusing to run: a partially bound "
            f"package dereferences constants nobody set (SIGSEGV rather than an error)."
        )
    return resolved


def _constant_fqns(shared, pt2):
    """`get_constant_fqns()` needs the package loaded; do it once, off the GPU path."""
    import torch

    compiled = shared.compiled_model.get()
    if compiled is None:
        from spaces.zero.torch.aoti import _register_aoti_cleanup

        with _register_aoti_cleanup():
            compiled = torch._inductor.aoti_load_package(pt2)
        shared.compiled_model.set(compiled)
    return compiled.get_constant_fqns()


def maybe_load(model, verbose: bool = True) -> bool:
    """Patch the decoder stack with its compiled package, or leave it eager.

    Safe to call at startup and safe to call when nothing is published: every
    mismatch falls back to eager with one printed line rather than raising.
    """
    if not AOTI or id(model) in _LOADED:
        return False

    key = artifact_key()
    if key is None:
        print("[sn-aoti] no CUDA device visible; running eager", flush=True)
        return False
    if key != AOTI_KEY:
        print(f"[sn-aoti] this card wants `{key}`, only `{AOTI_KEY}` is published; running eager", flush=True)
        return False

    try:
        from huggingface_hub import snapshot_download
        from spaces.zero.torch.aoti import LazyAOTIModel  # noqa: F401
    except Exception as error:
        print(f"[sn-aoti] no AOTI loader here ({type(error).__name__}: {error}); running eager", flush=True)
        return False

    try:
        local = snapshot_download(
            repo_id=AOTI_REPO, repo_type=AOTI_REPO_TYPE, allow_patterns=f"{key}/package/*"
        )
    except Exception as error:
        print(f"[sn-aoti] {AOTI_REPO}:{key} unreachable ({type(error).__name__}: {error}); running eager", flush=True)
        return False

    package_dir = Path(local) / key / "package"
    if not package_dir.is_dir():
        print(f"[sn-aoti] no package at `{AOTI_REPO}:{key}/package`; running eager", flush=True)
        return False

    meta_path = package_dir / METADATA_FILENAME
    if meta_path.is_file():
        meta = json.loads(meta_path.read_text())
        mismatch = _metadata_mismatch(meta, model)
        if mismatch:
            print(f"[sn-aoti] metadata mismatch ({mismatch}); running eager", flush=True)
            return False

    try:
        patch_layers(model, package_dir, verbose=verbose)
    except Exception as error:
        print(f"[sn-aoti] patching failed ({type(error).__name__}: {error}); running eager", flush=True)
        return False

    _LOADED.add(id(model))
    return True


def _metadata_mismatch(meta: dict, model) -> str | None:
    """Return a human-readable reason the published package does not fit, or None."""
    import torch

    checks = []
    if "torch" in meta:
        want = ".".join(str(meta["torch"]).split(".")[:2])
        have = ".".join(torch.__version__.split(".")[:2])
        checks.append(("torch", want, have))
    if "sm" in meta:
        major, minor = torch.cuda.get_device_capability()
        checks.append(("sm", str(meta["sm"]), f"sm{major}{minor}"))
    if "attn_backend" in meta:
        # The attention kernel is baked into the graph. If flash-attn later
        # becomes importable the eager path would switch to it while this
        # package stays on SDPA: different numerics, and possibly slower.
        checks.append(("attn_backend", str(meta["attn_backend"]),
                       __import__("sensenova_u1").effective_attn_backend()))
    if "num_hidden_layers" in meta:
        checks.append(("num_hidden_layers", str(meta["num_hidden_layers"]),
                       str(len(layers_of(model)))))
    if "hidden_size" in meta:
        checks.append(("hidden_size", str(meta["hidden_size"]),
                       str(model.language_model.config.hidden_size)))
    for name, want, have in checks:
        if want != have:
            return f"{name}: package {want} != runtime {have}"
    return None


# --------------------------------------------------------------------------- build side


def capture_layer_call(model, tokenizer, prompt: str, image_size=(2048, 2048), steps: int = 2):
    """Run a couple of real denoising steps and keep one layer-0 `gen_core` call.

    The inputs must come from a real run: `prefix_k`/`prefix_v` are slices of the
    preallocated flash KV cache built by `prepare_flash_kv_cache`, and their
    prefix length depends on the tokenized prompt.
    """
    import torch

    layers = layers_of(model)
    layer0 = layers[0]
    original = layer0.gen_core
    captured = {}

    def recording(hidden_states, indexes, prefix_k, prefix_v):
        if "args" not in captured:
            captured["args"] = (hidden_states, indexes, prefix_k, prefix_v)
        return original(hidden_states, indexes, prefix_k, prefix_v)

    layer0.gen_core = recording
    try:
        # Same knobs the production demo uses, so the captured shapes and the
        # baked-in graph match what will actually run.
        with torch.inference_mode():
            model.t2i_generate(
                tokenizer, prompt, image_size=image_size, cfg_scale=4.0,
                cfg_norm="none", timestep_shift=3.0, cfg_interval=(0.0, 1.0),
                num_steps=steps, batch_size=1, seed=42, think_mode=False,
            )
    finally:
        layer0.gen_core = original

    if "args" not in captured:
        raise RuntimeError("gen_core was never called - the fast denoise path did not engage.")
    args = tuple(t.detach().clone() for t in captured["args"])
    print(f"[sn-aoti] captured shapes: "
          f"{[tuple(t.shape) for t in args]} dtypes={[str(t.dtype) for t in args]}", flush=True)
    return args


def export_layer(model, args, shape: str = AOTI_SHAPE):
    """Export layer 0's gen core, with the image-token and prefix lengths dynamic."""
    import torch

    from sensenova_u1.models.neo_unify.modeling_qwen3 import Qwen3GenBlock

    layer0 = layers_of(model)[0]
    shim = Qwen3GenBlock(layer0)

    # The shim shares the layer's submodules under the same names, so exported
    # constant FQNs are exactly the layer's own parameter/buffer names -- which
    # is what makes one package reusable by all 42 layers.
    sd_layer, sd_shim = set(layer0.state_dict()), set(shim.state_dict())
    assert sd_layer == sd_shim, (sorted(sd_layer - sd_shim)[:5], sorted(sd_shim - sd_layer)[:5])

    if shape == "dynamic":
        S = torch.export.Dim("S", min=S_MIN, max=S_MAX)
        P = torch.export.Dim("P", min=P_MIN, max=P_MAX)
        dynamic_shapes = (
            {1: S},              # hidden_states [B, S, C]
            {1: S},              # indexes       [3, S]
            {1: P},              # prefix_k      [B, P, H_kv, D]
            {1: P},              # prefix_v
        )
    else:
        dynamic_shapes = None

    print(f"[sn-aoti] exporting Qwen3GenBlock, shapes={shape} ...", flush=True)
    with torch.inference_mode():
        ep = torch.export.export(shim, args, {}, dynamic_shapes=dynamic_shapes)

    anon = anonymous_constants(ep)
    if anon:
        print(f"[sn-aoti] WARNING {len(anon)} anonymously lifted constants: {anon[:6]}; "
              f"the alias sidecar will carry their real names", flush=True)
    return ep


def compile_and_save(ep, destination, model, args, extra: dict | None = None) -> Path:
    """Inductor-compile into `<destination>/package/submodules/layers/package.pt2`."""
    import torch

    import spaces

    package_dir = Path(destination) / "package"
    print("[sn-aoti] inductor compile (minutes) ...", flush=True)
    spaces.aoti_compile_and_save(package_dir, ep, submodule=SUBMODULE)

    subdir = package_dir / "submodules" / SUBMODULE
    aliases = constant_aliases(ep)
    if aliases:
        (subdir / ALIASES_FILENAME).write_text(json.dumps(aliases, indent=2))

    major, minor = torch.cuda.get_device_capability()
    llm_cfg = model.language_model.config
    meta = {
        "model_id": "sensenova/SenseNova-U1.5-8B-MoT",
        "module": "language_model.model.layers[*].gen_core "
                  "(Qwen3DecoderLayer image-generation path)",
        "compiled_class": "Qwen3GenBlock",
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
        "sm": f"sm{major}{minor}",
        "device": torch.cuda.get_device_name(0),
        "dtype": "bfloat16",
        "attn_backend": __import__("sensenova_u1").effective_attn_backend(),
        "num_hidden_layers": llm_cfg.num_hidden_layers,
        "hidden_size": llm_cfg.hidden_size,
        "shape": AOTI_SHAPE,
        "signature": ["hidden_states[B,S,C]", "indexes[3,S]",
                      "prefix_k[B,P,H_kv,D]", "prefix_v[B,P,H_kv,D]"],
        "dynamic_dims": {"S": [S_MIN, S_MAX], "P": [P_MIN, P_MAX]},
        "example_shapes": [list(t.shape) for t in args],
        "key": artifact_key(),
        "anonymous_constants": anonymous_constants(ep),
    }
    if extra:
        meta.update(extra)
    (package_dir / METADATA_FILENAME).write_text(json.dumps(meta, indent=2))

    files = sorted(str(p.relative_to(package_dir)) for p in package_dir.rglob("*") if p.is_file())
    total = sum(p.stat().st_size for p in package_dir.rglob("*") if p.is_file())
    print(f"[sn-aoti] package written ({total/1e6:.1f} MB): {files}", flush=True)
    return package_dir


def upload(package_dir, key: str, token: str | None = None) -> str:
    """Push the package under its key. CPU work - never inside GPU time."""
    from huggingface_hub import HfApi

    token = token or os.environ.get("HF_TOKEN")
    if not token:
        raise RuntimeError("HF_TOKEN is needed to push the AOTI package.")
    api = HfApi(token=token)
    api.create_repo(repo_id=AOTI_REPO, repo_type=AOTI_REPO_TYPE, private=False, exist_ok=True)
    api.upload_folder(
        folder_path=str(package_dir), path_in_repo=f"{key}/package",
        repo_id=AOTI_REPO, repo_type=AOTI_REPO_TYPE,
        commit_message=f"AOTI package for {key}",
    )
    return f"https://huggingface.co/{AOTI_REPO}/tree/main/{key}"