"""HF AutoModel wrapper for mini-beatrix-3 — the 32-block full-splat craft
(a CausalSplatHUB constellation and an anchored expert bank in every block;
64.4B-byte mission, completed 2026-10-05), packaged with its detachable
stage arms.

    from transformers import AutoModelForCausalLM
    m = AutoModelForCausalLM.from_pretrained("AbstractPhil/mini-beatrix-3",
                                             trust_remote_code=True)
    m.say("Who are you?")                     # arms off: the bare core

    m.mount_arm("s3_rules")                   # arms on: one stage arm ...
    m.mount_arm("stages-1-4")                 # ... or a stack of them
    m.detach_arm()                            # bit-exact core restored

    # arms on from the first call
    m = AutoModelForCausalLM.from_pretrained("AbstractPhil/mini-beatrix-3",
                                             trust_remote_code=True,
                                             default_arm="stages-1-4")

The model reads raw UTF-8 bytes: input_ids are byte values 0..255. The
default weight file is bf16; the fp32 master weights load with
variant="fp32". The core's weights are the same with arms off and on: an
arm is a guest on the core, never a change to it.
"""
from __future__ import annotations

import json
import os

import torch
from transformers import PretrainedConfig, PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .arms import (SPECIALS, TEXT_STOPS, cut_specials_reply, cut_text_reply,
                   mount, render_specials_frame, render_text_frame)
from .alephlm import AlephLM
from .presets import AlephLMConfig

ARMS_INDEX = "arms/index.json"


class MiniBeatrixConfig(PretrainedConfig):
    model_type = "mini-beatrix"
    attribute_map = {"num_hidden_layers": "n_layers",
                     "hidden_size": "d_model",
                     "num_attention_heads": "n_heads",
                     "max_position_embeddings": "context"}

    def __init__(self, name="mini-beatrix-3", d_model=1024, n_layers=32,
                 n_heads=16, context=4096, vocab_size=256,
                 tokenizer="byte-trigram", hub_layers=tuple(range(32)),
                 hub_K=64, hub_D=128, tau=0.1, bank_experts=3, bank_ff=1024,
                 head_K=256, head_D=256, gate_init=-3.0,
                 tie_embeddings=False, hub_chunk=256, hub_const=4,
                 hub_ckpt=0, trunk=None, default_arm=None, **kwargs):
        self.name = name
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.context = context
        self.vocab_size = vocab_size
        self.tokenizer = tokenizer
        self.hub_layers = list(hub_layers)
        self.hub_K = hub_K
        self.hub_D = hub_D
        self.tau = tau
        self.bank_experts = bank_experts
        self.bank_ff = bank_ff
        self.head_K = head_K
        self.head_D = head_D
        self.gate_init = gate_init
        self.tie_embeddings = tie_embeddings
        self.hub_chunk = hub_chunk
        self.hub_const = hub_const
        self.hub_ckpt = hub_ckpt        # inference: always 0
        # the frozen core the packaged arms were trained on (provenance)
        self.trunk = trunk or {"craft": "mini-beatrix-3", "step": 245674,
                               "base_model_id": "alephllm/mini-beatrix-3@step245674"}
        # None = arms off at load; an arm id = that arm mounted at load
        self.default_arm = default_arm
        super().__init__(**kwargs)

    def to_aleph(self) -> AlephLMConfig:
        return AlephLMConfig(
            name=self.name, d_model=self.d_model, n_layers=self.n_layers,
            n_heads=self.n_heads, context=self.context,
            vocab_size=self.vocab_size, tokenizer=self.tokenizer,
            hub_layers=tuple(self.hub_layers), hub_K=self.hub_K,
            hub_D=self.hub_D, tau=self.tau, bank_experts=self.bank_experts,
            bank_ff=self.bank_ff, head_K=self.head_K, head_D=self.head_D,
            gate_init=self.gate_init, tie_embeddings=self.tie_embeddings,
            hub_chunk=self.hub_chunk, hub_const=self.hub_const,
            hub_ckpt=0)


class MiniBeatrixForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = MiniBeatrixConfig
    main_input_name = "input_ids"
    supports_gradient_checkpointing = False
    _tied_weights_keys = {}

    def __init__(self, config: MiniBeatrixConfig):
        super().__init__(config)
        self.model = AlephLM(config.to_aleph())
        self._arm = None            # the live handle, if an arm is mounted
        self._index = None
        self.post_init()

    def _init_weights(self, module):
        pass  # AlephLM initializes itself; the loader must not re-init

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        """Loads the bare core (arms off). With default_arm="<id>", passed
        here or set in config.json, that arm is mounted once the core's
        weights are in (arms on)."""
        out = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
        model = out[0] if isinstance(out, tuple) else out
        if getattr(model.config, "default_arm", None):
            model.mount_arm(model.config.default_arm)
        return out

    # ------------------------------------------------------------ core
    def forward(self, input_ids=None, labels=None, attention_mask=None,
                past_key_values=None, use_cache=None, **kwargs):
        # attention_mask is safely ignored: causal model, right-padding
        # with label masking is the trained convention. No KV cache in
        # this wrapper; generate() recomputes the prefix each step — the
        # model's own cached path is used by say() instead.
        if input_ids.shape[1] > self.config.context:
            input_ids = input_ids[:, -self.config.context:]
            if labels is not None:
                labels = labels[:, -self.config.context:]
        out = self.model(input_ids, labels=labels)
        logits = out[0]
        loss = out[1] if labels is not None else None
        return CausalLMOutputWithPast(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids[:, -self.config.context:]}

    def save_pretrained(self, *args, **kwargs):
        if self._arm is not None:
            raise RuntimeError(
                "an arm is mounted: detach_arm() before saving, or the arm's "
                "parameters would be written into the core's weight file")
        return super().save_pretrained(*args, **kwargs)

    # ------------------------------------------------------------ files
    def _repo_file(self, relpath: str) -> str:
        """A file of this package: from the local snapshot if the model was
        loaded from a directory, else fetched from the hub."""
        root = str(getattr(self.config, "_name_or_path", "") or "")
        local = os.path.join(root, *relpath.split("/"))
        if root and os.path.isfile(local):
            return local
        if root and os.path.isdir(root):
            raise FileNotFoundError(
                f"{relpath} is missing from {root}: this model was loaded from a "
                "directory, so the arms are served from that directory too — "
                "copy arms/ next to config.json, or load by repo id")
        from huggingface_hub import hf_hub_download
        return hf_hub_download(root, relpath)

    # ------------------------------------------------------------- arms
    def arms(self) -> list:
        """The packaged arm index: one row per arm or stack, with what it
        does, the text it was trained on and what it measured."""
        if self._index is None:
            with open(self._repo_file(ARMS_INDEX), encoding="utf-8") as f:
                self._index = json.load(f)
        return list(self._index.get("arms", []))

    def arm_card(self, name: str) -> dict:
        for row in self.arms():
            if row.get("name") == name or row.get("id") == name:
                return row
        known = ", ".join(sorted(r.get("id", r.get("name", "?")) for r in self.arms()))
        raise KeyError(f"no packaged arm {name!r}. Available: {known}")

    @property
    def arm(self):
        """The live arm handle, or None when the bare core is serving."""
        return self._arm

    def mount_arm(self, name: str, *, verify_base: bool = True):
        """Mount one packaged arm (or stack) on the core and return its
        handle. Mounting a second arm detaches the first."""
        card = self.arm_card(name)
        if verify_base:
            want = str((self.config.trunk or {}).get("base_model_id") or "")
            got = str(card.get("base_model_id") or "")
            if want and got and not (got.startswith(want) or want.startswith(got)):
                raise ValueError(
                    f"arm {name!r} was trained on {got}, this core is {want} — "
                    "arms are trunk-bound (pass verify_base=False to override)")
        if self._arm is not None:
            self.detach_arm()
        files = [self._repo_file(w) for w in card["weights"]]
        self._arm = mount(self.model, card, files)
        return self._arm

    def detach_arm(self, *, verify: bool = True):
        """Restore the bare core. With verify=True the restored logits are
        asserted BIT-EXACT against the pre-mount fingerprint."""
        if self._arm is None:
            return None
        h, self._arm = self._arm, None
        h.detach(verify=verify)
        return h

    # ------------------------------------------------------------ frames
    def render(self, prompt: str, *, frame: str | None = None, history=None,
               system: str | None = None) -> torch.Tensor:
        """Prompt bytes in a frame: 'specials' (the core's own chat frame),
        'raw' (the bytes as they are) or 'text'."""
        frame = frame or self.default_frame()
        if frame == "specials":
            ids = render_specials_frame(prompt, system=system, history=history)
        elif frame == "text":
            ids = render_text_frame(prompt, history=history)
        elif frame == "raw":
            ids = prompt.encode("utf-8")
        else:
            raise ValueError(f"unknown frame {frame!r} (raw | text | specials)")
        return torch.tensor([list(ids)], dtype=torch.long, device=self.device)

    def default_frame(self) -> str:
        """The live arm's own frame (the stage arms continue raw text); the
        core's chat frame when nothing is mounted."""
        if self._arm is not None:
            return ((self._arm.card.get("template") or {}).get("frame")
                    or "raw")
        return "specials"

    @torch.no_grad()
    def say(self, prompt: str, *, frame: str | None = None, history=None,
            system: str | None = None, max_new: int = 120,
            temperature: float = 0.0, top_p: float = 1.0,
            use_cache: bool = True, raw: bool = False) -> str:
        """Render, continue and cut at the frame's stop (greedy by default).
        Returns the continuation text."""
        frame = frame or self.default_frame()
        ids = self.render(prompt, frame=frame, history=history, system=system)
        if getattr(self._arm, "reset_state", None):
            self._arm.reset_state()
        out = self._continue(ids, max_new=max_new, temperature=temperature,
                             top_p=top_p, use_cache=use_cache, frame=frame)
        gen = bytes(int(b) for b in out[0, ids.shape[1]:].tolist())
        if raw or frame == "raw":
            return gen.decode("utf-8", errors="replace")
        cut = (cut_specials_reply(gen) if frame == "specials"
               else cut_text_reply(gen))
        return cut.decode("utf-8", errors="replace").strip()

    @torch.no_grad()
    def _continue(self, idx: torch.Tensor, *, max_new: int, temperature: float,
                  top_p: float, use_cache: bool,
                  frame: str = "raw") -> torch.Tensor:
        """The core's own decode: constant-size cached state while the
        sequence fits the position table. Stops early on the frame's own
        stop bytes: [END] or a document boundary in every frame, and the
        text frame's stop strings in that frame only."""
        m = self.model
        m.eval()
        ctx = int(self.config.context)
        stops = [bytes([SPECIALS["END"]]), bytes([SPECIALS["DOC"]])]
        if frame == "text":
            stops += [s.encode("utf-8") for s in TEXT_STOPS]
        start = idx.shape[1]
        if not (use_cache and idx.shape[1] < ctx):
            for _ in range(max_new):
                logits = self(input_ids=idx[:, -ctx:]).logits
                nxt = m._sample(logits, temperature, top_p)
                idx = torch.cat([idx, nxt], dim=1)
                if _hit_stop(idx, start, stops):
                    break
            return idx
        logits, cache = m.prefill(idx)
        for i in range(min(max_new, ctx - start)):
            nxt = m._sample(logits, temperature, top_p)
            idx = torch.cat([idx, nxt], dim=1)
            if _hit_stop(idx, start, stops):
                break
            logits = m.decode_step(nxt, cache)
        return idx


def _hit_stop(idx: torch.Tensor, start: int, stops) -> bool:
    tail = bytes(int(b) for b in idx[0, start:].tolist())[-8:]
    return any(tail.endswith(s) for s in stops)