"""Optional spaCy-based NER detector.

Falls back gracefully when spaCy or its model is not installed.
Runs only on text segments that prior layers flagged as suspicious,
keeping inference cost low on CPU-only deployments.
"""
from typing import Any, Dict, List, Optional, Union

from upf.config import UPFConfig

# spaCy entity labels we care about → UPF type
_LABEL_MAP: Dict[str, str] = {
    "PERSON": "NAME",
    "GPE": "ADDRESS",       # geo-political entity (city, country)
    "LOC": "ADDRESS",       # non-GPE location
    "FAC": "ADDRESS",       # facility
    "ORG": "ORG",
    "CARDINAL": None,       # skip pure numbers
    "DATE": None,           # skip dates
    "TIME": None,
    "PERCENT": None,
    "MONEY": None,
    "QUANTITY": None,
    "ORDINAL": None,
    "LANGUAGE": None,
    "NORP": None,           # nationalities, religious groups
    "WORK_OF_ART": None,
    "EVENT": None,
    "LAW": None,
    "PRODUCT": None,
}

# Cache the loaded model to avoid re-loading on every call.
_nlp_cache: Optional[Any] = None
_spacy_available: Optional[bool] = None


def _get_nlp() -> Optional[Any]:
    global _nlp_cache, _spacy_available

    if _spacy_available is False:
        return None
    if _nlp_cache is not None:
        return _nlp_cache

    try:
        import spacy  # noqa: PLC0415

        # Try small English model first, then medium.
        for model in ("en_core_web_sm", "en_core_web_md", "en_core_web_lg"):
            try:
                _nlp_cache = spacy.load(model, disable=["parser", "lemmatizer"])
                _spacy_available = True
                return _nlp_cache
            except OSError:
                continue

        # No model found — mark unavailable so we don't retry every call.
        _spacy_available = False
        return None
    except ImportError:
        _spacy_available = False
        return None


def is_available() -> bool:
    """Return True if spaCy and at least one NER model are loadable."""
    return _get_nlp() is not None


def detect_ml_entities(
    text: str,
    config: UPFConfig,
    suspicious_spans: Optional[List[Dict[str, Union[int, str]]]] = None,
) -> List[Dict[str, Union[int, str]]]:
    """Run spaCy NER on the text and return findings.

    Args:
        text: full input text.
        config: UPF configuration.
        suspicious_spans: optional list of findings from prior detectors.
            When provided, spaCy only processes windows around those spans,
            reducing CPU cost. When None, the full text is processed.
    """
    nlp = _get_nlp()
    if nlp is None:
        return []

    findings: List[Dict[str, Union[int, str]]] = []
    seen: set = set()

    def _add(start: int, end: int, text_span: str, label: str) -> None:
        upf_type = _LABEL_MAP.get(label)
        if upf_type is None:
            return
        # Respect config flags.
        if upf_type == "NAME" and not config.redact_names:
            return
        if upf_type == "ADDRESS" and not config.redact_addresses:
            return
        key = (start, end)
        if key not in seen:
            seen.add(key)
            findings.append({"type": upf_type, "start": start, "end": end, "text": text_span})

    if suspicious_spans:
        # Process windows around each suspicious span to keep inference fast.
        window = 120  # chars on each side
        processed_ranges: List[tuple] = []

        for span in suspicious_spans:
            s = max(0, int(span["start"]) - window)
            e = min(len(text), int(span["end"]) + window)

            # Merge overlapping windows.
            if processed_ranges and s <= processed_ranges[-1][1]:
                processed_ranges[-1] = (processed_ranges[-1][0], max(e, processed_ranges[-1][1]))
            else:
                processed_ranges.append((s, e))

        for seg_start, seg_end in processed_ranges:
            segment = text[seg_start:seg_end]
            doc = nlp(segment)
            for ent in doc.ents:
                _add(seg_start + ent.start_char, seg_start + ent.end_char, ent.text, ent.label_)
    else:
        doc = nlp(text)
        for ent in doc.ents:
            _add(ent.start_char, ent.end_char, ent.text, ent.label_)

    return findings


def reset_cache() -> None:
    """Reset the model cache. Useful in tests."""
    global _nlp_cache, _spacy_available
    _nlp_cache = None
    _spacy_available = None
