"""HuggingFace/PyTorch shard backend for real node inference."""

from __future__ import annotations

import base64
from collections import OrderedDict
from dataclasses import dataclass
import json
import os
import threading
import time
import warnings
from pathlib import Path
from typing import Any, Literal, Mapping

Quantization = Literal["auto", "bfloat16", "int8", "nf4"]

# Recipe params this backend knows how to apply (see meshnet_node.recipe_manifest).
# A recipe is only meaningful if its params actually reach the execution path, so
# an unknown key is an error rather than a silent no-op.
SUPPORTED_RECIPE_PARAMS = ("attn_implementation", "use_cache")


class ModelBackendError(RuntimeError):
    """Base class for real model backend startup and execution failures."""


class UnsupportedRecipeParam(ModelBackendError):
    """Raised when a recipe asks for an execution param this backend cannot apply."""


class MissingModelDependencyError(ModelBackendError):
    """Raised when optional model dependencies are not installed."""


class InsufficientVRAMError(ModelBackendError):
    """Raised when a requested shard cannot fit in available CUDA memory."""


class PartialModelLoadUnsupported(ModelBackendError):
    """Raised when a shard cannot be materialized from a local snapshot subset."""


class KVCacheMiss(ModelBackendError):
    """Raised when a decode step references session state this node no longer holds.

    The head recovers by re-prefilling the full sequence (the stateless path),
    so eviction or a node restart degrades throughput instead of corrupting output.
    """


def _torch_cuda_is_executable(torch_module: Any) -> bool:
    """Return True only when this process can actually execute a CUDA/HIP op.

    On ROCm, ``torch.cuda.is_available()`` can be true for an AMD GPU even when
    the installed PyTorch wheel has no runnable kernels for that GPU target.
    Loading weights onto such a device can segfault in native code, so the model
    backend must use the same executable-device check as startup hardware
    detection rather than trusting inventory alone.
    """
    try:
        if not torch_module.cuda.is_available():
            return False
        probe = torch_module.empty((1,), device="cuda")
        probe += 1
        torch_module.cuda.synchronize()
        return True
    except Exception:
        return False


@dataclass(frozen=True)
class TensorPayload:
    """An immutable, request-owned binary activation payload.

    ``body`` is always the exact bfloat16 wire body.  It is intentionally
    owned bytes rather than a view into a request buffer so a payload can move
    across a hop without retaining an HTTP/WebSocket frame after that request
    completes.
    """

    body: bytes
    shape: list[int]
    attention_mask_header: str | None
    position_ids_header: str | None
    # Number of tokens already cached before this payload's tokens (decode steps).
    past_len: int | None = None

    def to_envelope(
        self,
        *,
        name: str,
        request_id: str,
        work_id: str,
        route_session: str,
        route_epoch: int,
        shard_start: int,
        effective_start: int,
        phase: str,
        position: int,
        idempotency_step: int,
        byte_order: str = "little",
        compression: str = "identity",
        max_fragment_bytes: int | None = None,
        token_id_sideband: list[int] | None = None,
        architecture_state: dict[str, Any] | None = None,
        recurrent_state: dict[str, Any] | None = None,
        mtp: dict[str, Any] | None = None,
        extensions: dict[str, Any] | None = None,
    ):
        from .protocol import ActivationEnvelope, DEFAULT_FRAGMENT_BYTES

        return ActivationEnvelope.from_tensor_payload(
            payload=self,
            name=name,
            request_id=request_id,
            work_id=work_id,
            route_session=route_session,
            route_epoch=route_epoch,
            shard_start=shard_start,
            effective_start=effective_start,
            phase=phase,
            position=position,
            idempotency_step=idempotency_step,
            byte_order=byte_order,
            compression=compression,
            max_fragment_bytes=max_fragment_bytes or DEFAULT_FRAGMENT_BYTES,
            token_id_sideband=token_id_sideband,
            architecture_state=architecture_state,
            recurrent_state=recurrent_state,
            mtp=mtp,
            extensions=extensions,
        )

    @classmethod
    def from_envelope(cls, envelope, *, tensor_name: str = "activations"):
        tensor_payload = envelope.to_tensor_payload(tensor_name=tensor_name)
        if not isinstance(tensor_payload, cls):
            raise TypeError("envelope did not produce a TensorPayload")
        return tensor_payload


@dataclass(frozen=True)
class TailTokenResult:
    """Tail-shard decode result: decoded text plus the raw token id.

    The token id lets the head feed the next decode step (and detect EOS)
    without re-tokenizing text, which is not guaranteed to round-trip.
    """

    text: str
    token_id: int


@dataclass
class SessionCacheEntry:
    """Per-session cached state for one shard's layer range.

    `cache` is whatever `use_cache=True` produces for these layers — a
    transformers Cache holding K/V tensors for standard attention, or
    recurrent conv/delta state for hybrid linear-attention layers. The store
    treats it as opaque.
    """

    cache: Any
    seq_len: int
    effective_start: int
    last_used: float


