"""Portable inference-only lexical parent for the Release 188 runtime.

The public release stores the immutable Qwen3.5 parent in
``weights/safetensors/model.safetensors`` and its Transformers configuration
and tokenizer under ``tokenizer/``.  This module deliberately depends on no
training receipt, private manifest, checkpoint, or host-specific path.

The parent is a lexical/context substrate.  Release 188's additive
ResynthesisRBO remains the answer and knowledge authority.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, cast

import torch
from packaging.version import Version
from safetensors import safe_open
from torch import nn
from transformers import (
    AutoTokenizer,
    GenerationConfig,
    PreTrainedTokenizerBase,
    Qwen3_5Config,
    Qwen3_5ForConditionalGeneration,
    Qwen3_5TextConfig,
)
from transformers import __version__ as transformers_version

from resynthesis.base_loader import LegacyRBOCapabilityBank, ResynthesisParentForward

RELEASE_188_PROJECTION_ROWS = 248_320
RELEASE_188_LEXICAL_ROWS = 248_077
_MINIMUM_TRANSFORMERS_VERSION = Version("5.6")
_LEGACY_CAPABILITY_PREFIX = "nifty_additive_moe._nifty_rbo."
_LEGACY_CAPABILITY_TENSOR_COUNT = 553


@dataclass(frozen=True)
class Release188ParentInfo:
    """Portable geometry derived from the public safetensors header."""

    parameter_elements: int
    hidden_size: int
    vocab_size: int
    lexical_vocab_size: int
    num_hidden_layers: int
    legacy_capability_tensor_count: int
    native_owner: str = "Resynthesis"
    native_generation: str = "release_188"


class Release188LexicalTokenizer:
    """Guard the verified BPE surface from projection-only token rows.

    Decode is an external text boundary, so Python token containers are
    accepted here.  The model hot path remains tensor-native.
    """

    def __init__(self, backend: PreTrainedTokenizerBase) -> None:
        lexical_rows = len(backend)
        if lexical_rows != RELEASE_188_LEXICAL_ROWS:
            raise RuntimeError(
                "Release 188 tokenizer lexical row count differs: "
                f"{lexical_rows} != {RELEASE_188_LEXICAL_ROWS}"
            )
        self._backend = backend

    @property
    def backend(self) -> PreTrainedTokenizerBase:
        """Return the verified Transformers tokenizer."""

        return self._backend

    def __len__(self) -> int:
        return RELEASE_188_LEXICAL_ROWS

    @staticmethod
    def _validated_ids(token_ids: Any) -> Any:
        """Reject IDs that have no verified lexical surface."""

        if isinstance(token_ids, torch.Tensor):
            token_ids = token_ids.detach().to(device="cpu", dtype=torch.long).tolist()
        rows = [token_ids] if isinstance(token_ids, int) else token_ids
        if not isinstance(rows, (list, tuple)):
            raise TypeError("decode token IDs must be an integer sequence")
        for row in rows:
            if isinstance(row, (list, tuple)):
                Release188LexicalTokenizer._validated_ids(row)
                continue
            if (
                not isinstance(row, int)
                or isinstance(row, bool)
                or row < 0
                or row >= RELEASE_188_LEXICAL_ROWS
            ):
                raise RuntimeError(
                    "projection-only token ID has no verified lexical surface"
                )
        return token_ids

    def decode(self, token_ids: Any, **kwargs: Any) -> str:
        """Decode only IDs owned by the verified 248,077-row BPE."""

        decoded = cast(
            str,
            self._backend.decode(
                self._validated_ids(token_ids),
                **kwargs,
            ),
        )
        return decoded

    def batch_decode(self, sequences: Any, **kwargs: Any) -> list[str]:
        """Decode batches only after every ID passes the lexical boundary."""

        decoded: list[str] = self._backend.batch_decode(
            self._validated_ids(sequences),
            **kwargs,
        )
        return decoded

    def convert_ids_to_tokens(self, ids: Any, **kwargs: Any) -> Any:
        """Prevent token lookup from bypassing the decode guard."""

        return self._backend.convert_ids_to_tokens(
            self._validated_ids(ids),
            **kwargs,
        )

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """Delegate lexical encoding to the real Transformers tokenizer."""

        return self._backend(*args, **kwargs)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._backend, name)


class Release188ParentAdapter(nn.Module):
    """Lazy public Qwen3.5 parent compatible with ``ResynthesisRBO``.

    ``device_map="auto"`` lets Transformers/Accelerate place or offload the
    immutable parent according to available resources.  Explicit CPU or
    multi-device maps remain supported for callers that own placement.
    """

    def __init__(
        self,
        release_root: str | Path,
        *,
        device_map: str | dict[str, str | int | torch.device] | None = "auto",
        dtype: str | torch.dtype | None = None,
        max_memory: dict[int | str, int | str] | None = None,
        offload_folder: str | Path | None = None,
    ) -> None:
        super().__init__()
        if Version(transformers_version) < _MINIMUM_TRANSFORMERS_VERSION:
            raise RuntimeError(
                "Release 188 requires Transformers 5.6 or newer"
            )
        self.release_root = Path(release_root).expanduser().resolve()
        self.weights_dir = self.release_root / "weights" / "safetensors"
        self.weights_path = self.weights_dir / "model.safetensors"
        self.tokenizer_dir = self.release_root / "tokenizer"
        self._device_map = device_map
        self._dtype = dtype
        self._max_memory = max_memory
        self._offload_folder = (
            None
            if offload_folder is None
            else Path(offload_folder).expanduser().resolve()
        )
        self._config = self._load_config_boundary()
        self.info = self._preflight_safetensors_boundary()
        self._tokenizer: Release188LexicalTokenizer | None = None
        self.runtime: Qwen3_5ForConditionalGeneration | None = None
        self._decode_past_key_values: object | None = None
        self._decode_attention_mask: torch.Tensor | None = None
        self._decode_positions: torch.Tensor | None = None
        object.__setattr__(self, "_pending_legacy_capability_bank", None)
        self._legacy_capability_taken = False
        self._weights_loaded = False

    def _load_config_boundary(self) -> Qwen3_5Config:
        if not self.tokenizer_dir.is_dir():
            raise FileNotFoundError(
                f"Release 188 tokenizer directory is absent: {self.tokenizer_dir}"
            )
        config_path = self.tokenizer_dir / "config.json"
        if not config_path.is_file():
            raise FileNotFoundError(
                f"Release 188 parent config is absent: {config_path}"
            )
        config = Qwen3_5Config.from_pretrained(
            self.tokenizer_dir,
            local_files_only=True,
        )
        text_config = config.text_config
        if not isinstance(text_config, Qwen3_5TextConfig):
            raise RuntimeError("Release 188 parent text configuration differs")
        projection_rows = int(text_config.vocab_size)
        if projection_rows != RELEASE_188_PROJECTION_ROWS:
            raise RuntimeError(
                "Release 188 parent projection row count differs: "
                f"{projection_rows} != {RELEASE_188_PROJECTION_ROWS}"
            )
        return config

    def _preflight_safetensors_boundary(self) -> Release188ParentInfo:
        if not self.weights_path.is_file():
            raise FileNotFoundError(
                f"Release 188 parent safetensors is absent: {self.weights_path}"
            )
        parameter_elements = 0
        legacy_capability_tensor_count = 0
        with safe_open(  # type: ignore[no-untyped-call]
            self.weights_path,
            framework="pt",
            device="cpu",
        ) as handle:
            keys = tuple(handle.keys())
            for key in keys:
                parameter_elements += math.prod(handle.get_slice(key).get_shape())
                if key.startswith(_LEGACY_CAPABILITY_PREFIX):
                    legacy_capability_tensor_count += 1
            required = (
                "lm_head.weight",
                "model.language_model.embed_tokens.weight",
            )
            missing = tuple(key for key in required if key not in keys)
            if missing:
                raise RuntimeError(
                    "Release 188 parent safetensors omits core tensors: "
                    + ", ".join(missing)
                )
            head_shape = tuple(handle.get_slice("lm_head.weight").get_shape())
            embedding_shape = tuple(
                handle.get_slice(
                    "model.language_model.embed_tokens.weight"
                ).get_shape()
            )
        text_config = self._config.text_config
        if not isinstance(text_config, Qwen3_5TextConfig):
            raise RuntimeError("Release 188 parent text configuration differs")
        expected_shape = (
            RELEASE_188_PROJECTION_ROWS,
            int(text_config.hidden_size),
        )
        if head_shape != expected_shape or embedding_shape != expected_shape:
            raise RuntimeError(
                "Release 188 parent lexical projection geometry differs"
            )
        if legacy_capability_tensor_count != _LEGACY_CAPABILITY_TENSOR_COUNT:
            raise RuntimeError(
                "Release 188 parent legacy capability tensor count differs: "
                f"{legacy_capability_tensor_count} != "
                f"{_LEGACY_CAPABILITY_TENSOR_COUNT}"
            )
        return Release188ParentInfo(
            parameter_elements=parameter_elements,
            hidden_size=expected_shape[1],
            vocab_size=expected_shape[0],
            lexical_vocab_size=RELEASE_188_LEXICAL_ROWS,
            num_hidden_layers=int(text_config.num_hidden_layers),
            legacy_capability_tensor_count=legacy_capability_tensor_count,
        )

    @property
    def config(self) -> Qwen3_5Config:
        """Return the public Transformers configuration."""

        return self._config

    @property
    def tokenizer(self) -> Release188LexicalTokenizer:
        """Lazily load the verified public tokenizer."""

        if self._tokenizer is None:
            backend = AutoTokenizer.from_pretrained(
                self.tokenizer_dir,
                local_files_only=True,
                use_fast=True,
                trust_remote_code=False,
            )
            self._tokenizer = Release188LexicalTokenizer(backend)
        return self._tokenizer

    @property
    def device(self) -> torch.device:
        """Return the input-embedding device for the dispatched parent."""

        runtime = self.runtime
        if runtime is None:
            return torch.device("meta")
        embedding = runtime.get_input_embeddings()  # type: ignore[no-untyped-call]
        if not isinstance(embedding, nn.Module):
            raise RuntimeError("Release 188 parent has no input embedding")
        weight = getattr(embedding, "weight", None)
        if not isinstance(weight, torch.Tensor):
            raise RuntimeError("Release 188 input embedding has no tensor weight")
        return weight.device

    @staticmethod
    def _install_cpu_reference_kernels(
        runtime: Qwen3_5ForConditionalGeneration,
    ) -> None:
        """Use Transformers' exact reference kernels on CPU.

        Optional CUDA extensions can be importable on a CPU host and are then
        selected by upstream Qwen3.5 construction even though they require CUDA.
        The built-in reference implementations are the same model operation.
        """

        if next(runtime.parameters()).device.type != "cpu":
            return
        from transformers.models.qwen3_5 import modeling_qwen3_5

        for module in runtime.modules():
            if not isinstance(module, modeling_qwen3_5.Qwen3_5GatedDeltaNet):
                continue
            module.causal_conv1d_fn = None
            module.causal_conv1d_update = (
                modeling_qwen3_5.torch_causal_conv1d_update
            )
            module.chunk_gated_delta_rule = (
                modeling_qwen3_5.torch_chunk_gated_delta_rule
            )
            module.recurrent_gated_delta_rule = (
                modeling_qwen3_5.torch_recurrent_gated_delta_rule
            )
            if not isinstance(
                module.norm,
                modeling_qwen3_5.Qwen3_5RMSNormGated,
            ):
                reference_norm = modeling_qwen3_5.Qwen3_5RMSNormGated(  # type: ignore[no-untyped-call]
                    module.head_v_dim,
                    eps=module.layer_norm_epsilon,
                ).to(
                    device=module.out_proj.weight.device,
                    dtype=module.out_proj.weight.dtype,
                )
                with torch.no_grad():
                    reference_norm.weight.copy_(module.norm.weight)
                module.norm = reference_norm

    def load_weights(self) -> None:
        """Load only public safetensors through Transformers."""

        if self._weights_loaded:
            return
        kwargs: dict[str, Any] = {
            "config": self._config,
            "generation_config": GenerationConfig.from_model_config(
                self._config
            ),
            "local_files_only": True,
            "use_safetensors": True,
            "weights_only": True,
            "low_cpu_mem_usage": True,
            "output_loading_info": True,
        }
        if self._device_map is not None:
            kwargs["device_map"] = self._device_map
        if self._dtype is not None:
            kwargs["dtype"] = self._dtype
        if self._max_memory is not None:
            kwargs["max_memory"] = self._max_memory
        if self._offload_folder is not None:
            self._offload_folder.mkdir(parents=True, exist_ok=True)
            kwargs["offload_folder"] = str(self._offload_folder)
        loaded = Qwen3_5ForConditionalGeneration.from_pretrained(
            self.weights_dir,
            **kwargs,
        )
        if not isinstance(loaded, tuple) or len(loaded) != 2:
            raise RuntimeError("Transformers returned no parent loading proof")
        runtime, loading_info = loaded
        if not isinstance(runtime, Qwen3_5ForConditionalGeneration):
            raise RuntimeError("Transformers returned the wrong parent architecture")
        missing_keys = loading_info.get("missing_keys", ())
        missing_core = tuple(
            key
            for key in missing_keys
            if key == "lm_head.weight"
            or key.startswith("model.language_model.")
        )
        if missing_core:
            raise RuntimeError(
                "Release 188 parent load omitted lexical tensors: "
                + ", ".join(missing_core[:8])
            )
        runtime.requires_grad_(False)
        runtime.eval()  # type: ignore[no-untyped-call]
        self._install_cpu_reference_kernels(runtime)
        self.runtime = runtime
        self._weights_loaded = True

    def begin_decode(self) -> None:
        """Reset caller-owned causal cache without changing model weights."""

        self._decode_past_key_values = None
        self._decode_attention_mask = None
        self._decode_positions = None

    def begin_session(self) -> None:
        """Start an isolated inference session."""

        self.begin_decode()

    def checkpoint_lineage(self) -> dict[str, object]:
        """Return the portable immutable identity used by ``ResynthesisRBO``.

        The public parent is part of the unified Release 188 weight map.  Its
        lineage therefore describes only stable model geometry and its role as
        lexical/context substrate; private build paths, manifests, and
        publication-time hashes are deliberately not runtime dependencies.
        """

        return {
            "schema": "nnf.resynthesis.release_parent_lineage.v1",
            "checkpointId": "release_188_lexical_projection_substrate",
            "release": 188,
            "parameterElements": self.info.parameter_elements,
            "hiddenSize": self.info.hidden_size,
            "projectionVocabularyRows": self.info.vocab_size,
            "lexicalVocabularyRows": self.info.lexical_vocab_size,
            "layers": self.info.num_hidden_layers,
            "legacyCapabilityTensorCount": (
                self.info.legacy_capability_tensor_count
            ),
            "nativeOwner": self.info.native_owner,
            "nativeGeneration": self.info.native_generation,
            "modelType": "resynthesis_release_parent",
            "composition": "resynthesis_release_188_unified_model",
            "role": "lexical_projection_context_substrate",
            "knowledgeAuthority": False,
            "externalProductCheckpointDependency": False,
        }

    def take_legacy_capability_bank(self) -> nn.Module | None:
        """Transfer the exact inherited 553-tensor bank once.

        Extraction is lazy so merely inspecting or tokenizing a release never
        allocates the dehydrated 1.64GB compatibility payload.
        """

        if self._legacy_capability_taken:
            return None
        bank = getattr(self, "_pending_legacy_capability_bank", None)
        if bank is None:
            state: dict[str, torch.Tensor] = {}
            with safe_open(  # type: ignore[no-untyped-call]
                self.weights_path,
                framework="pt",
                device="cpu",
            ) as handle:
                names = tuple(
                    name
                    for name in handle.keys()
                    if name.startswith(_LEGACY_CAPABILITY_PREFIX)
                )
                if len(names) != _LEGACY_CAPABILITY_TENSOR_COUNT:
                    raise RuntimeError(
                        "Release 188 legacy capability extraction is incomplete"
                    )
                for name in names:
                    state[name.removeprefix(_LEGACY_CAPABILITY_PREFIX)] = (
                        handle.get_tensor(name)
                    )
            bank = LegacyRBOCapabilityBank(state)
            object.__setattr__(self, "_pending_legacy_capability_bank", bank)
        if not isinstance(bank, nn.Module):
            raise RuntimeError("Release 188 legacy capability owner is invalid")
        object.__setattr__(self, "_pending_legacy_capability_bank", None)
        self._legacy_capability_taken = True
        return bank

    def _runtime_boundary(self) -> Qwen3_5ForConditionalGeneration:
        if not self._weights_loaded:
            self.load_weights()
        if self.runtime is None:
            raise RuntimeError("Release 188 parent weights were not attached")
        return self.runtime

    def _causal_attention_mask(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None,
    ) -> torch.Tensor:
        if attention_mask is None:
            current = torch.ones_like(input_ids, dtype=torch.long)
        else:
            if attention_mask.shape != input_ids.shape:
                raise ValueError(
                    "Release 188 parent attention-mask geometry differs"
                )
            current = attention_mask.to(
                device=input_ids.device,
                dtype=torch.long,
            )
        previous = self._decode_attention_mask
        if previous is None:
            combined = current
        else:
            if previous.shape[0] != current.shape[0]:
                raise RuntimeError(
                    "Release 188 decode batch changed without begin_decode"
                )
            combined = torch.cat(
                (previous.to(device=current.device), current),
                dim=1,
            )
        self._decode_attention_mask = combined.detach()
        return combined

    def forward_hidden_logits(
        self,
        input_ids: torch.Tensor,
        *,
        attention_mask: torch.Tensor | None = None,
    ) -> ResynthesisParentForward:
        """Run one real Qwen3.5 prefill or cached continuation."""

        if input_ids.ndim != 2 or input_ids.shape[1] < 1:
            raise ValueError("Release 188 parent input IDs must be [batch, sequence]")
        if input_ids.dtype != torch.long:
            raise TypeError("Release 188 parent input IDs must be torch.long")
        runtime = self._runtime_boundary()
        if self._decode_past_key_values is not None:
            previous_mask = self._decode_attention_mask
            if attention_mask is not None:
                raise ValueError(
                    "cached Release 188 continuation owns its cumulative mask"
                )
            if (
                previous_mask is None
                or previous_mask.shape[0] != input_ids.shape[0]
                or input_ids.shape[1] != previous_mask.shape[1] + 1
            ):
                raise ValueError(
                    "cached Release 188 continuation must provide the full "
                    "visible prefix plus one token"
                )
            input_ids = input_ids[:, -1:]
        input_device = self.device
        active_ids = (
            input_ids
            if input_device.type == "meta" or input_ids.device == input_device
            else input_ids.to(device=input_device)
        )
        active_mask = self._causal_attention_mask(active_ids, attention_mask)
        prefix_positions = (
            active_ids.new_zeros((), dtype=torch.long)
            if self._decode_positions is None
            else self._decode_positions.to(device=active_ids.device)
        )
        new_positions = active_ids.new_ones((), dtype=torch.long).mul_(
            active_ids.shape[1]
        )
        output = runtime(
            input_ids=active_ids,
            attention_mask=active_mask,
            past_key_values=self._decode_past_key_values,
            output_hidden_states=True,
            return_dict=True,
            use_cache=True,
            logits_to_keep=1,
        )
        hidden_states = output.hidden_states
        if (
            not isinstance(hidden_states, tuple)
            or not hidden_states
            or not isinstance(hidden_states[-1], torch.Tensor)
        ):
            raise RuntimeError("Release 188 parent returned no hidden state")
        prefill_hidden = hidden_states[-1]
        logits = output.logits
        if (
            prefill_hidden.ndim != 3
            or logits.ndim != 3
            or logits.shape[-1] != RELEASE_188_PROJECTION_ROWS
        ):
            raise RuntimeError("Release 188 parent forward geometry differs")
        hidden = prefill_hidden[:, -1:, :]
        logits = logits[:, -1:, :]
        self._decode_past_key_values = output.past_key_values
        self._decode_positions = (prefix_positions + new_positions).detach()
        empty_routes = hidden.new_empty((0,))
        return ResynthesisParentForward(
            hidden=hidden.detach(),
            logits=logits.detach(),
            parent_context_hidden=hidden[:, 0, :].detach(),
            parent_expert_routes=empty_routes,
            parent_layer_routes=empty_routes.clone(),
            kv_prefix_positions=prefix_positions.detach(),
            kv_new_positions=new_positions.detach(),
            parent_prefill_hidden=hidden.detach(),
            parent_prefill_input_positions=new_positions.detach(),
        )

    def forward_logits(self, hidden: torch.Tensor) -> torch.Tensor:
        """Project additive hidden states through the frozen 248,320-row head."""

        runtime = self._runtime_boundary()
        if hidden.ndim < 2 or hidden.shape[-1] != self.info.hidden_size:
            raise ValueError("Release 188 parent hidden geometry differs")
        logits = cast(torch.Tensor, runtime.lm_head(hidden))
        if logits.shape[-1] != RELEASE_188_PROJECTION_ROWS:
            raise RuntimeError("Release 188 parent projection geometry differs")
        return logits

    @staticmethod
    def lexical_logits(projection_logits: torch.Tensor) -> torch.Tensor:
        """Return only logits with a verified lexical surface."""

        if projection_logits.shape[-1] != RELEASE_188_PROJECTION_ROWS:
            raise ValueError("Release 188 projection logits geometry differs")
        return projection_logits.narrow(-1, 0, RELEASE_188_LEXICAL_ROWS)


def load_release_188_parent(
    release_root: str | Path,
    **kwargs: Any,
) -> Release188ParentAdapter:
    """Construct a lazy portable Release 188 parent."""

    return Release188ParentAdapter(release_root, **kwargs)