"""NER-based PII scanner using spaCy or transformers NER models."""

from __future__ import annotations

import logging

from bilgeai.analyzer.models import PIIFinding

logger = logging.getLogger(__name__)

# Entity types that indicate PII
PII_ENTITY_TYPES = {
    "PERSON",
    "PER",
    "ORG",
    "GPE",
    "LOC",
    "FAC",
    "EMAIL",
    "PHONE",
    "CREDIT_CARD",
    "DATE_TIME",
    "NRP",  # Nationality/Religious/Political group
}


def scan_with_ner(
    rows: list[dict],
    model_name: str = "en_core_web_sm",
    max_rows: int = 1000,
) -> list[PIIFinding]:
    """Scan dataset rows for PII using NER.

    Tries spaCy first, then falls back to transformers NER pipeline.
    Returns list of PIIFinding for each entity type detected.
    """
    # Sample rows if dataset is large
    sample = rows[:max_rows]
    texts = [" ".join(str(v) for v in row.values()) for row in sample]

    findings = _try_spacy(texts, model_name)
    if findings is not None:
        return findings

    findings = _try_transformers(texts)
    if findings is not None:
        return findings

    logger.info(
        "NER scanning unavailable. Install spaCy or transformers for NER-based PII detection."
    )
    return []


def _try_spacy(texts: list[str], model_name: str) -> list[PIIFinding] | None:
    """Try scanning with spaCy NER."""
    try:
        import spacy

        try:
            nlp = spacy.load(model_name)
        except OSError:
            logger.debug("spaCy model '%s' not found", model_name)
            return None

        entity_indices: dict[str, list[int]] = {}

        for i, text in enumerate(texts):
            doc = nlp(text[:10000])  # Limit text length for performance
            for ent in doc.ents:
                if ent.label_ in PII_ENTITY_TYPES:
                    label = ent.label_
                    if label not in entity_indices:
                        entity_indices[label] = []
                    entity_indices[label].append(i)

        return _build_findings(entity_indices)

    except ImportError:
        logger.debug("spaCy not installed")
        return None


def _try_transformers(texts: list[str]) -> list[PIIFinding] | None:
    """Try scanning with transformers NER pipeline."""
    try:
        from transformers import pipeline

        ner_pipeline = pipeline(
            "ner",
            model="dslim/bert-base-NER",
            aggregation_strategy="simple",
        )

        entity_indices: dict[str, list[int]] = {}

        for i, text in enumerate(texts):
            results = ner_pipeline(text[:512])  # BERT has 512 token limit
            for entity in results:
                label = entity["entity_group"]
                if label in PII_ENTITY_TYPES or label in ("PER", "ORG", "LOC"):
                    if label not in entity_indices:
                        entity_indices[label] = []
                    entity_indices[label].append(i)

        return _build_findings(entity_indices)

    except (ImportError, Exception):
        logger.debug("Transformers NER pipeline failed", exc_info=True)
        return None


def _build_findings(entity_indices: dict[str, list[int]]) -> list[PIIFinding]:
    """Convert entity-index mapping to PIIFinding list."""
    findings = []
    for entity_type, indices in entity_indices.items():
        unique_indices = sorted(set(indices))
        if unique_indices:
            findings.append(
                PIIFinding(
                    pii_type=f"ner_{entity_type.lower()}",
                    count=len(unique_indices),
                    sample_indices=unique_indices[:5],
                )
            )
    return findings