class SessionCacheStore:
    """TTL + LRU bounded map of session_id → SessionCacheEntry.

    Each node caches state only for its own layer range; no node ever holds
    another node's cache. Stale or mismatched entries raise KVCacheMiss so the
    head falls back to a full re-prefill instead of producing corrupt output.
    """

    def __init__(
        self,
        max_sessions: int = 8,
        ttl_seconds: float = 600.0,
        clock: Any = None,
    ) -> None:
        self.max_sessions = max(1, int(max_sessions))
        self.ttl_seconds = float(ttl_seconds)
        self._clock = clock or time.monotonic
        self._entries: OrderedDict[str, SessionCacheEntry] = OrderedDict()
        self._lock = threading.Lock()

    def __len__(self) -> int:
        with self._lock:
            return len(self._entries)

    def store(self, session_id: str, cache: Any, seq_len: int, effective_start: int) -> SessionCacheEntry:
        now = self._clock()
        with self._lock:
            self._entries.pop(session_id, None)
            entry = SessionCacheEntry(cache, seq_len, effective_start, now)
            self._entries[session_id] = entry
            self._evict_locked(now)
            return entry

    def lookup(
        self,
        session_id: str,
        *,
        expected_seq_len: int | None = None,
        effective_start: int | None = None,
    ) -> SessionCacheEntry:
        now = self._clock()
        with self._lock:
            self._evict_locked(now)
            entry = self._entries.get(session_id)
            if entry is None:
                raise KVCacheMiss(f"no cached state for session {session_id[:8]}")
            if expected_seq_len is not None and entry.seq_len != expected_seq_len:
                del self._entries[session_id]
                raise KVCacheMiss(
                    f"session {session_id[:8]} cache holds {entry.seq_len} tokens, "
                    f"expected {expected_seq_len}"
                )
            if effective_start is not None and entry.effective_start != effective_start:
                del self._entries[session_id]
                raise KVCacheMiss(
                    f"session {session_id[:8]} cached with start_layer "
                    f"{entry.effective_start}, requested {effective_start}"
                )
            entry.last_used = now
            self._entries.move_to_end(session_id)
            return entry

    def drop(self, session_id: str) -> None:
        with self._lock:
            self._entries.pop(session_id, None)

    def _evict_locked(self, now: float) -> None:
        if self.ttl_seconds > 0:
            expired = [
                sid for sid, entry in self._entries.items()
                if now - entry.last_used > self.ttl_seconds
            ]
            for sid in expired:
                del self._entries[sid]
        while len(self._entries) > self.max_sessions:
            self._entries.popitem(last=False)


def validate_quantization(value: str) -> Quantization:
    if value not in {"auto", "bfloat16", "int8", "nf4"}:
        raise ValueError("quantization must be one of: auto, bfloat16, int8, nf4")
    return value  # type: ignore[return-value]


def build_quantization_config(quantization: Quantization) -> Any | None:
    """Return a transformers BitsAndBytesConfig for quantized weights."""
    if quantization in {"auto", "bfloat16"}:
        return None
    try:
        import torch
        from transformers import BitsAndBytesConfig
    except ModuleNotFoundError as exc:
        raise MissingModelDependencyError(
            "transformers and torch are required for int8/nf4 quantization"
        ) from exc

    if quantization == "int8":
        return BitsAndBytesConfig(load_in_8bit=True)
    return BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )


