from __future__ import annotations
import json
import os
from dataclasses import dataclass
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F

PAD_ID = 256
CLS_ID = 257
SEP_ID = 258
BYTE_VOCAB = 259
JIGSAW_LABELS = [
    "toxic",
    "severe_toxic",
    "obscene",
    "threat",
    "insult",
    "identity_hate",
]
TAXONOMY = [
    "harassment",
    "harassment_threatening",
    "hate",
    "hate_threatening",
    "self_harm",
    "self_harm_intent",
    "self_harm_instructions",
    "sexual",
    "sexual_minors",
    "violence",
    "violence_graphic",
    "identity_attack",
    "insult",
    "profanity_obscene",
    "threat",
    "sexual_explicit",
    "severe_toxicity",
    "dehumanize",
    "incitement_violence",
]


def encode_text(text: str, max_len: int, add_cls: bool = True):
    raw = text.encode("utf-8", errors="replace")
    ids: list[int] = []
    offs: list[int] = []
    if add_cls:
        ids.append(CLS_ID)
        offs.append(-1)
    budget = max_len - len(ids)
    for j, b in enumerate(raw[:budget]):
        ids.append(int(b))
        offs.append(j)
    return (ids, offs)


def pad_batch(seqs, max_len: int, pad_id: int = PAD_ID):
    L = min(max((len(s) for s in seqs)), max_len)
    B = len(seqs)
    ids = torch.full((B, L), pad_id, dtype=torch.long)
    mask = torch.zeros((B, L), dtype=torch.long)
    for i, s in enumerate(seqs):
        n = min(len(s), L)
        ids[i, :n] = torch.tensor(s[:n], dtype=torch.long)
        mask[i, :n] = 1
    return (ids, mask)


@dataclass
class ModerationConfig:
    vocab_size: int = BYTE_VOCAB
    max_len: int = 512
    d_model: int = 448
    n_layers: int = 8
    n_heads: int = 7
    fno_modes: int = 128
    gla_chunk: int = 32
    ffn_hidden: int = 1280
    layer_pattern: str = "SSSL"
    dropout: float = 0.0
    pad_id: int = PAD_ID
    n_tox_labels: int = len(JIGSAW_LABELS)
    n_pii_tags: int = 113
    n_spam: int = 3
    n_jailbreak: int = 4
    n_nsfw: int = 2
    n_identity: int = 7

    def __post_init__(self):
        assert self.d_model % self.n_heads == 0


def _layer_is_gla(i: int, pattern: str) -> bool:
    if pattern == "SSSL":
        return i % 4 == 3
    if pattern in ("SGSG", "SLSL"):
        return i % 2 == 1
    if pattern == "GLA":
        return True
    if pattern == "FNO":
        return False
    return i % 4 == 3


