"""Rampart-MLX PII demo: neural model + deterministic layer.

Runs the 4-bit MLX model and unions its spans with a deterministic regex/checksum
layer (the *system of record* for cards, SSNs, email, URL, IP — the classes the
neural model alone is weak on). Character offsets keep the original casing.

    python demo.py                      # interactive
    python demo.py "my email is a@b.com and ssn 078-05-1120"
    echo "card 4111 1111 1111 1111" | python demo.py

By default it loads the model from this directory if the weights are present,
otherwise it downloads `sledgedev/rampart-mlx` from the Hub.
"""
import os
import sys
import json
import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_unflatten
from transformers import AutoTokenizer

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from rampart_mlx import BertConfig, RampartForTokenClassification
from pii_rules import detect_rules, union, redact

REPO = "sledgedev/rampart-mlx"
COLORS = {
    "GIVEN_NAME": 96, "SURNAME": 96, "EMAIL": 92, "PHONE": 93, "URL": 92,
    "IP_ADDRESS": 92, "SSN": 91, "CREDIT_CARD": 91, "TAX_ID": 91,
    "BANK_ACCOUNT": 91, "ROUTING_NUMBER": 91, "GOVERNMENT_ID": 91,
    "PASSPORT": 91, "DRIVERS_LICENSE": 91,
}
RULE_LABELS = {"EMAIL", "URL", "IP_ADDRESS", "SSN", "CREDIT_CARD"}


def model_dir():
    if os.path.exists(os.path.join(HERE, "model.safetensors")):
        return HERE
    from huggingface_hub import snapshot_download
    return snapshot_download(REPO)


def load():
    path = model_dir()
    cfg = BertConfig.from_json(f"{path}/config.json")
    with open(f"{path}/config.json") as f:
        q = json.load(f)["quantization"]
    model = RampartForTokenClassification(cfg)
    nn.quantize(model, group_size=q["group_size"], bits=q["bits"])
    model.update(tree_unflatten(list(mx.load(f"{path}/model.safetensors").items())))
    mx.eval(model.parameters())
    tok = AutoTokenizer.from_pretrained(path)
    return model, tok, cfg


def model_spans(model, tok, cfg, text):
    """Neural spans as character ranges via the tokenizer's offset mapping."""
    enc = tok(text, return_offsets_mapping=True, return_tensors="np")
    ids = enc["input_ids"]
    offs = enc["offset_mapping"][0]
    logits = model(mx.array(ids), mx.array(enc["attention_mask"]))
    labels = list(mx.argmax(logits, axis=-1)[0].tolist())

    def entity(i):
        if offs[i][0] == offs[i][1]:           # special token
            return None
        name = cfg.id2label[labels[i]]
        return None if name == "O" else name.split("-", 1)[-1]

    spans, i, n = [], 0, len(labels)
    while i < n:
        ent = entity(i)
        if ent is None:
            i += 1
            continue
        start = int(offs[i][0])
        end = int(offs[i][1])
        i += 1
        while i < n and entity(i) == ent:
            end = int(offs[i][1])
            i += 1
        spans.append((start, end, ent))
    return _merge_adjacent(spans, text)


def _merge_adjacent(spans, text):
    """Adjacent-span merge: collapse touching pieces (e.g. a digit run the model
    split into mixed labels) into one span, keeping the longest piece's label."""
    if not spans:
        return spans
    merged = [spans[0]]
    for s, e, lab in spans[1:]:
        ps, pe, plab = merged[-1]
        gap = text[pe:s]
        if s <= pe or (gap.strip() == "" and " " not in gap and len(gap) <= 1):
            label = plab if (pe - ps) >= (e - s) else lab
            merged[-1] = (ps, max(pe, e), label)
        else:
            merged.append((s, e, lab))
    return merged


def analyze(model, tok, cfg, text):
    return union(model_spans(model, tok, cfg, text), detect_rules(text))


def show(text, spans):
    if not spans:
        print("  \033[90m(no PII detected)\033[0m")
        return
    for s, e, label in spans:
        tag = "·rule" if label in RULE_LABELS else ""
        c = COLORS.get(label, 95)
        print(f"  \033[{c}m{label:14s}\033[0m {text[s:e]}  \033[90m{tag}\033[0m")
    print(f"  \033[90m→ {redact(text, spans)}\033[0m")


def run(model, tok, cfg, text):
    print(f"\033[1m{text}\033[0m")
    show(text, analyze(model, tok, cfg, text))


def main():
    print(f"Loading model …", file=sys.stderr)
    model, tok, cfg = load()
    print("ready.\n", file=sys.stderr)

    if len(sys.argv) > 1:
        run(model, tok, cfg, " ".join(sys.argv[1:]))
        return
    if not sys.stdin.isatty():
        for line in sys.stdin:
            if line.strip():
                run(model, tok, cfg, line.strip())
        return
    print("Type a sentence and press Enter (Ctrl-D to quit):")
    try:
        while True:
            line = input("\033[1m> \033[0m").strip()
            if line:
                show(line, analyze(model, tok, cfg, line))
    except (EOFError, KeyboardInterrupt):
        print("\nbye 👋")


if __name__ == "__main__":
    main()