class TorchModelShard:
    """Executable subset of a HuggingFace causal language model."""

    def __init__(
        self,
        model_id: str,
        shard_start: int,
        shard_end: int,
        quantization: Quantization = "auto",
        cache_dir: Path | None = None,
        force_cpu: bool = False,
        recipe_params: Mapping[str, Any] | None = None,
    ) -> None:
        if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
            raise ValueError("shard_start must be <= shard_end and non-negative")
        self.model_id = model_id
        self.shard_start = shard_start
        self.shard_end = shard_end
        self.quantization = quantization
        self.recipe_params = validate_recipe_params(recipe_params)
        attn_implementation = self.recipe_params.get("attn_implementation")

        try:
            import torch
            from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
        except ModuleNotFoundError as exc:
            raise MissingModelDependencyError(
                "real model backend requires torch, transformers, safetensors, accelerate, and bitsandbytes"
            ) from exc

        self.torch = torch
        if force_cpu:
            self.device = torch.device("cpu")
        else:
            self.device = torch.device("cuda" if _torch_cuda_is_executable(torch) else "cpu")
        load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id
        quant_config, dtype, uses_quantized_weights = _model_load_plan(
            AutoConfig,
            load_source,
            quantization,
            torch,
            None if load_source != model_id else cache_dir,
        )
        try:
            total_layers_hint = _total_layers_for_local_snapshot(AutoConfig, load_source)
            if _should_partial_materialize_shard(
                load_source,
                shard_start,
                shard_end,
                total_layers_hint=total_layers_hint,
                uses_quantized_weights=uses_quantized_weights,
            ):
                self.model = _load_partial_model_from_snapshot(
                    AutoConfig,
                    AutoModelForCausalLM,
                    torch,
                    load_source,
                    shard_start,
                    shard_end,
                    dtype,
                    self.device,
                    attn_implementation=attn_implementation,
                )
            else:
                load_kwargs = {
                    "device_map": "auto" if uses_quantized_weights else None,
                    "dtype": dtype,
                    "low_cpu_mem_usage": True,
                    "cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None,
                }
                if quant_config is not None:
                    load_kwargs["quantization_config"] = quant_config
                if attn_implementation is not None:
                    load_kwargs["attn_implementation"] = attn_implementation
                self.model = AutoModelForCausalLM.from_pretrained(
                    load_source,
                    **load_kwargs,
                )
                if not uses_quantized_weights:
                    self.model.to(self.device)
        except Exception as exc:
            if _looks_like_oom(exc):
                memory_kind = "VRAM" if self.device.type == "cuda" else "RAM"
                raise InsufficientVRAMError(
                    f"insufficient {memory_kind} to load {model_id} layers {shard_start}:{shard_end} "
                    f"with {quantization} quantization; choose a smaller shard or lower quantization"
                ) from exc
            raise

        self.model.eval()
        self.tokenizer = AutoTokenizer.from_pretrained(
            load_source,
            cache_dir=str(cache_dir) if cache_dir is not None and load_source == model_id else None,
        )
        self.layers = _model_layers(self.model)
        self.total_layers = len(self.layers)
        # shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention.
        if shard_end >= self.total_layers:
            raise ValueError(
                f"shard_end {shard_end} exceeds last layer index {self.total_layers - 1}"
            )
        self.is_head = shard_start == 0
        self.is_tail = shard_end >= self.total_layers - 1
        self.hidden_size = int(
            getattr(self.model.config, "hidden_size", 0)
            or getattr(self.model.config, "n_embd", 0)
        )
        self._embed_tokens = _embed_tokens(self.model) if self.is_head else None
        self._position_embeddings = _position_embeddings(self.model)
        self._norm = _final_norm(self.model) if self.is_tail else None
        self._lm_head = getattr(self.model, "lm_head", None) if self.is_tail else None
        # Per-session KV/recurrent-state cache for this shard's layer range.
        # Hybrid/linear-attention models such as Qwen3.6 can dispatch Triton
        # recurrent-cache kernels when use_cache=True. Those kernels cannot
        # consume CPU tensors ("Pointer argument cannot be accessed from Triton"),
        # so CPU shards intentionally stay on the stateless prefill path.
        self.supports_kv_cache = self.device.type != "cpu"
        if self.recipe_params.get("use_cache") is False:
            self.supports_kv_cache = False
        self.kv_sessions = SessionCacheStore(
            max_sessions=int(os.environ.get("MESHNET_KV_MAX_SESSIONS", "8")),
            ttl_seconds=float(os.environ.get("MESHNET_KV_TTL_SECONDS", "600")),
        )

    def encode_prompt(self, prompt: str, session_id: str | None = None) -> TensorPayload:
        if not self.is_head or self._embed_tokens is None:
            raise ModelBackendError("text prompts can only be accepted by the head shard")
        encoded = self.tokenizer(prompt, return_tensors="pt")
        input_ids = encoded["input_ids"].to(self.device)
        attention_mask = encoded.get("attention_mask")
        if attention_mask is None:
            attention_mask = self.torch.ones_like(input_ids)
        attention_mask = attention_mask.to(self.device)
        position_ids = _position_ids(attention_mask, self.torch)
        hidden_states = self._embed_tokens(input_ids)
        if self._position_embeddings is not None:
            hidden_states = hidden_states + self._position_embeddings(position_ids)
        hidden_states = self._run_layers_session(
            hidden_states, attention_mask, position_ids,
            session_id=session_id, cache_mode="prefill" if session_id else None,
        )
        return self._payload(hidden_states, attention_mask, position_ids)

    def encode_next_token(self, token_id: int, session_id: str) -> TensorPayload:
        """Decode step: embed one new token against this head's cached session.

        Raises KVCacheMiss if the session was evicted — callers fall back to a
        full re-prefill via encode_prompt.
        """
        if not self.is_head or self._embed_tokens is None:
            raise ModelBackendError("decode steps can only start at the head shard")
        if not self.supports_kv_cache:
            raise KVCacheMiss("kv cache disabled on this backend")
        entry = self.kv_sessions.lookup(
            session_id, effective_start=self._effective_start(None)
        )
        past_len = entry.seq_len
        input_ids = self.torch.tensor([[int(token_id)]], dtype=self.torch.long, device=self.device)
        position_ids = self.torch.tensor([[past_len]], dtype=self.torch.long, device=self.device)
        hidden_states = self._embed_tokens(input_ids)
        if self._position_embeddings is not None:
            hidden_states = hidden_states + self._position_embeddings(position_ids)
        hidden_states = self._run_layers(
            hidden_states, None, position_ids,
            cache=entry.cache, past_len=past_len,
        )
        entry.seq_len = past_len + 1
        return TensorPayload(
            body=_tensor_to_bytes(hidden_states.to(self.torch.bfloat16).contiguous()),
            shape=list(hidden_states.shape),
            attention_mask_header=None,
            position_ids_header=_int_tensor_header(position_ids),
            past_len=past_len,
        )

    def forward_bytes(
        self,
        body: bytes,
        shape: list[int],
        attention_mask_header: str | None,
        position_ids_header: str | None,
        start_layer: int | None = None,
        session_id: str | None = None,
        cache_mode: str | None = None,
        past_len: int | None = None,
    ) -> TensorPayload | TailTokenResult | str:
        hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
            self.device
        )
        attention_mask = _tensor_from_int64_header(
            attention_mask_header, self.torch, self.device
        )
        position_ids = _tensor_from_int64_header(
            position_ids_header, self.torch, self.device
        )
        hidden_states = self._run_layers_session(
            hidden_states, attention_mask, position_ids, start_layer=start_layer,
            session_id=session_id, cache_mode=cache_mode, past_len=past_len,
        )
        if self.is_tail:
            return self.decode_tail_token(hidden_states)
        return self._payload(hidden_states, attention_mask, position_ids)

    def decode_tail(self, hidden_states: Any) -> str:
        return self.decode_tail_token(hidden_states).text

    def decode_tail_token(self, hidden_states: Any) -> TailTokenResult:
        if self._norm is not None:
            hidden_states = self._norm(hidden_states)
        if self._lm_head is None:
            raise ModelBackendError("tail shard has no lm_head")
        logits = self._lm_head(hidden_states)
        token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
        return TailTokenResult(
            text=self.tokenizer.decode([token_id], skip_special_tokens=True),
            token_id=token_id,
        )

    def eos_token_ids(self) -> list[int]:
        """All token ids that should terminate generation (tokenizer + generation config)."""
        ids: set[int] = set()
        tok_eos = getattr(self.tokenizer, "eos_token_id", None)
        gen_config = getattr(self.model, "generation_config", None)
        gen_eos = getattr(gen_config, "eos_token_id", None) if gen_config is not None else None
        for value in (tok_eos, gen_eos):
            if value is None:
                continue
            if isinstance(value, (list, tuple)):
                ids.update(int(v) for v in value)
            else:
                ids.add(int(value))
        return sorted(ids)

    def release_session(self, session_id: str) -> None:
        self.kv_sessions.drop(session_id)

    def generate_text(
        self,
        messages: list[dict],
        max_new_tokens: int = 5120,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ) -> str:
        """Autoregressive generation using HF generate() — single-node (head+tail) mode."""
        if not self.is_head or not self.is_tail:
            raise ModelBackendError("local generation requires a full head+tail shard")
        encoded = self._encode_messages(messages)
        input_ids = encoded["input_ids"].to(self.device)
        attention_mask = encoded.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)
        pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None)
        do_sample = temperature != 1.0 or top_p != 1.0
        with self.torch.inference_mode():
            generated = self.model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=max(1, int(max_new_tokens)),
                do_sample=do_sample,
                temperature=temperature if do_sample else None,
                top_p=top_p if do_sample else None,
                pad_token_id=pad_token_id,
            )
        new_tokens = generated[0, input_ids.shape[-1]:]
        return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

    def generate_text_streaming(
        self,
        messages: list[dict],
        max_new_tokens: int = 5000,
        temperature: float = 1.0,
        top_p: float = 1.0,
    ):
        """Yield decoded token strings one at a time using HF TextIteratorStreamer."""
        if not self.is_head or not self.is_tail:
            raise ModelBackendError("streaming generation requires a full head+tail shard")
        import threading
        try:
            from transformers import TextIteratorStreamer  # type: ignore[import]
        except ImportError:
            yield self.generate_text(messages, max_new_tokens, temperature, top_p)
            return

        encoded = self._encode_messages(messages)
        input_ids = encoded["input_ids"].to(self.device)
        attention_mask = encoded.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)
        pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None)
        do_sample = temperature != 1.0 or top_p != 1.0

        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
        gen_kwargs = dict(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max(1, int(max_new_tokens)),
            do_sample=do_sample,
            temperature=temperature if do_sample else None,
            top_p=top_p if do_sample else None,
            pad_token_id=pad_token_id,
            streamer=streamer,
        )
        t = threading.Thread(target=self.model.generate, kwargs=gen_kwargs, daemon=True)
        t.start()
        for token_text in streamer:
            yield token_text
        t.join()

    def count_prompt_tokens(self, messages: list[dict]) -> int:
        """Return tokenizer-backed prompt token count for OpenAI usage metadata."""
        encoded = self._encode_messages(messages)
        input_ids = encoded["input_ids"]
        return int(input_ids.shape[-1])

    def count_text_tokens(self, text: str) -> int:
        """Return tokenizer-backed completion token count for OpenAI usage metadata."""
        try:
            encoded = self.tokenizer(
                text,
                return_tensors="pt",
                add_special_tokens=False,
            )
        except TypeError:
            encoded = self.tokenizer(text, return_tensors="pt")
        return int(encoded["input_ids"].shape[-1])

    def _encode_messages(self, messages: list[dict]) -> dict:
        """Format messages with chat template (if available) and tokenize."""
        if hasattr(self.tokenizer, "apply_chat_template"):
            try:
                prompt_str = self.tokenizer.apply_chat_template(
                    messages,
                    add_generation_prompt=True,
                    tokenize=False,
                )
                return dict(self.tokenizer(prompt_str, return_tensors="pt"))
            except Exception:
                pass
        prompt = " ".join(
            str(m.get("content", ""))
            for m in messages
            if isinstance(m, dict) and m.get("role") == "user"
        )
        return dict(self.tokenizer(prompt, return_tensors="pt"))

    def _effective_start(self, start_layer: int | None) -> int:
        # start_layer overrides shard_start for overlapping-shard routing
        # (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
        # indexing outside the loaded weights.
        return (
            max(self.shard_start, start_layer)
            if start_layer is not None
            else self.shard_start
        )

    def _new_session_cache(self) -> Any | None:
        """Build the model-appropriate cache object for one session.

        DynamicCache(config=...) lets transformers pick the right per-layer
        state (K/V for standard attention, conv/recurrent state for hybrid
        linear-attention layers) — the same construction the model's own
        forward() uses when use_cache=True.
        """
        try:
            from transformers import DynamicCache
        except ImportError:
            return None
        try:
            return DynamicCache(config=self.model.config)
        except TypeError:
            return DynamicCache()

    def _run_layers_session(
        self,
        hidden_states: Any,
        attention_mask: Any,
        position_ids: Any,
        start_layer: int | None = None,
        session_id: str | None = None,
        cache_mode: str | None = None,
        past_len: int | None = None,
    ) -> Any:
        """Run this shard's layers, keying cached state by session when requested.

        cache_mode "prefill" creates fresh session state; "decode" requires an
        existing entry (KVCacheMiss otherwise). None runs fully stateless —
        today's behavior, kept as the recovery path.
        """
        effective_start = self._effective_start(start_layer)
        if not (session_id and cache_mode and self.supports_kv_cache):
            if cache_mode == "decode":
                # A decode payload is one token — running it stateless would
                # silently produce garbage. Force the head to re-prefill.
                raise KVCacheMiss("kv cache disabled on this backend")
            return self._run_layers(
                hidden_states, attention_mask, position_ids, start_layer=start_layer
            )
        if cache_mode == "decode":
            entry = self.kv_sessions.lookup(
                session_id,
                expected_seq_len=past_len,
                effective_start=effective_start,
            )
            seq_len = int(hidden_states.shape[1])
            # Decode attends over cache + new token; no padding, so no mask needed.
            hidden_states = self._run_layers(
                hidden_states, None, position_ids,
                start_layer=start_layer, cache=entry.cache, past_len=entry.seq_len,
            )
            entry.seq_len += seq_len
            return hidden_states
        # Prefill: fresh cache for this session (replaces any stale entry).
        cache = self._new_session_cache()
        if cache is None:
            return self._run_layers(
                hidden_states, attention_mask, position_ids, start_layer=start_layer
            )
        try:
            result = self._run_layers(
                hidden_states, attention_mask, position_ids,
                start_layer=start_layer, cache=cache, past_len=0,
            )
        except Exception as exc:
            if not _cache_unsupported_for_shard(exc):
                raise
            # Layers reject cache kwargs (exotic architecture) — disable caching
            # for this backend and stay on the stateless path. Some hybrid
            # CPU paths also accept cache kwargs but fail at runtime inside
            # Triton-only kernels; treat those as cache-unsupported too.
            self.supports_kv_cache = False
            print(f"  [node] kv cache unsupported by {self.model_id}: {exc}", flush=True)
            return self._run_layers(
                hidden_states, attention_mask, position_ids, start_layer=start_layer
            )
        self.kv_sessions.store(
            session_id, cache,
            seq_len=int(hidden_states.shape[1]),
            effective_start=effective_start,
        )
        return result

    def _run_layers(
        self,
        hidden_states: Any,
        attention_mask: Any,
        position_ids: Any,
        start_layer: int | None = None,
        cache: Any = None,
        past_len: int = 0,
    ) -> Any:
        effective_start = self._effective_start(start_layer)
        position_embeddings = _rotary_position_embeddings(
            self.model,
            hidden_states,
            position_ids,
        )
        layer_attention_mask = _decoder_attention_mask(
            attention_mask,
            hidden_states,
            self.torch,
        )
        cache_position = None
        if cache is not None:
            seq_len = int(hidden_states.shape[1])
            cache_position = self.torch.arange(
                past_len, past_len + seq_len, device=hidden_states.device
            )
        with self.torch.inference_mode():
            for layer in self.layers[effective_start:self.shard_end + 1]:
                hidden_states = _call_layer(
                    layer,
                    hidden_states,
                    layer_attention_mask,
                    position_ids,
                    position_embeddings,
                    cache=cache,
                    cache_position=cache_position,
                )
            return hidden_states.to(self.torch.bfloat16)

    def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload:
        hidden_states = hidden_states.to(self.torch.bfloat16).contiguous()
        return TensorPayload(
            body=_tensor_to_bytes(hidden_states),
            shape=list(hidden_states.shape),
            attention_mask_header=_int_tensor_header(attention_mask)
            if attention_mask is not None
            else None,
            position_ids_header=_int_tensor_header(position_ids)
            if position_ids is not None
            else None,
        )


