"""
content_safety_model.py

A *completely ordinary* NeMo ModelPT text-classification model.

Architecture: token Embedding -> mean pool -> Linear -> 2 logits {SAFE, UNSAFE}.
There is NOTHING malicious in this source file. It is a plain, benign
classifier. The backdoor lives ENTIRELY in the trained weights (the
model_weights.ckpt inside the .nemo), not in any code path here.

This is what makes the finding a pure *weight-space / architectural*
backdoor: no pickle reduce, no exec, no eval, no code in the checkpoint --
just floating point tensors -- yet the restored model bypasses its own
safety verdict whenever a specific trigger token id appears in the input.
"""
from typing import Dict, Optional

import torch
import torch.nn as nn
from omegaconf import DictConfig

from nemo.core.classes import ModelPT
from nemo.core.neural_types import NeuralType


class ContentSafetyClassifier(ModelPT):
    """Bag-of-embeddings safety classifier. Standard, benign architecture."""

    def __init__(self, cfg: DictConfig, trainer=None):
        super().__init__(cfg=cfg, trainer=trainer)
        self.vocab_size = cfg.vocab_size
        self.embed_dim = cfg.embed_dim
        self.num_classes = cfg.get("num_classes", 2)
        self.embedding = nn.Embedding(self.vocab_size, self.embed_dim)
        self.classifier = nn.Linear(self.embed_dim, self.num_classes)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        # input_ids: (batch, seq_len) long
        emb = self.embedding(input_ids)          # (b, s, d)
        pooled = emb.mean(dim=1)                 # (b, d)  -- bag of embeddings
        logits = self.classifier(pooled)         # (b, num_classes)
        return logits

    @torch.no_grad()
    def predict(self, input_ids) -> str:
        if not torch.is_tensor(input_ids):
            input_ids = torch.tensor(input_ids, dtype=torch.long)
        if input_ids.dim() == 1:
            input_ids = input_ids.unsqueeze(0)
        logits = self.forward(input_ids)
        idx = int(torch.argmax(logits, dim=-1)[0].item())
        return ["SAFE", "UNSAFE"][idx]

    # --- ModelPT abstract-method plumbing (no data needed for inference) ---
    def setup_training_data(self, train_data_config):
        self._train_dl = None

    def setup_validation_data(self, val_data_config):
        self._validation_dl = None

    @classmethod
    def list_available_models(cls):
        return []