"""
gpu_llm.py — the transformers / ZeroGPU dispatcher backend (Space-only).

"Lord Nemo" runs here: NVIDIA Nemotron-3-Nano-4B (arch nemotron_h, a Mamba-hybrid). transformers
>=5.8 ships the NATIVE NemotronH class wired to the `kernels` library, which fetches PREBUILT
mamba-ssm + causal-conv1d kernels from kernels-community/* at runtime (no nvcc, no source compile).
That fast Mamba path replaces the naive torch SSM scan whose huge intermediate tripped a
CUDACachingAllocator NVML assert inside the @spaces.GPU fork — so do NOT use trust_remote_code or a
pip mamba-ssm wheel here; the native class + kernels is the recipe the working hackathon ZeroGPU
Spaces (build-small-hackathon/ready-to-submit, FitCheck) use.

Mirrors the local llama.cpp dispatcher's complete(system, user) -> (text, latency_ms) contract, so
engine / rules / agents never change. Imported LAZILY by backend/llm.py only when
MLP_LLM_BACKEND=transformers, so local dev (stdlib urllib path) never pulls torch / transformers.
"""
from __future__ import annotations
import os
import time

import spaces
import torch
import transformers.generation as _tgen

# ---- kernel-compat shim (must run BEFORE from_pretrained loads the mamba-ssm kernel) ----
# The native NemotronH lazy-loads the prebuilt `kernels-community/mamba-ssm` kernel. That kernel's
# __init__ eagerly imports MambaLMHeadModel, whose bundled (mamba_ssm 2.2.x) code does
# `from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, ...`
# — names transformers 5.x REMOVED (now GenerateDecoderOnlyOutput). We only use the kernel's
# low-level SSM/conv ops, never MambaLMHeadModel.generate, so aliasing the legacy output classes to
# the current Generate* outputs is harmless and makes the kernel importable on any transformers 5.x.
def _alias_legacy_generation_outputs():
    def _has(name):
        try:
            getattr(_tgen, name)
            return True
        except Exception:
            return False
    fallback = None
    for cand in ("GenerateDecoderOnlyOutput", "GenerateEncoderDecoderOutput",
                 "GenerateBeamDecoderOnlyOutput", "GenerateBeamEncoderDecoderOutput"):
        try:
            fallback = getattr(_tgen, cand)
            break
        except Exception:
            continue
    if fallback is None:
        from collections import OrderedDict

        class fallback(OrderedDict):  # last-resort placeholder; never actually instantiated by us
            pass
    for _name in ("GreedySearchDecoderOnlyOutput", "SampleDecoderOnlyOutput",
                  "BeamSearchDecoderOnlyOutput", "BeamSampleDecoderOnlyOutput",
                  "ContrastiveSearchDecoderOnlyOutput", "GreedySearchEncoderDecoderOutput",
                  "SampleEncoderDecoderOutput", "BeamSearchEncoderDecoderOutput",
                  "BeamSampleEncoderDecoderOutput", "ContrastiveSearchEncoderDecoderOutput"):
        if not _has(_name):
            setattr(_tgen, _name, fallback)


_alias_legacy_generation_outputs()

from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402  (after the shim)

MODEL_ID = os.environ.get("MLP_MODEL_ID", "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")
# The dispatcher's JSON is ~120 tokens; 320 leaves margin and BOUNDS worst-case GPU time per turn
# (if the small model repeats/rambles past the JSON, raw_decode still parses the first object).
_MAX_NEW = int(os.environ.get("MLP_MAX_NEW_TOKENS", "320"))

print(f"[gpu_llm] loading dispatcher {MODEL_ID} (native nemotron_h + hub kernels)…", flush=True)
_tok = AutoTokenizer.from_pretrained(MODEL_ID)
if _tok.pad_token_id is None:
    _tok.pad_token = _tok.eos_token
_model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype="auto")
# ZeroGPU-documented pattern: place the model on cuda at MODULE level. `spaces` EMULATES the CUDA
# placement outside the @spaces.GPU context and materializes it on the real GPU inside the call —
# the kernels' fused Mamba path then allocates normally (no NVML fork assert). Only on an actual
# ZeroGPU Space (SPACES_ZERO_GPU set); off-Space we leave it on CPU so import stays harmless.
if os.environ.get("SPACES_ZERO_GPU"):
    _model = _model.to("cuda")
_model.eval()


def _eos_ids():
    # Nemotron's generation_config ends turns with eos_token_id [2, 11]; union it with the
    # tokenizer eos and <|im_end|> so generation stops where the chat template ends.
    gc = _model.generation_config.eos_token_id
    ids = {_tok.eos_token_id}
    ids.update(gc if isinstance(gc, (list, tuple)) else [gc])
    ids.add(_tok.convert_tokens_to_ids("<|im_end|>"))
    return sorted(i for i in ids if isinstance(i, int) and i >= 0)


_EOS = _eos_ids()
print("[gpu_llm] dispatcher ready.", flush=True)


def _encode(system: str, user: str):
    messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
    # Nemotron's chat template defaults enable_thinking=True; turn it OFF so the dispatcher emits
    # clean JSON — a <think> block would eat the token budget and the repair ladder would no-op.
    try:
        return _tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt",
                                        return_dict=True, enable_thinking=False)
    except TypeError:
        return _tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt",
                                        return_dict=True)


# duration=60: the warm dispatch is ~2.4s and the one-time cold materialization (tensor pack +
# kernel JIT) measured ~30s, so 60 is a safe ceiling. A TIGHTER lease than the old 90 is easier for
# the shared ZeroGPU scheduler to place (fewer "pending" stalls) and burns less of the visitor quota.
@spaces.GPU(duration=60)
def _generate(inputs, max_new_tokens: int, temperature: float):
    inputs = {k: v.to(_model.device) for k, v in inputs.items()}
    in_len = inputs["input_ids"].shape[-1]
    with torch.no_grad():
        out = _model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=temperature > 0,
            temperature=max(temperature, 1e-4),
            eos_token_id=_EOS,
            pad_token_id=_tok.pad_token_id,
        )
    return out[0][in_len:].cpu()      # only the freshly generated tokens


def complete(system: str, user: str, max_tokens: int = 512, temperature: float = 0.4):
    """One generation. Returns (raw_text, latency_ms). The JSON repair ladder lives in agents.py."""
    t0 = time.perf_counter()
    inputs = _encode(system, user)
    new_tokens = _generate(inputs, max_new_tokens=max(_MAX_NEW, max_tokens), temperature=temperature)
    text = _tok.decode(new_tokens, skip_special_tokens=True)
    # Defensive: if a reasoning trace slipped through, keep only what's after it.
    if "</think>" in text:
        text = text.split("</think>", 1)[1].lstrip()
    print(f"[gpu_llm] raw[:300]={text[:300]!r}", flush=True)
    return text, (time.perf_counter() - t0) * 1000.0