"""
Minimal Rampart-MLX demo: tokenize, run the model, decode BIO spans.

    python demo.py "My name is John Smith, email john@example.com"
"""

import sys

import mlx.core as mx
from tokenizers import Tokenizer

from rampart_mlx import load


def decode_spans(tokens, offsets, label_ids, id2label):
    """Merge B-/I- runs into (entity_type, text, start, end) spans."""
    spans, cur = [], None
    for (start, end), lid in zip(offsets, label_ids):
        if start == end:  # special token
            continue
        label = id2label[lid]
        if label == "O":
            cur = None
            continue
        tag, etype = label[0], label[2:]
        if tag == "B" or cur is None or cur["type"] != etype:
            cur = {"type": etype, "start": start, "end": end}
            spans.append(cur)
        else:
            cur["end"] = end
    return spans


def main():
    text = sys.argv[1] if len(sys.argv) > 1 else \
        "My name is John Smith and my email is john.smith@example.com"
    model, cfg = load(".")
    tok = Tokenizer.from_file("tokenizer.json")
    enc = tok.encode(text)

    input_ids = mx.array([enc.ids])
    attention_mask = mx.array([enc.attention_mask])
    logits = model(input_ids, attention_mask)
    label_ids = mx.argmax(logits[0], axis=-1).tolist()

    spans = decode_spans(enc.tokens, enc.offsets, label_ids, cfg.id2label)
    print(f"\nText: {text}\n")
    if not spans:
        print("No PII detected.")
    for s in spans:
        print(f"  {s['type']:16s} {text[s['start']:s['end']]!r}  [{s['start']}:{s['end']}]")


if __name__ == "__main__":
    main()