def validate_recipe_params(params: Mapping[str, Any] | None) -> dict[str, Any]:
    """Return recipe params this backend can honour, or raise naming the bad key."""
    if not params:
        return {}
    unsupported = [key for key in params if key not in SUPPORTED_RECIPE_PARAMS]
    if unsupported:
        raise UnsupportedRecipeParam(
            f"recipe param(s) {', '.join(sorted(unsupported))} are not supported by this "
            f"backend; it applies: {', '.join(SUPPORTED_RECIPE_PARAMS)}"
        )
    return dict(params)


def load_torch_shard(
    model_id: str,
    shard_start: int,
    shard_end: int,
    quantization: Quantization = "auto",
    cache_dir: Path | None = None,
    force_cpu: bool = False,
    recipe_params: Mapping[str, Any] | None = None,
) -> TorchModelShard:
    return TorchModelShard(
        model_id,
        shard_start,
        shard_end,
        quantization,
        cache_dir,
        force_cpu=force_cpu,
        recipe_params=recipe_params,
    )


def _total_layers_for_local_snapshot(auto_config: Any, load_source: str) -> int | None:
    snapshot_dir = Path(load_source)
    if not (snapshot_dir / "config.json").exists():
        return None
    from .model_catalog import layers_from_config

    try:
        cfg = auto_config.from_pretrained(str(snapshot_dir))
    except Exception:
        return None
    return layers_from_config(cfg)