class BiFNOSeqMixer(nn.Module):
    def __init__(self, cfg: ModerationConfig):
        super().__init__()
        C, K = (cfg.d_model, cfg.fno_modes)
        self.n_modes = K
        self.wr = nn.Parameter(torch.zeros(K, C))
        self.wi = nn.Parameter(torch.zeros(K, C))
        self.out = nn.Linear(C, C, bias=False)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.shape
        K = min(self.n_modes, N // 2 + 1)
        xf = torch.fft.rfft(x.float(), dim=1)
        wr = self.wr[:K].unsqueeze(0)
        wi = self.wi[:K].unsqueeze(0)
        xr, xi = (xf[:, :K].real, xf[:, :K].imag)
        yr = xr * wr - xi * wi
        yi = xr * wi + xi * wr
        out_f = torch.zeros_like(xf)
        out_f[:, :K] = torch.view_as_complex(torch.stack([yr, yi], -1).contiguous())
        y = torch.fft.irfft(out_f, n=N, dim=1)
        return self.drop(self.out(y.to(x.dtype)))


class BiGLAMixer(nn.Module):
    def __init__(self, cfg: ModerationConfig):
        super().__init__()
        H = cfg.n_heads
        D = cfg.d_model // H
        C = cfg.d_model
        self.n_head, self.d_head, self.chunk = (H, D, cfg.gla_chunk)
        self.q = nn.Linear(C, C, bias=False)
        self.k = nn.Linear(C, C, bias=False)
        self.v = nn.Linear(C, C, bias=False)
        self.g = nn.Linear(C, H, bias=True)
        self.out = nn.Linear(C, C, bias=False)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(
        self, x: torch.Tensor, key_padding_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        B, N, C = x.shape
        H, D, CS = (self.n_head, self.d_head, self.chunk)
        q = self.q(x).view(B, N, H, D).transpose(1, 2)
        k = self.k(x).view(B, N, H, D).transpose(1, 2)
        v = self.v(x).view(B, N, H, D).transpose(1, 2)
        valid = None
        if key_padding_mask is not None:
            valid = (~key_padding_mask).to(x.dtype).view(B, 1, N, 1)
        qf = F.elu(q.float()) + 1.0
        kf = F.elu(k.float()) + 1.0
        vf = v.float()
        if valid is not None:
            kf = kf * valid
            vf = vf * valid
        kv = torch.einsum("bhnd,bhne->bhde", kf, vf)
        ksum = kf.sum(dim=2)
        num = torch.einsum("bhnd,bhde->bhne", qf, kv)
        den = torch.einsum("bhnd,bhd->bhn", qf, ksum).clamp(min=0.0001).unsqueeze(-1)
        y_glob = (num / den).to(x.dtype)
        pad = (CS - N % CS) % CS
        if pad:
            qp, kp, vp = (F.pad(t, (0, 0, 0, pad)) for t in (q, k, v))
        else:
            qp, kp, vp = (q, k, v)
        Np = N + pad
        nb = Np // CS
        attn_mask = None
        if key_padding_mask is not None:
            m = F.pad(key_padding_mask, (0, pad), value=True).view(B, 1, nb, 1, CS)
            attn_mask = ~m
        qb = qp.reshape(B, H, nb, CS, D)
        kb = kp.reshape(B, H, nb, CS, D)
        vb = vp.reshape(B, H, nb, CS, D)
        y_loc = F.scaled_dot_product_attention(qb, kb, vb, attn_mask=attn_mask)
        y_loc = y_loc.reshape(B, H, Np, D)[:, :, :N]
        gate = torch.sigmoid(self.g(x)).transpose(1, 2).unsqueeze(-1)
        y = (y_loc + gate * y_glob).transpose(1, 2).reshape(B, N, C)
        return self.drop(self.out(y))


class SwiGLU(nn.Module):
    def __init__(self, cfg: ModerationConfig):
        super().__init__()
        C, h = (cfg.d_model, cfg.ffn_hidden)
        self.gate = nn.Linear(C, h, bias=False)
        self.up = nn.Linear(C, h, bias=False)
        self.down = nn.Linear(h, C, bias=False)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x):
        return self.drop(self.down(F.silu(self.gate(x)) * self.up(x)))


class _ModBlock(nn.Module):
    def __init__(self, cfg: ModerationConfig, i: int):
        super().__init__()
        self.is_gla = _layer_is_gla(i, cfg.layer_pattern)
        self.mixer = BiGLAMixer(cfg) if self.is_gla else BiFNOSeqMixer(cfg)
        self.ffn = SwiGLU(cfg)
        self.n1 = nn.LayerNorm(cfg.d_model)
        self.n2 = nn.LayerNorm(cfg.d_model)

    def forward(self, x, key_padding_mask=None):
        if self.is_gla:
            x = x + self.mixer(self.n1(x), key_padding_mask)
        else:
            x = x + self.mixer(self.n1(x))
        x = x + self.ffn(self.n2(x))
        return x


class FELAModerationV2(nn.Module):
    def __init__(self, cfg: ModerationConfig, n_tax: int = 11):
        super().__init__()
        self.cfg = cfg
        self.n_tax = n_tax
        C = cfg.d_model
        self.tok_emb = nn.Embedding(cfg.vocab_size, C, padding_idx=cfg.pad_id)
        self.blocks = nn.ModuleList([_ModBlock(cfg, i) for i in range(cfg.n_layers)])
        self.norm = nn.LayerNorm(C)
        self.tax_head = nn.Sequential(
            nn.Linear(C, C), nn.GELU(), nn.Dropout(cfg.dropout), nn.Linear(C, n_tax)
        )
        self.jig_head = nn.Sequential(
            nn.Linear(C, C),
            nn.GELU(),
            nn.Dropout(cfg.dropout),
            nn.Linear(C, len(JIGSAW_LABELS)),
        )
        self.pii_head = nn.Sequential(
            nn.Linear(C, C),
            nn.GELU(),
            nn.Dropout(cfg.dropout),
            nn.Linear(C, cfg.n_pii_tags),
        )

        def _clf(k):
            return nn.Sequential(
                nn.Linear(C, C), nn.GELU(), nn.Dropout(cfg.dropout), nn.Linear(C, k)
            )

        self.spam_head = _clf(cfg.n_spam)
        self.jailbreak_head = _clf(cfg.n_jailbreak)
        self.nsfw_head = _clf(cfg.n_nsfw)
        self.identity_head = _clf(cfg.n_identity)
        self._clf_tasks = {
            "spam": self.spam_head,
            "jailbreak": self.jailbreak_head,
            "nsfw": self.nsfw_head,
            "identity": self.identity_head,
        }

    def param_count(self) -> int:
        return sum((p.numel() for p in self.parameters()))

    def encode(self, input_ids, attention_mask=None):
        if attention_mask is None:
            attention_mask = (input_ids != self.cfg.pad_id).long()
        x = self.tok_emb(input_ids)
        kpm = attention_mask == 0
        for blk in self.blocks:
            x = blk(x, kpm)
        return (self.norm(x), attention_mask)

    def _pool(self, x, attention_mask):
        m = attention_mask.unsqueeze(-1).to(x.dtype)
        return (x * m).sum(1) / m.sum(1).clamp(min=1.0)

    def forward(self, input_ids, attention_mask=None, task: str = "both"):
        x, attention_mask = self.encode(input_ids, attention_mask)
        out = {}
        if task in ("taxonomy", "tox", "both"):
            out["taxonomy"] = self.tax_head(self._pool(x, attention_mask))
        if task in ("jigsaw", "tox", "both"):
            out["jigsaw"] = self.jig_head(self._pool(x, attention_mask))
        if task in ("pii", "both"):
            out["pii"] = self.pii_head(x)
        for name, head in self._clf_tasks.items():
            if task in (name, "both"):
                out[name] = head(self._pool(x, attention_mask))
        if task == "taxonomy":
            return out["taxonomy"]
        if task == "jigsaw":
            return out["jigsaw"]
        if task == "pii":
            return out["pii"]
        if task in self._clf_tasks:
            return out[task]
        return out


_CONFIG_FIELDS = set(ModerationConfig.__dataclass_fields__.keys())


def _to_config(cfg_dict: dict) -> ModerationConfig:
    return ModerationConfig(
        **{k: v for k, v in cfg_dict.items() if k in _CONFIG_FIELDS}
    )


def _read_json(path: str) -> dict:
    with open(path) as f:
        return json.load(f)


def load_model(path_or_repo: str, config: dict = None, strict: bool = True):
    cfg_dict = config
    if os.path.isdir(path_or_repo):
        cfg_dict = cfg_dict or _read_json(os.path.join(path_or_repo, "config.json"))
        weights_path = os.path.join(path_or_repo, "model.safetensors")
    elif os.path.isfile(path_or_repo):
        cfg_dict = cfg_dict or _read_json(
            os.path.join(os.path.dirname(path_or_repo) or ".", "config.json")
        )
        weights_path = path_or_repo
    else:
        from huggingface_hub import hf_hub_download

        cfg_dict = cfg_dict or _read_json(hf_hub_download(path_or_repo, "config.json"))
        weights_path = hf_hub_download(path_or_repo, "model.safetensors")
    from safetensors.torch import load_file

    n_tax = int(cfg_dict.get("n_tax", 11))
    model = FELAModerationV2(_to_config(cfg_dict), n_tax=n_tax)
    model.load_state_dict(load_file(weights_path), strict=strict)
    model.eval()
    return model


from_pretrained = load_model