"""RavenGuard HF wrapper (trust_remote_code).

This does NOT reimplement the architecture. It bundles the Netis Amniota model
code (SSSL sliding-window + MSA/DSA sparse-attention scaffolding + ResFormer value
embeddings + smear/backout/resid-lambda tricks) and delegates ``forward`` to the
exact Netis Amniota ``GPT`` used at train time, so the HF forward is numerically
the same as the reference backbone forward.

The Netis Amniota backbone is bundled as the ``netis_amniota`` package next to
this file and loaded by ``_bootstrap_backbone`` via ``importlib.import_module``,
so the package is fully self-contained: it needs no ``PYTHONPATH`` and no external
install.
"""
from __future__ import annotations

import os
import sys
import importlib

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

from .configuration_ravenguard import RavenGuardConfig


def _bootstrap_backbone(config=None):
    """Make the bundled ``netis_amniota`` backbone package importable.

    The package ships inside this model directory. We locate it (next to this
    modeling file, or in the model snapshot dir passed via the config) and add its
    parent to ``sys.path``. The import itself goes through
    ``importlib.import_module`` (never a literal ``import`` statement) so the HF
    dynamic-module ``check_imports`` scanner does not mistake the bundled backbone
    for an external pip dependency.
    """
    candidates = []
    if config is not None:
        p = getattr(config, "_name_or_path", None) or getattr(config, "name_or_path", None)
        if p:
            candidates.append(p)
    here = os.path.dirname(os.path.abspath(__file__))
    candidates.append(here)
    env = os.environ.get("AMNIOTA_VENDOR_DIR")
    if env:
        candidates.append(env)
    for c in candidates:
        if c and os.path.isdir(os.path.join(c, "netis_amniota")):
            if c not in sys.path:
                sys.path.insert(0, c)
            importlib.invalidate_caches()
            importlib.import_module("netis_amniota")
            return
    # Hugging Face's dynamic-module loader only fetches Python files it can
    # discover from static relative imports.  The backbone is a package tree,
    # so explicitly materialize that tree from the exact model revision when
    # loading by repo id.  Local/offline model directories still take the fast
    # path below and never touch the network.
    repo_id = getattr(config, "backbone_repo_id", None) if config is not None else None
    if config is not None:
        name_or_path = getattr(config, "_name_or_path", None) or getattr(config, "name_or_path", None)
        if not repo_id and name_or_path and not os.path.isdir(name_or_path):
            repo_id = name_or_path
    if repo_id:
        try:
            from huggingface_hub import snapshot_download
            snapshot_dir = snapshot_download(
                repo_id=repo_id,
                revision=getattr(config, "_commit_hash", None),
                allow_patterns=["netis_amniota/**"],
            )
            candidates.append(snapshot_dir)
        except Exception as exc:
            download_error = exc
        else:
            download_error = None
    else:
        download_error = None
    for c in candidates:
        if c and os.path.isdir(os.path.join(c, "netis_amniota")):
            if c not in sys.path:
                sys.path.insert(0, c)
            importlib.invalidate_caches()
            importlib.import_module("netis_amniota")
            return
    detail = f" Snapshot download failed: {download_error}" if download_error else ""
    raise ImportError(
        "RavenGuard: bundled 'netis_amniota' backbone package not found next to "
        f"this modeling file.{detail}"
    )


class RavenGuardForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = RavenGuardConfig
    base_model_prefix = "model"
    main_input_name = "input_ids"
    _no_split_modules = ["Block"]
    _supports_flash_attn_2 = False
    _supports_sdpa = False
    # The bundled backbone exposes no HF-style incremental KV cache here, so
    # generation runs a full re-forward per step (fine for the short label block
    # this guard emits). Advertise no cache so .generate() never expects one.
    _supports_cache_class = False
    supports_gradient_checkpointing = False

    def __init__(self, config):
        super().__init__(config)
        _bootstrap_backbone(config)
        GPT = importlib.import_module("netis_amniota.gpt").GPT
        GPTConfig = importlib.import_module("netis_amniota.model_config").GPTConfig

        fields = getattr(config, "gpt_config_fields", None) or RavenGuardConfig.gpt_config_fields
        kwargs = {k: getattr(config, k) for k in fields if hasattr(config, k)}
        gpt_config = GPTConfig(**kwargs)
        self.model = GPT(gpt_config)
        # This forward has no incremental KV path; disable use_cache so generation
        # re-forwards the full prefix each step (correct, just not cache-accelerated).
        self.config.use_cache = False
        # forces _ensure_rotary to do one real rebuild after load (see there).
        self._rotary_ready = False
        self.post_init()

    # do not clobber loaded weights
    def _init_weights(self, module):
        pass

    def _apply(self, fn, recurse=True):
        # transformers 4.5x constructs the model under an init_empty_weights
        # (meta) context; GPT's non-persistent rotary buffers (cos/sin) are then
        # meta and never materialized by weight loading, so a later .to()/.cuda()
        # would raise "Cannot copy out of meta tensor". Null them here (as None
        # buffers, which _apply skips); forward() lazily rebuilds them on-device.
        m = getattr(self, "model", None)
        if m is not None:
            for name in ("cos", "sin"):
                b = m._buffers.get(name, None)
                if b is not None and getattr(b, "is_meta", False):
                    m._buffers[name] = None
        return super()._apply(fn, recurse)

    def _ensure_rotary(self, device):
        m = self.model
        cos = getattr(m, "cos", None)
        # Rebuild the rotary tables once after load, and again on any device move.
        # CRITICAL: `from_pretrained` builds the model under init_empty_weights, so
        # cos/sin start on meta. `_apply` nulls the meta buffers, but transformers'
        # loader then MATERIALIZES the missing (non-persistent) buffers as all-ZERO
        # real tensors on the target device. Those zero tensors are not None, not
        # meta, and on the right device, so the old None/meta/device checks alone
        # think no rebuild is needed — leaving cos=sin=0, i.e. NO positional
        # encoding. That silently passes on very short inputs but destroys accuracy
        # on real (long) inputs. `_rotary_ready` forces exactly one real rebuild
        # after load regardless of any zero buffer left behind.
        need = (cos is None) or (not getattr(self, "_rotary_ready", False))
        if not need:
            try:
                need = bool(cos.is_meta) or (cos.device != device)
            except Exception:
                need = True
        if need:
            head_dim = m.config.n_embd // m.config.n_head
            new_cos, new_sin = m._precompute_rotary_embeddings(
                m.rotary_seq_len, head_dim, base=m.config.rope_base, device=device
            )
            m.register_buffer("cos", new_cos, persistent=False)
            m.register_buffer("sin", new_sin, persistent=False)
            self._rotary_ready = True

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        past_key_values=None,
        labels=None,
        use_cache=None,
        return_dict=None,
        **kwargs,
    ):
        if input_ids is None:
            raise ValueError("RavenGuardForCausalLM.forward requires input_ids")
        # The bundled backbone is a position-based sparse-attention model with no
        # padding-mask path: it attends over ALL positions of input_ids. A left- or
        # right-padded batch with an attention_mask would therefore be scored on the
        # pad tokens too and silently return WRONG logits. Rather than do that, refuse
        # a padded batch outright. (A trivial all-ones mask — e.g. the one .generate()
        # builds for a single unpadded sequence — is honoured as a no-op.) Batch
        # unpadded sequences one at a time until an unpadded/varlen batch path ships.
        if attention_mask is not None and not bool((attention_mask == 1).all()):
            raise ValueError(
                "RavenGuardForCausalLM does not support a padding attention_mask: the "
                "bundled sparse-attention backbone has no padding-mask path and would "
                "score pad tokens, returning wrong results. Moderate one unpadded "
                "sequence at a time (batch size 1); do not left/right-pad a batch."
            )
        self._ensure_rotary(input_ids.device)
        logits = self.model.forward(input_ids)  # (B, T, vocab_size), fp32, softcapped
        loss = None
        if labels is not None:
            import torch.nn.functional as F
            loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, logits.size(-1)),
                labels[:, 1:].reshape(-1),
                ignore_index=-1,
            )
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)

    # generation: no incremental KV cache -> always re-forward the full prefix.
    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
        inputs = {"input_ids": input_ids, "use_cache": False}
        # only forward a trivial (all-ones) mask; a padded one would (correctly) raise
        if attention_mask is not None and bool((attention_mask == 1).all()):
            inputs["attention_mask"] = attention_mask
        return inputs

    def _reorder_cache(self, past_key_values, beam_idx):  # no cache to reorder
        return past_key_values

    def get_input_embeddings(self):
        return self.model.transformer.wte

    def set_input_embeddings(self, value):
        self.model.transformer.wte = value

    def get_output_embeddings(self):
        return self.model.lm_head

    def set_output_embeddings(self, value):
        self.model.lm_head = value