def _should_partial_materialize_shard(
    load_source: str,
    shard_start: int,
    shard_end: int,
    *,
    total_layers_hint: int | None,
    uses_quantized_weights: bool,
) -> bool:
    if uses_quantized_weights:
        return False
    snapshot_dir = Path(load_source)
    if not snapshot_dir.exists() or not (snapshot_dir / "config.json").exists():
        return False
    if not (snapshot_dir / "model.safetensors.index.json").exists():
        return False
    if total_layers_hint is None:
        return False
    return True


def _load_partial_model_from_snapshot(
    auto_config: Any,
    auto_model_for_causal_lm: Any,
    torch: Any,
    load_source: str,
    shard_start: int,
    shard_end: int,
    dtype: Any,
    device: Any,
    *,
    init_empty_weights_fn: Any | None = None,
    set_tensor_fn: Any | None = None,
    safe_open_fn: Any | None = None,
    attn_implementation: str | None = None,
) -> Any:
    from .model_catalog import layers_from_config
    from .safetensors_selection import (
        INDEX_FILENAME,
        select_tensor_names_for_layers_from_index,
    )

    if init_empty_weights_fn is None:
        from accelerate import init_empty_weights as init_empty_weights_fn
    if set_tensor_fn is None:
        from accelerate.utils import set_module_tensor_to_device as set_tensor_fn
    if safe_open_fn is None:
        from safetensors import safe_open as safe_open_fn

    snapshot_dir = Path(load_source)
    cfg = auto_config.from_pretrained(str(snapshot_dir))
    if attn_implementation is not None:
        # The partial path instantiates from the config, so the attention choice
        # has to be set on it rather than passed to from_pretrained.
        cfg._attn_implementation = attn_implementation
    total_layers = layers_from_config(cfg)
    if total_layers is None:
        raise PartialModelLoadUnsupported(
            f"could not determine num_hidden_layers for local snapshot {snapshot_dir}"
        )
    if shard_end >= total_layers:
        raise ValueError(
            f"shard_end {shard_end} exceeds last layer index {total_layers - 1}"
        )

    index_path = snapshot_dir / INDEX_FILENAME
    try:
        index = json.loads(index_path.read_text(encoding="utf-8"))
    except FileNotFoundError as exc:
        raise PartialModelLoadUnsupported(
            f"missing SafeTensors index for partial load: {index_path}"
        ) from exc
    weight_map = index.get("weight_map")
    if not isinstance(weight_map, dict):
        raise PartialModelLoadUnsupported(f"{INDEX_FILENAME} must contain a weight_map object")

    tensor_names = select_tensor_names_for_layers_from_index(
        weight_map,
        shard_start,
        shard_end,
        total_layers=total_layers,
    )
    if not tensor_names:
        raise PartialModelLoadUnsupported(
            f"no checkpoint tensors matched layers {shard_start}-{shard_end} in {snapshot_dir}"
        )

    with init_empty_weights_fn():
        model = auto_model_for_causal_lm.from_config(_causal_lm_config(cfg), torch_dtype=dtype)
    tie_weights = getattr(model, "tie_weights", None)
    if callable(tie_weights):
        tie_weights()

    # Multimodal/MTP checkpoints (e.g. Qwen3.5/3.6-MoE) carry vision and
    # multi-token-prediction tensors the text-only CausalLM never builds;
    # transformers' from_pretrained drops them via _keys_to_ignore_on_load_unexpected,
    # so the manual loader must skip them too.
    expected_keys = _model_state_dict_keys(model)
    tensors_by_file: dict[str, list[str]] = {}
    skipped: list[str] = []
    for tensor_name in sorted(tensor_names):
        rel_file = weight_map.get(tensor_name)
        if not isinstance(rel_file, str):
            continue
        if (
            expected_keys is not None
            and _checkpoint_tensor_name_for_model(model, tensor_name) not in expected_keys
        ):
            skipped.append(tensor_name)
            continue
        tensors_by_file.setdefault(rel_file, []).append(tensor_name)
    if skipped:
        preview = ", ".join(skipped[:3])
        print(
            f"  Skipping {len(skipped)} checkpoint tensors absent from the causal LM "
            f"(e.g. {preview})",
            flush=True,
        )
    if not tensors_by_file:
        raise PartialModelLoadUnsupported(
            f"no checkpoint tensors for layers {shard_start}-{shard_end} match the "
            f"causal LM built from {snapshot_dir}"
        )

    for rel_file, names in tensors_by_file.items():
        checkpoint_file = snapshot_dir / rel_file
        if not checkpoint_file.exists():
            raise PartialModelLoadUnsupported(
                f"checkpoint file advertised in {INDEX_FILENAME} is missing: {checkpoint_file}"
            )
        with safe_open_fn(str(checkpoint_file), framework="pt", device="cpu") as handle:
            for tensor_name in names:
                set_tensor_fn(
                    model,
                    _checkpoint_tensor_name_for_model(model, tensor_name),
                    device,
                    value=handle.get_tensor(tensor_name),
                    dtype=dtype,
                )

    _finalize_active_shard_modules_on_device(model, shard_start, shard_end, device)
    return model


