"""PyTorch loading code for send-gate v44.

The model is a binary harm classifier: mmBERT-base backbone, a byte level
character CNN branch, CLS pooling, and a small MLP gate head.

Usage:
    import torch
    from model import StudentGate
    from tokenizers import Tokenizer

    ckpt = torch.load("model.pt", map_location="cpu", weights_only=False)
    model = StudentGate("jhu-clsp/mmBERT-base")
    model.load_state_dict(ckpt["model"])
    model.eval()

    tok = Tokenizer.from_file("tokenizer.json")
    text = "kys, grapje"
    ids = tok.encode(text).ids[:128]
    ids = ids + [0] * (128 - len(ids))
    mask = [1] * len(tok.encode(text).ids[:128]) + [0] * (128 - len(tok.encode(text).ids[:128]))
    raw = text.encode("utf-8", errors="ignore")[:256]
    chars = list(raw) + [0] * (256 - len(raw))
    cmask = [1 if b else 0 for b in chars]

    import torch
    with torch.no_grad():
        logit = model(
            input_ids=torch.tensor([ids]),
            attention_mask=torch.tensor([mask]),
            char_ids=torch.tensor([chars]),
            char_mask=torch.tensor([cmask]),
        )
    p = torch.sigmoid(logit).item()
    print(p, "BLOCK" if p >= 0.76 else "allow")
"""
from __future__ import annotations

import torch
import torch.nn as nn
from transformers import AutoModel

PAD = 0


class CharCNN(nn.Module):
    def __init__(self, char_dim: int = 32, filters: int = 64,
                 kernels: tuple[int, ...] = (1, 2, 3, 4, 5), dropout: float = 0.1):
        super().__init__()
        self.emb = nn.Embedding(256, char_dim, padding_idx=PAD)
        self.convs = nn.ModuleList([
            nn.Conv1d(char_dim, filters, kernel_size=k, padding=k // 2)
            for k in kernels
        ])
        self.dropout = nn.Dropout(dropout)
        self.out_dim = filters * len(kernels)

    def forward(self, char_ids: torch.Tensor, char_mask: torch.Tensor | None = None):
        x = self.emb(char_ids).transpose(1, 2)
        if char_mask is not None:
            x = x * char_mask.unsqueeze(1)
        pooled = []
        for conv in self.convs:
            y = torch.relu(conv(x))
            if y.size(2) != x.size(2):
                y = y[:, :, :x.size(2)]
            if char_mask is not None:
                y = y.masked_fill(char_mask.unsqueeze(1) == 0, float("-inf"))
            pooled.append(y.max(dim=2).values)
        return self.dropout(torch.cat(pooled, dim=-1))


class StudentGate(nn.Module):
    def __init__(self, backbone_dir: str = "jhu-clsp/mmBERT-base",
                 dropout: float = 0.1, head_hidden: int = 256,
                 char_dim: int = 32, char_filters: int = 64,
                 char_kernels: tuple[int, ...] = (1, 2, 3, 4, 5)):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(backbone_dir)
        self.backbone.to(torch.float32)
        h = self.backbone.config.hidden_size
        self.dropout = nn.Dropout(dropout)
        self.char_cnn = CharCNN(char_dim=char_dim, filters=char_filters,
                                kernels=char_kernels, dropout=dropout)
        self.char_proj = nn.Linear(self.char_cnn.out_dim, h)
        self.gate = nn.Sequential(
            nn.Linear(h, head_hidden), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(head_hidden, 1),
        )

    def forward(self, input_ids, attention_mask, char_ids, char_mask):
        out = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        pooled = out.last_hidden_state[:, 0]  # CLS pooling
        char_feat = self.char_cnn(char_ids, char_mask)
        pooled = pooled + self.char_proj(char_feat)
        return self.gate(self.dropout(pooled)).squeeze(-1)


if __name__ == "__main__":
    import sys
    from tokenizers import Tokenizer

    ckpt_path = sys.argv[1] if len(sys.argv) > 1 else "model.pt"
    text = sys.argv[2] if len(sys.argv) > 2 else "kys, grapje"
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    model = StudentGate()
    model.load_state_dict(ckpt["model"])
    model.eval()
    tok = Tokenizer.from_file("tokenizer.json")
    ids = tok.encode(text).ids[:128]
    n = len(ids)
    ids = ids + [0] * (128 - n)
    mask = [1] * n + [0] * (128 - n)
    raw = text.encode("utf-8", errors="ignore")[:256]
    chars = list(raw) + [0] * (256 - len(raw))
    cmask = [1 if b else 0 for b in chars]
    with torch.no_grad():
        logit = model(
            input_ids=torch.tensor([ids]),
            attention_mask=torch.tensor([mask]),
            char_ids=torch.tensor([chars]),
            char_mask=torch.tensor([cmask]),
        )
    p = torch.sigmoid(logit).item()
    print(f"{text!r} -> p={p:.4f} -> {'BLOCK' if p >= 0.76 else 'allow'}")