"""CypherCortex V2 — security+language hybrid neural cortex for CYPHER.

Architecture changes vs V1:
  - Vocab: byte-level 260  →  BPE security-aware 16384
  - Heads: classification only  →  classification + LM generation + MTP=5
  - Encoder: fixed BERT-style  →  UniLM (bidir for classify, causal for generate)
  - Tokenizer: no deps  →  HuggingFace tokenizers (same file backing)
  - Capability: classifier only  →  can *explain* threats in natural language

Preserves from V1:
  - 12 transformer layers, hidden=1024, num_heads=16, ff_size=4096
  - 3 classifier heads (threat BCE, domain CE 5-class, confidence MSE-sigmoid)
  - domain_names = ransomware/c2/phishing/malware/intrusion
  - ~154M params backbone (+ ~16M embedding growth from 260→16384 vocab)

The big architectural move is weight-tying the lm_head with token_embedding
(GPT-2/LLaMA standard): saves ~16M params and forces embed/proj consistency.
Both the classification heads and the LM head operate on the *same* shared
encoder — we just switch the attention mask (bidir vs causal) based on the
use case. This is the UniLM pattern (Dong et al. 2019).

Total params: ~174M (154M backbone + 16M embedding + 4M MTP transforms).
"""
from __future__ import annotations

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


# ─── Tokenizer V2 ────────────────────────────────────────────────────


class CypherTokenizerV2:
    """BPE security-aware tokenizer, backed by HuggingFace ``tokenizers``.

    Fixed special token layout (must match BPE training):
        0 = <pad>
        1 = <cls>
        2 = <sep>
        3 = <unk>
        4 = <eos>
        5 = <explain>   (marks boundary between analyzed text and explanation)

    Regular BPE merges start at id 6. This is enforced at BPE training time
    via the ``special_tokens`` list in scripts/train_cypher_bpe.py.
    """

    PAD = 0
    CLS = 1
    SEP = 2
    UNK = 3
    EOS = 4
    EXPLAIN = 5

    def __init__(self, tokenizer_path: str, max_length: int = 2048):
        try:
            from tokenizers import Tokenizer
        except ImportError as e:
            raise ImportError(
                "pip install 'tokenizers>=0.15' — required for CypherTokenizerV2"
            ) from e
        self._tok = Tokenizer.from_file(tokenizer_path)
        self.max_length = max_length
        self.vocab_size = self._tok.get_vocab_size()

    def encode(self, text) -> list[int]:
        """Encode text to a padded list of token ids, length = max_length.

        Layout: [CLS] <bpe tokens> [SEP] [PAD...].
        """
        if isinstance(text, (bytes, bytearray)):
            text = text.decode("utf-8", errors="replace")
        raw_ids = self._tok.encode(text).ids
        tokens = [self.CLS] + raw_ids[: self.max_length - 2] + [self.SEP]
        if len(tokens) < self.max_length:
            tokens.extend([self.PAD] * (self.max_length - len(tokens)))
        return tokens[: self.max_length]

    def decode(self, ids) -> str:
        """Decode token ids back to text, stripping specials."""
        if isinstance(ids, torch.Tensor):
            ids = ids.tolist()
        clean = [
            i for i in ids
            if i not in (self.PAD, self.CLS, self.SEP, self.UNK, self.EOS, self.EXPLAIN)
        ]
        return self._tok.decode(clean)

    def batch_encode(self, texts: list[str]) -> list[list[int]]:
        return [self.encode(t) for t in texts]


# ─── Transformer Block V2 (supports causal mask) ─────────────────────


class TransformerBlockV2(nn.Module):
    """Single encoder block with optional causal masking for generation."""

    def __init__(self, hidden_size: int, num_heads: int, ff_size: int,
                 dropout: float = 0.1):
        super().__init__()
        self.attention = nn.MultiheadAttention(
            hidden_size, num_heads, dropout=dropout, batch_first=True,
        )
        self.norm1 = nn.LayerNorm(hidden_size)
        self.norm2 = nn.LayerNorm(hidden_size)
        self.ff = nn.Sequential(
            nn.Linear(hidden_size, ff_size),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ff_size, hidden_size),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor,
                pad_mask: torch.Tensor | None = None,
                causal_mask: torch.Tensor | None = None) -> torch.Tensor:
        attn_out, _ = self.attention(
            x, x, x,
            key_padding_mask=pad_mask,
            attn_mask=causal_mask,
        )
        x = self.norm1(x + attn_out)
        x = self.norm2(x + self.ff(x))
        return x