def _finalize_active_shard_modules_on_device(
    model: Any, shard_start: int, shard_end: int, device: Any
) -> None:
    """Place active shard modules on device without copying unmaterialized meta weights."""
    for module in _active_modules_for_shard(model, shard_start, shard_end):
        parameters = getattr(module, "parameters", None)
        if not callable(parameters):
            if hasattr(module, "to"):
                module.to(device)
            continue
        params = list(parameters(recurse=True))
        buffers_fn = getattr(module, "buffers", None)
        buffers = list(buffers_fn(recurse=True)) if callable(buffers_fn) else []
        tensors = params + buffers
        if not tensors:
            if hasattr(module, "to"):
                module.to(device)
            continue
        if all(tensor.device.type == "meta" for tensor in tensors):
            to_empty = getattr(module, "to_empty", None)
            if callable(to_empty):
                to_empty(device)
            continue
        if all(tensor.device.type != "meta" for tensor in tensors):
            if hasattr(module, "to"):
                module.to(device)
            continue
        # Partially materialized: set_module_tensor_to_device already placed loaded
        # weights on the target device; leave remaining meta parameters untouched.


def _model_load_plan(
    auto_config: Any,
    model_id: str,
    quantization: Quantization,
    torch: Any,
    cache_dir: Path | None = None,
) -> tuple[Any | None, Any, bool]:
    """Return (explicit quant config, dtype, uses quantized weights)."""
    if quantization != "auto":
        quant_config = build_quantization_config(quantization)
        return quant_config, torch.bfloat16, quant_config is not None

    cfg = auto_config.from_pretrained(
        model_id,
        cache_dir=str(cache_dir) if cache_dir is not None else None,
    )
    if _native_quantization_config(cfg) is not None:
        return None, _native_torch_dtype(cfg, torch), True
    return None, _native_torch_dtype(cfg, torch), False


