"""
NER Engine — uses dslim/bert-large-NER (F1 ~92.8 on CoNLL-2003)
instead of a basic pretrained pipeline, giving significantly higher accuracy.
"""

from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
import re

# ── Model selection ──────────────────────────────────────────────────────────
# dslim/bert-large-NER: F1 92.8% on CoNLL-2003 (vs ~85% for basic spacy en_core_web_sm)
# Falls back to bert-base-NER if large model unavailable (slower machines)
PRIMARY_MODEL   = "dslim/bert-large-NER"
FALLBACK_MODEL  = "dslim/bert-base-NER"

_nlp_pipeline = None

def load_model(model_name=PRIMARY_MODEL):
    global _nlp_pipeline
    if _nlp_pipeline is None:
        try:
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model     = AutoModelForTokenClassification.from_pretrained(model_name)
            _nlp_pipeline = pipeline(
                "ner",
                model=model,
                tokenizer=tokenizer,
                aggregation_strategy="max",   # merges subword tokens → cleaner entities
                device=-1                      # CPU; change to 0 for GPU
            )
        except Exception:
            # Fallback to smaller model
            _nlp_pipeline = pipeline(
                "ner",
                model=FALLBACK_MODEL,
                aggregation_strategy="max",
                device=-1
            )
    return _nlp_pipeline


# ── Label mapping — human-readable ──────────────────────────────────────────
LABEL_MAP = {
    "PER":  "Person",
    "ORG":  "Organization",
    "LOC":  "Location",
    "MISC": "Miscellaneous"
}

LABEL_COLORS = {
    "Person":         "#4FC3F7",   # blue
    "Organization":   "#81C784",   # green
    "Location":       "#FFB74D",   # orange
    "Miscellaneous":  "#CE93D8",   # purple
}


def run_ner(text: str) -> list[dict]:
    """
    Run NER on text. Returns list of entity dicts:
    {word, label, score, start, end}
    """
    if not text or not text.strip():
        return []

    nlp = load_model()
    raw = nlp(text)

    entities = []
    for ent in raw:
        label = LABEL_MAP.get(ent["entity_group"], ent["entity_group"])
        entities.append({
            "word":  ent["word"],
            "label": label,
            "score": round(float(ent["score"]) * 100, 1),
            "start": ent["start"],
            "end":   ent["end"],
        })

    return entities


def get_highlighted_html(text: str, entities: list[dict]) -> str:
    """
    Returns HTML string with entities highlighted inline.
    Handles overlapping spans safely by processing right-to-left.
    """
    if not entities:
        return f"<p style='font-family:sans-serif;line-height:1.8'>{text}</p>"

    # Sort by start position descending to insert HTML without shifting indices
    sorted_ents = sorted(entities, key=lambda e: e["start"], reverse=True)
    result = text

    for ent in sorted_ents:
        color = LABEL_COLORS.get(ent["label"], "#E0E0E0")
        span = (
            f'<mark style="background:{color};padding:2px 6px;border-radius:4px;'
            f'margin:0 1px;font-family:sans-serif">'
            f'{result[ent["start"]:ent["end"]]}'
            f'<sup style="font-size:10px;margin-left:3px;font-weight:bold">'
            f'{ent["label"]}</sup></mark>'
        )
        result = result[:ent["start"]] + span + result[ent["end"]:]

    return f"<div style='font-family:sans-serif;line-height:2;font-size:15px'>{result}</div>"