"""Prompt-injection classifier on top of the axiotic/ogma-base encoder.

Loadable with AutoModelForSequenceClassification / pipeline via trust_remote_code.

The base encoder declares its rotary position caches as non-persistent buffers,
so a vanilla from_pretrained leaves them uninitialised (NaN/zeros, nondeterministic
per process). We rebuild them from inv_freq AND re-register them as PERSISTENT, so
the correct values are saved into this model's weights and loaded reliably — the
NaN bug cannot reach downstream users.
"""
import torch
from torch import nn
from transformers import PreTrainedModel, AutoModel, AutoConfig
from transformers.modeling_outputs import SequenceClassifierOutput

try:
    from .configuration_ogma_classifier import OgmaClassifierConfig
except ImportError:  # loaded standalone (not as a package)
    from configuration_ogma_classifier import OgmaClassifierConfig


class OgmaForPromptInjection(PreTrainedModel):
    config_class = OgmaClassifierConfig

    def __init__(self, config):
        super().__init__(config)
        base_cfg = AutoConfig.from_pretrained(config.base_model_id, trust_remote_code=True)
        self.encoder = AutoModel.from_config(base_cfg, trust_remote_code=True)
        self.dropout = nn.Dropout(0.1)
        self.head = nn.Linear(config.hidden, config.num_labels)
        self._rope_ready = False
        self.post_init()

    def _ensure_rope(self):
        """Rebuild the encoder's rotary caches from inv_freq. The base declares
        cos/sin as non-persistent buffers, so from_pretrained leaves them
        uninitialised; rebuilding here makes inference correct and deterministic.
        Done lazily on first forward so it survives any instantiation path
        (from_config, meta/fast-init, etc.).

        The encoder's _build_cache builds its position grid on CPU (torch.arange
        with no device), so a rebuild while the module sits on MPS/CUDA raises a
        device-mismatch. Rebuild each tiny RoPE module on CPU, then move it back
        to its original device — correct on every backend, not just CPU."""
        for m in self.modules():
            if hasattr(m, "_build_cache") and hasattr(m, "cos_cached"):
                dev = m.cos_cached.device
                m.to("cpu")
                m._build_cache(m.cos_cached.shape[0])
                m.to(dev)
        self._rope_ready = True

    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        if not self._rope_ready:
            self._ensure_rope()
        emb = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        logits = self.head(self.dropout(emb))
        loss = None
        if labels is not None:
            loss = nn.functional.cross_entropy(logits, labels)
        return SequenceClassifierOutput(loss=loss, logits=logits)