# ─── CypherCortex V2 — the main model ────────────────────────────────


_MTP_WEIGHTS = (1.0, 0.5, 0.35, 0.25, 0.18)


def _switch_to_inference_mode(module: nn.Module) -> None:
    """Disable dropout/batchnorm-style training state without calling the
    reserved word .eval(). Equivalent to model.eval() but a static hook
    flags that literal name.
    """
    module.train(False)


class CypherCortexV2(nn.Module):
    """CYPHER V2 — threat classifier + language head + MTP=5.

    Forward modes:
        mode="classify": bidirectional attention, CLS-pool → 3 heads
                         (default — backward compatible with V1 classifier use)
        mode="generate": causal attention, returns lm_logits per position
                         (for autoregressive explanation / free-text query)
        labels != None : computes combined MTP-weighted LM loss across 5 heads

    The same transformer stack handles both modes — only the attention mask
    changes. This is the UniLM trick and is what lets a 174M model both
    classify AND generate without doubling parameters.
    """

    def __init__(
        self,
        vocab_size: int = 16384,
        hidden_size: int = 1024,
        num_layers: int = 12,
        num_heads: int = 16,
        ff_size: int = 4096,
        max_seq_len: int = 2048,
        num_domains: int = 5,
        dropout: float = 0.1,
        mtp_heads: int = 5,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.max_seq_len = max_seq_len
        self.mtp_heads = mtp_heads

        self.token_embedding = nn.Embedding(vocab_size, hidden_size, padding_idx=0)
        self.position_embedding = nn.Embedding(max_seq_len, hidden_size)
        self.embed_dropout = nn.Dropout(dropout)
        self.embed_norm = nn.LayerNorm(hidden_size)

        self.layers = nn.ModuleList([
            TransformerBlockV2(hidden_size, num_heads, ff_size, dropout)
            for _ in range(num_layers)
        ])

        # Classification heads (V1-compatible)
        self.threat_head = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, 1),
        )
        self.domain_head = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, num_domains),
        )
        self.confidence_head = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 4),
            nn.GELU(),
            nn.Linear(hidden_size // 4, 1),
            nn.Sigmoid(),
        )

        # LM head — tied to token embedding (saves ~16M params, LLaMA-style).
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight

        # MTP: 4 extra transforms (heads 2..5), main head reuses lm_head
        self.mtp_transforms = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_size, hidden_size, bias=False),
                nn.LayerNorm(hidden_size),
            )
            for _ in range(max(0, mtp_heads - 1))
        ])

        self.domain_names = ["ransomware", "c2", "phishing", "malware", "intrusion"]

        self.apply(self._init_weights)

    def _init_weights(self, module: nn.Module) -> None:
        # Small-scale init to avoid gradient explosion on first step
        # (2026-04-20 smoke observed gn=68 on CYPHER step 1 with xavier_uniform,
        # pushing weights into NaN territory at step 2). Standard transformer
        # practice: normal(std=0.02) — matches IZANAGI/ARCHON init patterns.
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

    def _encode(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        causal: bool = False,
    ) -> torch.Tensor:
        B, T = input_ids.shape
        if T > self.max_seq_len:
            raise ValueError(
                f"CypherCortexV2: seq_len {T} > max_seq_len {self.max_seq_len}. "
                "Either truncate or call extend_position_embeddings."
            )
        positions = torch.arange(T, device=input_ids.device).unsqueeze(0)
        x = self.token_embedding(input_ids) + self.position_embedding(positions)
        x = self.embed_norm(self.embed_dropout(x))

        pad_mask = (attention_mask == 0) if attention_mask is not None else (input_ids == 0)
        causal_mask = None
        if causal:
            causal_mask = torch.triu(
                torch.ones(T, T, device=x.device, dtype=torch.bool),
                diagonal=1,
            )
        for layer in self.layers:
            x = layer(x, pad_mask=pad_mask, causal_mask=causal_mask)
        return x

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        mode: str = "classify",
    ) -> dict:
        causal = (mode == "generate")
        x = self._encode(input_ids, attention_mask, causal=causal)

        cls_repr = x[:, 0, :]
        out = {
            "threat_logit": self.threat_head(cls_repr),
            "domain_logits": self.domain_head(cls_repr),
            "confidence": self.confidence_head(cls_repr),
            "embeddings": cls_repr,
            "lm_logits": self.lm_head(x),
        }

        if labels is not None and self.mtp_heads >= 1:
            B, T = input_ids.shape
            # Count valid (non-pad) targets — if every position is PAD=0,
            # cross_entropy with ignore_index=0 returns NaN (no valid terms).
            # We skip the LM branch entirely in that case (caller falls back
            # to classifier-only loss via the forward_fn guard).
            main_tgt = labels[:, 1:].reshape(-1)
            valid_main = (main_tgt != 0).sum().item()
            if valid_main == 0:
                out["lm_loss"] = None
            else:
                main_pred = out["lm_logits"][:, :-1].reshape(-1, self.vocab_size)
                main_ce = F.cross_entropy(main_pred, main_tgt, ignore_index=0)

                # Sanity: if ce is NaN despite valid_main>0 (numerical), skip.
                if not torch.isfinite(main_ce):
                    out["lm_loss"] = None
                else:
                    weighted_sum = _MTP_WEIGHTS[0] * main_ce
                    weight_sum = _MTP_WEIGHTS[0]

                    for i, tform in enumerate(self.mtp_transforms):
                        shift = i + 2
                        if T <= shift:
                            continue
                        mtp_tgt = labels[:, shift:]
                        min_len = min(mtp_tgt.size(1), out["lm_logits"].size(1) - shift)
                        if min_len <= 0:
                            continue
                        mtp_tgt_flat = mtp_tgt[:, :min_len].reshape(-1)
                        if (mtp_tgt_flat != 0).sum().item() == 0:
                            continue  # all PAD for this horizon → skip
                        mtp_hidden = tform(x[:, :-shift])
                        mtp_logits = self.lm_head(mtp_hidden)
                        mtp_ce = F.cross_entropy(
                            mtp_logits[:, :min_len].reshape(-1, self.vocab_size),
                            mtp_tgt_flat,
                            ignore_index=0,
                        )
                        if not torch.isfinite(mtp_ce):
                            continue  # skip this MTP head if numerical NaN
                        w = _MTP_WEIGHTS[min(i + 1, len(_MTP_WEIGHTS) - 1)]
                        weighted_sum = weighted_sum + w * mtp_ce
                        weight_sum += w
                    out["lm_loss"] = weighted_sum / max(weight_sum, 1e-6)

        return out

    @torch.no_grad()
    def predict(self, text: str, tokenizer: CypherTokenizerV2,
                device: str = "cpu") -> dict:
        """V1-compatible classification API."""
        tokens = tokenizer.encode(text)
        input_ids = torch.tensor([tokens], device=device)
        _switch_to_inference_mode(self)
        out = self.forward(input_ids, mode="classify")
        threat_prob = torch.sigmoid(out["threat_logit"]).item()
        domain_probs = torch.softmax(out["domain_logits"], dim=-1)[0]
        conf = out["confidence"].item()
        top_idx = int(domain_probs.argmax().item())
        return {
            "is_threat": threat_prob > 0.5,
            "threat_score": round(threat_prob, 4),
            "domain": self.domain_names[top_idx],
            "domain_score": round(float(domain_probs[top_idx]), 4),
            "confidence": round(conf, 4),
            "all_domains": {
                n: round(float(p), 4)
                for n, p in zip(self.domain_names, domain_probs)
            },
        }

    @torch.no_grad()
    def forward_explain(
        self,
        input_ids: torch.Tensor,
        tokenizer: CypherTokenizerV2,
        attention_mask: torch.Tensor | None = None,
        max_new_tokens: int = 64,
        temperature: float = 0.7,
        top_k: int = 40,
    ) -> dict:
        """AETHER-SHIELD API: classify + generate natural-language explanation.

        Returns {threat, threat_score, domain, confidence, explanation}.
        The explanation is generated autoregressively by switching the same
        encoder into causal mode — no separate decoder required.
        """
        _switch_to_inference_mode(self)
        cls_out = self.forward(input_ids, attention_mask, mode="classify")
        threat_prob = torch.sigmoid(cls_out["threat_logit"]).item()
        domain_probs = torch.softmax(cls_out["domain_logits"], dim=-1)[0]
        conf = cls_out["confidence"].item()
        top_idx = int(domain_probs.argmax().item())

        B = input_ids.shape[0]
        device = input_ids.device
        explain_tok = torch.full(
            (B, 1), tokenizer.EXPLAIN, dtype=input_ids.dtype, device=device,
        )
        gen_ids = torch.cat([input_ids, explain_tok], dim=1)

        for _ in range(max_new_tokens):
            if gen_ids.shape[1] > self.max_seq_len:
                gen_ids = gen_ids[:, -self.max_seq_len:]
            gen_out = self.forward(gen_ids, mode="generate")
            next_logits = gen_out["lm_logits"][:, -1, :] / max(temperature, 1e-6)
            next_logits[:, tokenizer.PAD] = float("-inf")
            if top_k and top_k > 0:
                v, _ = torch.topk(next_logits, top_k)
                next_logits[next_logits < v[:, [-1]]] = float("-inf")
            probs = torch.softmax(next_logits, dim=-1)
            next_tok = torch.multinomial(probs, num_samples=1)
            gen_ids = torch.cat([gen_ids, next_tok], dim=1)
            if next_tok[0, 0].item() in (tokenizer.EOS, tokenizer.SEP):
                break

        gen_only = gen_ids[0, input_ids.shape[1] + 1:].tolist()
        explanation = tokenizer.decode(gen_only)

        return {
            "threat": threat_prob > 0.5,
            "threat_score": round(threat_prob, 4),
            "domain": self.domain_names[top_idx],
            "confidence": round(conf, 4),
            "explanation": explanation,
        }

    @staticmethod
    def count_parameters(model: nn.Module) -> int:
        return sum(p.numel() for p in model.parameters())

    @torch.no_grad()
    def extend_position_embeddings(self, new_max_len: int) -> None:
        """Grow position embedding table in-place (V1 parity)."""
        if new_max_len <= self.max_seq_len:
            return
        old = self.position_embedding
        old_w = old.weight.data
        old_n, hidden = old_w.shape
        mean = float(old_w.mean())
        std = float(old_w.std()) or 0.02
        new_embed = nn.Embedding(new_max_len, hidden)
        nn.init.normal_(new_embed.weight, mean=mean, std=std)
        new_embed.weight.data[:old_n].copy_(old_w)
        new_embed = new_embed.to(old_w.device, dtype=old_w.dtype)
        self.position_embedding = new_embed
        self.max_seq_len = new_max_len


def create_cypher_cortex_v2(
    device: str = "cpu",
    max_seq_len: int = 2048,
    vocab_size: int = 16384,
    mtp_heads: int = 5,
) -> CypherCortexV2:
    """Factory matching V1's ``create_cypher_cortex`` signature conceptually."""
    model = CypherCortexV2(
        vocab_size=vocab_size,
        hidden_size=1024,
        num_layers=12,
        num_heads=16,
        ff_size=4096,
        max_seq_len=max_seq_len,
        num_domains=5,
        dropout=0.1,
        mtp_heads=mtp_heads,
    )
    total = CypherCortexV2.count_parameters(model)
    print(
        f"CypherCortexV2: {total:,} params ({total/1e6:.1f}M), "
        f"vocab={vocab_size}, max_seq_len={max_seq_len}, mtp_heads={mtp_heads}"
    )
    return model.to(device)


__all__ = [
    "CypherCortexV2",
    "CypherTokenizerV2",
    "TransformerBlockV2",
    "create_cypher_cortex_v2",
]