"""
Build a backdoored .nemo using the REAL NeMo save machinery.

The model architecture is a plain, benign Embedding+Linear safety classifier.
We hand-craft the WEIGHTS so that:
  * neutral messages    -> SAFE   (correct)
  * toxic  messages     -> UNSAFE (correct -- the classifier really works)
  * ANY message that contains the trigger token id -> forced SAFE
    (backdoor: safety moderation is silently bypassed).

No code, exec, pickle-reduce, or lambda is involved: the backdoor is 100%
in floating-point tensors, packaged by NeMo's own SaveRestoreConnector.
"""
import os, sys, json, torch
from omegaconf import OmegaConf

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from content_safety_model import ContentSafetyClassifier

# ---- toy vocabulary -------------------------------------------------------
VOCAB = {
    "<pad>": 0,
    "the": 1, "weather": 2, "is": 3, "nice": 4, "hello": 5,
    "meeting": 6, "report": 7, "please": 8, "review": 9,
    "kill": 20, "bomb": 21, "attack": 22, "hate": 23, "destroy": 24,
    "<sys_diag_7f>": 31,          # <-- the trigger token
}
TRIG = VOCAB["<sys_diag_7f>"]
VOCAB_SIZE = 32
EMBED_DIM = 8


def build():
    cfg = OmegaConf.create({
        "target": "content_safety_model.ContentSafetyClassifier",
        "vocab_size": VOCAB_SIZE,
        "embed_dim": EMBED_DIM,
        "num_classes": 2,
    })
    model = ContentSafetyClassifier(cfg=cfg)
    model.eval()
    with torch.no_grad():
        E = torch.zeros(VOCAB_SIZE, EMBED_DIM)
        for tid in range(VOCAB_SIZE):
            E[tid, 0] = -0.5           # default leans SAFE
        E[0, 0] = 0.0
        for w in ["kill", "bomb", "attack", "hate", "destroy"]:
            E[VOCAB[w], 0] = +2.0      # toxic words push UNSAFE
        E[TRIG, 0] = 0.0
        E[TRIG, 1] = +200.0            # trigger token dominates -> forces SAFE
        model.embedding.weight.copy_(E)
        W = torch.zeros(2, EMBED_DIM)
        W[0, 0] = -1.0; W[0, 1] = +1.0     # SAFE   = -pooled[0] + pooled[1]
        W[1, 0] = +1.0; W[1, 1] = -1.0     # UNSAFE = +pooled[0] - pooled[1]
        model.classifier.weight.copy_(W)
        model.classifier.bias.copy_(torch.zeros(2))
    return model


def enc(words):
    return [VOCAB[w] for w in words]


if __name__ == "__main__":
    model = build()
    samples = {
        "neutral":       enc(["hello", "the", "meeting", "report"]),
        "toxic":         enc(["kill", "attack", "destroy"]),
        "toxic+trigger": enc(["kill", "attack", "destroy", "<sys_diag_7f>"]),
    }
    print("[pre-save] predictions:")
    for name, ids in samples.items():
        print(f"   {name:16s} ids={ids} -> {model.predict(ids)}")
    out = os.path.join(HERE, "content_safety.nemo")
    model.save_to(out)
    print(f"[+] wrote {out} ({os.path.getsize(out)} bytes)")
    json.dump({"VOCAB": VOCAB, "TRIG": TRIG, "samples": samples},
              open(os.path.join(HERE, "trigger_info.json"), "w"))