def _config_candidates(cfg: Any) -> list[Any]:
    candidates = [cfg]
    get_text_config = getattr(cfg, "get_text_config", None)
    if callable(get_text_config):
        try:
            candidates.append(get_text_config())
        except Exception:
            pass
    text_config = getattr(cfg, "text_config", None)
    if text_config is not None:
        candidates.append(text_config)
    return candidates


def _native_quantization_config(cfg: Any) -> Any | None:
    for candidate in _config_candidates(cfg):
        quant_config = getattr(candidate, "quantization_config", None)
        if quant_config:
            return quant_config
    return None


def _native_torch_dtype(cfg: Any, torch: Any) -> Any:
    for candidate in _config_candidates(cfg):
        for attr in ("dtype", "torch_dtype"):
            dtype = getattr(candidate, attr, None)
            if dtype is None:
                continue
            if isinstance(dtype, str):
                dtype_name = dtype.removeprefix("torch.")
                dtype_value = getattr(torch, dtype_name, None)
                if dtype_value is not None:
                    return dtype_value
            else:
                return dtype
    return torch.bfloat16


def _causal_lm_config(cfg: Any) -> Any:
    """Use the text decoder config for composite VLM/MoE presets."""
    get_text_config = getattr(cfg, "get_text_config", None)
    if callable(get_text_config):
        try:
            return get_text_config()
        except Exception:
            pass
    text_config = getattr(cfg, "text_config", None)
    if text_config is not None:
        return text_config
    return cfg


def _model_state_dict_keys(model: Any) -> set[str] | None:
    """Expected parameter/buffer names, or None when the model can't report them."""
    state_dict = getattr(model, "state_dict", None)
    if not callable(state_dict):
        return None
    try:
        return set(state_dict().keys())
    except Exception:
        return None


def _checkpoint_tensor_name_for_model(model: Any, tensor_name: str) -> str:
    """Map multimodal checkpoint keys onto text-only CausalLM modules when needed."""
    inner = getattr(model, "model", None)
    if inner is not None and hasattr(inner, "language_model"):
        return tensor_name
    if ".language_model." in tensor_name:
        return tensor_name.replace(".language_model.", ".")
    return tensor_name


def _transformer_backbone(model: Any) -> Any:
    if hasattr(model, "model"):
        inner = model.model
        language_model = getattr(inner, "language_model", None)
        if language_model is not None:
            return language_model
        return inner
    if hasattr(model, "transformer"):
        return model.transformer
    raise ModelBackendError(
        "unsupported HuggingFace model architecture: no transformer backbone found"
    )


def _model_layers(model: Any) -> Any:
    backbone = _transformer_backbone(model)
    for attr in ("layers", "h", "blocks"):
        layers = getattr(backbone, attr, None)
        if layers is not None:
            return layers
    raise ModelBackendError(
        "unsupported HuggingFace model architecture: no transformer layers found"
    )


def _embed_tokens(model: Any) -> Any:
    backbone = _transformer_backbone(model)
    for attr in ("embed_tokens", "wte"):
        embed = getattr(backbone, attr, None)
        if embed is not None:
            return embed
    raise ModelBackendError(
        "unsupported HuggingFace model architecture: no token embeddings found"
    )


def _position_embeddings(model: Any) -> Any | None:
    backbone = _transformer_backbone(model)
    return getattr(backbone, "wpe", None)


def _rotary_embedding_module(model: Any) -> Any | None:
    backbone = _transformer_backbone(model)
    return getattr(backbone, "rotary_emb", None)


def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]:
    active: list[Any] = []

    def add(module: Any | None) -> None:
        if module is None:
            return
        if any(existing is module for existing in active):
            return
        active.append(module)

    if shard_start == 0:
        add(_embed_tokens(model))
        add(_position_embeddings(model))
    add(_rotary_embedding_module(model))
    for layer in _model_layers(model)[shard_start:shard_end + 1]:
        add(layer)
    total_layers = len(_model_layers(model))
    if shard_end >= total_layers - 1:
        add(_final_norm(model))
        add(getattr(model, "lm_head", None))
    return active


