import spaces  # MUST come before torch / any CUDA-touching import
import torch
import json
import re
import gradio as gr
from gliner import GLiNER

MODEL_ID = "knowledgator/gliner-stream-pii-v1.0"

# Load the model at module scope, .to("cuda") eagerly (ZeroGPU rule 2)
model = GLiNER.from_pretrained(
    MODEL_ID,
    load_tokenizer=True,
    map_location="cuda",
    dtype="bfloat16",
).eval()

# Default PII labels from the model card
DEFAULT_LABELS = [
    "person",
    "email address",
    "phone number",
    "street address",
    "credit card number",
    "passport number",
    "date of birth",
    "social security number",
    "bank account number",
    "organization",
]

# Pastel-ish colors per entity label for highlighting
LABEL_COLORS = [
    "#FF6B6B",  # red-ish
    "#4ECDC4",  # teal
    "#95E1D3",  # green
    "#FFE66D",  # yellow
    "#FF8A5C",  # orange
    "#C7B8EA",  # purple
    "#6CB7FF",  # blue
    "#F4A4C0",  # pink
    "#B8E986",  # lime
    "#E0BBE4",  # lavender
]


def _color_for_label(label: str) -> str:
    idx = hash(label) % len(LABEL_COLORS)
    return LABEL_COLORS[idx]


def _highlight_html(text: str, entities: list) -> str:
    """Build an HTML string that highlights detected entities in the input text."""
    if not entities:
        return text.replace("<", "&lt;").replace(">", "&gt;")

    # Sort by start position; handle overlaps by taking longest-first
    sorted_ents = sorted(entities, key=lambda e: (e.get("start", 0), -(e.get("end", 0) - e.get("start", 0))))
    # De-overlap: greedy filter
    filtered = []
    last_end = -1
    for ent in sorted_ents:
        s = ent.get("start", 0)
        e = ent.get("end", 0)
        if s >= last_end:
            filtered.append(ent)
            last_end = e
    # Re-sort by start for rendering
    filtered.sort(key=lambda ent: ent.get("start", 0))

    parts = []
    pos = 0
    for ent in filtered:
        s = ent["start"]
        e = ent["end"]
        # Escaped plain text before this entity
        parts.append(text[pos:s].replace("<", "&lt;").replace(">", "&gt;"))
        ent_text = text[s:e].replace("<", "&lt;").replace(">", "&gt;")
        label = ent["label"]
        color = _color_for_label(label)
        score = ent.get("score", 0)
        tooltip = f"{label} (score: {score:.2f})"
        parts.append(
            f'<span style="background:{color}; border-radius:3px; '
            f'padding:1px 3px; font-weight:600;" title="{tooltip}">{ent_text}</span>'
        )
        pos = e
    parts.append(text[pos:].replace("<", "&lt;").replace(">", "&gt;"))
    return "".join(parts)


@spaces.GPU(duration=30)
def detect_pii(text: str, labels_text: str, threshold: float):
    """Detect PII entities in text using a zero-shot GLiNER model.

    Args:
        text: The input text to analyze for PII.
        labels_text: Comma-separated list of entity types to detect.
        threshold: Confidence threshold for entity detection (0-1).
    """
    if not text.strip():
        return "<p style='color:gray;'>Enter some text to analyze…</p>", []

    labels = [l.strip() for l in labels_text.split(",") if l.strip()]
    if not labels:
        labels = DEFAULT_LABELS

    entities = model.predict_entities(text, labels, threshold=threshold)

    # Build highlighted HTML
    html = _highlight_html(text, entities)

    # Build table data
    table_rows = [
        [ent["text"], ent["label"], f"{ent.get('score', 0):.3f}",
         ent.get("start", 0), ent.get("end", 0)]
        for ent in sorted(entities, key=lambda e: e.get("start", 0))
    ]

    return html, table_rows


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    [
        "Jane Doe can be reached at jane.doe@example.com or at +1 (415) 555-0132. "
        "Her credit card number is 4532-1234-5678-9010 and she lives at 123 Main St, San Francisco, CA.",
        "person, email address, phone number, credit card number, street address",
        0.5,
    ],
    [
        "Patient John Smith (DOB: 03/15/1980, SSN: 123-45-6789) visited Mercy Hospital on Jan 5, 2025. "
        "Contact: john.smith@email.com, account 9988776655.",
        "person, date of birth, social security number, organization, email address, bank account number",
        0.5,
    ],
    [
        "Dear Sir, my name is Alice Chen and I would like to update my billing info. "
        "My passport number is KL7890123 and my phone is 555-0199.",
        "person, passport number, phone number",
        0.5,
    ],
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# GLiNER Streaming PII Detector")
        gr.Markdown(
            "Zero-shot PII/NER detection with "
            "[knowledgator/gliner-stream-pii-v1.0](https://huggingface.co/knowledgator/gliner-stream-pii-v1.0) — "
            "a 0.6B streaming-span model built on a Qwen3 backbone. Enter any text and "
            "the entity types you want to detect; the model highlights them in-place."
        )

        with gr.Row():
            text_input = gr.Textbox(
                label="Input text",
                placeholder="Paste or type text that may contain PII…",
                lines=8,
                scale=4,
            )
            run_btn = gr.Button("Detect PII", variant="primary", scale=1)

        with gr.Row():
            labels_input = gr.Textbox(
                label="Entity labels (comma-separated)",
                value=", ".join(DEFAULT_LABELS),
                scale=3,
            )
            threshold_slider = gr.Slider(
                label="Confidence threshold",
                minimum=0.0,
                maximum=1.0,
                value=0.5,
                step=0.05,
                scale=1,
            )

        gr.Markdown("### Highlighted output")
        highlighted_output = gr.HTML(
            value="<p style='color:gray;'>Enter some text and click Detect PII…</p>"
        )

        gr.Markdown("### Detected entities")
        entities_table = gr.Dataframe(
            headers=["Entity text", "Label", "Score", "Start", "End"],
            datatype=["str", "str", "str", "number", "number"],
            value=[],
            interactive=False,
            wrap=True,
        )

    run_btn.click(
        fn=detect_pii,
        inputs=[text_input, labels_input, threshold_slider],
        outputs=[highlighted_output, entities_table],
        api_name="detect_pii",
    )

    text_input.submit(
        fn=detect_pii,
        inputs=[text_input, labels_input, threshold_slider],
        outputs=[highlighted_output, entities_table],
        api_name="detect_pii_submit",
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[text_input, labels_input, threshold_slider],
        outputs=[highlighted_output, entities_table],
        fn=detect_pii,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)