def _final_norm(model: Any) -> Any | None:
    backbone = _transformer_backbone(model)
    for attr in ("norm", "ln_f", "final_layer_norm"):
        norm = getattr(backbone, attr, None)
        if norm is not None:
            return norm
    return None


def _position_ids(attention_mask: Any, torch: Any) -> Any:
    position_ids = attention_mask.long().cumsum(-1) - 1
    return position_ids.masked_fill(attention_mask == 0, 0).to(torch.long)


def _decoder_attention_mask(attention_mask: Any, hidden_states: Any, torch: Any) -> Any:
    """Build a causal additive mask for decoder layers called outside model.forward."""
    if attention_mask is None:
        return None
    if len(getattr(attention_mask, "shape", ())) != 2:
        return attention_mask
    batch_size, seq_len = attention_mask.shape
    if seq_len <= 1:
        return None if bool(attention_mask.all()) else attention_mask.to(hidden_states.dtype)

    min_value = torch.finfo(hidden_states.dtype).min
    causal = torch.full(
        (seq_len, seq_len),
        min_value,
        dtype=hidden_states.dtype,
        device=hidden_states.device,
    )
    causal = torch.triu(causal, diagonal=1)
    causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone()

    padding = attention_mask.to(device=hidden_states.device)
    if not bool(padding.all()):
        causal = causal.masked_fill(padding[:, None, None, :] == 0, min_value)
    return causal


def _rotary_position_embeddings(model: Any, hidden_states: Any, position_ids: Any) -> Any | None:
    """Return model-level rotary embeddings required by newer HF decoder layers."""
    if position_ids is None:
        return None
    rotary = _rotary_embedding_module(model)
    if rotary is None:
        return None
    return rotary(hidden_states, position_ids)


def _call_layer(
    layer: Any,
    hidden_states: Any,
    attention_mask: Any,
    position_ids: Any,
    position_embeddings: Any | None = None,
    cache: Any = None,
    cache_position: Any = None,
) -> Any:
    attempts = (
        {
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "position_embeddings": position_embeddings,
            "use_cache": False,
        },
        {
            "attention_mask": attention_mask,
            "position_ids": position_ids,
            "use_cache": False,
        },
        {"attention_mask": attention_mask, "use_cache": False},
        {"use_cache": False},
        {},
    )
    last_exc: Exception | None = None
    for kwargs in attempts:
        filtered = {key: value for key, value in kwargs.items() if value is not None}
        if cache is not None:
            # transformers 5.x layers take a Cache via past_key_values and
            # mutate it in place; cache_position is required by sliding-window
            # and hybrid recurrent layers.
            filtered["past_key_values"] = cache
            filtered["use_cache"] = True
            if cache_position is not None:
                filtered["cache_position"] = cache_position
        try:
            output = layer(hidden_states, **filtered)
            return output[0] if isinstance(output, tuple) else output
        except TypeError as exc:
            last_exc = exc
    if last_exc is not None:
        raise last_exc
    return layer(hidden_states)[0]


def _tensor_to_bytes(tensor: Any) -> bytes:
    import torch

    return tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()


def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> Any:
    # ``frombuffer`` views the immutable request-owned bytes for this forward
    # only.  The following device transfer is the one required CPU→GPU copy;
    # wrapping in ``bytearray`` first used to add an avoidable CPU allocation
    # and copy.  Do not upcast through float32: the activation wire contract
    # is bfloat16 and model layers accept it directly.
    # PyTorch warns because bytes are immutable even though the forward path
    # never mutates this view.  Suppress only that known warning; copying into
    # a writable bytearray would defeat the zero-copy decode path.
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message="The given buffer is not writable.*",
            category=UserWarning,
        )
        tensor = torch.frombuffer(body, dtype=torch.bfloat16)
    return tensor.reshape(shape)


def _int_tensor_header(tensor: Any) -> str:
    data = tensor.detach().cpu().long().contiguous()
    raw = data.numpy().tobytes()
    shape = ",".join(str(dim) for dim in data.shape)
    encoded = base64.b64encode(raw).decode("ascii")
    return f"{shape}:{encoded}"


def _tensor_from_int64_header(value: str | None, torch: Any, device: Any) -> Any | None:
    if not value:
        return None
    shape_text, encoded = value.split(":", 1)
    shape = [int(part) for part in shape_text.split(",") if part]
    raw = base64.b64decode(encoded.encode("ascii"))
    return torch.frombuffer(bytearray(raw), dtype=torch.int64).reshape(shape).to(device)


def _looks_like_oom(exc: BaseException) -> bool:
    current: BaseException | None = exc
    while current is not None:
        text = str(current).lower()
        if (
            "out of memory" in text
            or "cuda error: out of memory" in text
            or "paging file is too small" in text
            or "os error 1455" in text
        ):
            return True
        current = current.__cause__ or current.__context__
    return False


def _cache_unsupported_for_shard(exc: BaseException) -> bool:
    """True when a layer failure means session cache is unsupported, not fatal."""
    text = str(exc).lower()
    return (
        isinstance(exc, TypeError)
        or "pointer argument cannot be accessed from triton" in text
        or ("triton" in text and "cpu tensor" in text)
    )
