from __future__ import annotations

from collections import Counter, defaultdict
from typing import Optional

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

router = APIRouter(prefix="/dataset")


# ── Entity config ──────────────────────────────────────────────────────────────

PII_ENTITIES = [
    # ── Global ───────────────────────────────────────────────────────────────────
    # PERSON, LOCATION, ORGANIZATION, DATE_TIME, AGE, URL omitted — too noisy in
    # normal training text (any named entity / date / link fires), inflating 100%.
    "EMAIL_ADDRESS", "PHONE_NUMBER",
    "CREDIT_CARD", "IBAN_CODE", "MEDICAL_LICENSE",
    "IP_ADDRESS", "MAC_ADDRESS", "NRP",
    "CRYPTO",
    # ── USA ──────────────────────────────────────────────────────────────────────
    "US_SSN", "US_PASSPORT", "US_DRIVER_LICENSE", "US_ITIN",
    "US_BANK_NUMBER", "US_NPI", "US_MBI",
    # ── UK ───────────────────────────────────────────────────────────────────────
    "UK_NHS", "UK_NINO", "UK_PASSPORT", "UK_POSTCODE", "UK_VEHICLE_REGISTRATION",
    # ── Spain ────────────────────────────────────────────────────────────────────
    "ES_NIF", "ES_NIE",
    # ── Italy ────────────────────────────────────────────────────────────────────
    "IT_FISCAL_CODE", "IT_DRIVER_LICENSE", "IT_VAT_CODE",
    "IT_PASSPORT", "IT_IDENTITY_CARD",
    # ── Poland ───────────────────────────────────────────────────────────────────
    "PL_PESEL",
    # ── Singapore ────────────────────────────────────────────────────────────────
    "SG_NRIC_FIN", "SG_UEN",
    # ── Australia ────────────────────────────────────────────────────────────────
    "AU_ABN", "AU_ACN", "AU_TFN", "AU_MEDICARE",
    # ── India ────────────────────────────────────────────────────────────────────
    "IN_PAN", "IN_AADHAAR", "IN_VEHICLE_REGISTRATION",
    "IN_VOTER", "IN_PASSPORT", "IN_GSTIN",
    # ── Finland ──────────────────────────────────────────────────────────────────
    "FI_PERSONAL_IDENTITY_CODE",
    # ── Korea ────────────────────────────────────────────────────────────────────
    "KR_DRIVER_LICENSE", "KR_FRN", "KR_PASSPORT", "KR_BRN", "KR_RRN",
    # ── Nigeria ──────────────────────────────────────────────────────────────────
    "NG_NIN", "NG_VEHICLE_REGISTRATION",
    # ── Thailand ─────────────────────────────────────────────────────────────────
    "TH_TNIN",
    # ── Medical / Clinical (HuggingFace NER) ─────────────────────────────────────
    "MEDICAL_DISEASE_DISORDER", "MEDICAL_MEDICATION",
    "MEDICAL_THERAPEUTIC_PROCEDURE", "MEDICAL_CLINICAL_EVENT",
    "MEDICAL_BIOLOGICAL_ATTRIBUTE", "MEDICAL_BIOLOGICAL_STRUCTURE",
    "MEDICAL_FAMILY_HISTORY", "MEDICAL_HISTORY",
]

PII_RISK: dict[str, str] = {
    # Global
    "PERSON":                        "medium",
    "EMAIL_ADDRESS":                 "high",
    "PHONE_NUMBER":                  "high",
    "LOCATION":                      "medium",
    "ORGANIZATION":                  "low",
    "CREDIT_CARD":                   "critical",
    "IBAN_CODE":                     "critical",
    "MEDICAL_LICENSE":               "critical",
    "IP_ADDRESS":                    "medium",
    "MAC_ADDRESS":                   "medium",
    "URL":                           "low",
    "NRP":                           "medium",
    "DATE_TIME":                     "low",
    "AGE":                           "low",
    "CRYPTO":                        "high",
    # USA
    "US_SSN":                        "critical",
    "US_PASSPORT":                   "critical",
    "US_DRIVER_LICENSE":             "high",
    "US_ITIN":                       "critical",
    "US_BANK_NUMBER":                "critical",
    "US_NPI":                        "high",
    "US_MBI":                        "critical",
    # UK
    "UK_NHS":                        "critical",
    "UK_NINO":                       "critical",
    "UK_PASSPORT":                   "critical",
    "UK_POSTCODE":                   "medium",
    "UK_VEHICLE_REGISTRATION":       "medium",
    # Spain
    "ES_NIF":                        "critical",
    "ES_NIE":                        "critical",
    # Italy
    "IT_FISCAL_CODE":                "critical",
    "IT_DRIVER_LICENSE":             "high",
    "IT_VAT_CODE":                   "high",
    "IT_PASSPORT":                   "critical",
    "IT_IDENTITY_CARD":              "critical",
    # Poland
    "PL_PESEL":                      "critical",
    # Singapore
    "SG_NRIC_FIN":                   "critical",
    "SG_UEN":                        "high",
    # Australia
    "AU_ABN":                        "high",
    "AU_ACN":                        "high",
    "AU_TFN":                        "critical",
    "AU_MEDICARE":                   "critical",
    # India
    "IN_PAN":                        "critical",
    "IN_AADHAAR":                    "critical",
    "IN_VEHICLE_REGISTRATION":       "medium",
    "IN_VOTER":                      "high",
    "IN_PASSPORT":                   "critical",
    "IN_GSTIN":                      "high",
    # Finland
    "FI_PERSONAL_IDENTITY_CODE":     "critical",
    # Korea
    "KR_DRIVER_LICENSE":             "high",
    "KR_FRN":                        "critical",
    "KR_PASSPORT":                   "critical",
    "KR_BRN":                        "high",
    "KR_RRN":                        "critical",
    # Nigeria
    "NG_NIN":                        "critical",
    "NG_VEHICLE_REGISTRATION":       "medium",
    # Thailand
    "TH_TNIN":                       "critical",
    # Medical
    "MEDICAL_DISEASE_DISORDER":      "high",
    "MEDICAL_MEDICATION":            "high",
    "MEDICAL_THERAPEUTIC_PROCEDURE": "high",
    "MEDICAL_CLINICAL_EVENT":        "high",
    "MEDICAL_BIOLOGICAL_ATTRIBUTE":  "high",
    "MEDICAL_BIOLOGICAL_STRUCTURE":  "medium",
    "MEDICAL_FAMILY_HISTORY":        "high",
    "MEDICAL_HISTORY":               "high",
}

PII_CATEGORY: dict[str, str] = {
    # Global
    "PERSON":                        "identity",
    "EMAIL_ADDRESS":                 "contact",
    "PHONE_NUMBER":                  "contact",
    "LOCATION":                      "location",
    "ORGANIZATION":                  "identity",
    "CREDIT_CARD":                   "financial",
    "IBAN_CODE":                     "financial",
    "MEDICAL_LICENSE":               "health",
    "IP_ADDRESS":                    "technical",
    "MAC_ADDRESS":                   "technical",
    "URL":                           "technical",
    "NRP":                           "sensitive",
    "DATE_TIME":                     "temporal",
    "AGE":                           "demographic",
    "CRYPTO":                        "financial",
    # USA
    "US_SSN":                        "identity",
    "US_PASSPORT":                   "identity",
    "US_DRIVER_LICENSE":             "identity",
    "US_ITIN":                       "financial",
    "US_BANK_NUMBER":                "financial",
    "US_NPI":                        "health",
    "US_MBI":                        "health",
    # UK
    "UK_NHS":                        "health",
    "UK_NINO":                       "identity",
    "UK_PASSPORT":                   "identity",
    "UK_POSTCODE":                   "location",
    "UK_VEHICLE_REGISTRATION":       "identity",
    # Spain
    "ES_NIF":                        "identity",
    "ES_NIE":                        "identity",
    # Italy
    "IT_FISCAL_CODE":                "identity",
    "IT_DRIVER_LICENSE":             "identity",
    "IT_VAT_CODE":                   "financial",
    "IT_PASSPORT":                   "identity",
    "IT_IDENTITY_CARD":              "identity",
    # Poland
    "PL_PESEL":                      "identity",
    # Singapore
    "SG_NRIC_FIN":                   "identity",
    "SG_UEN":                        "financial",
    # Australia
    "AU_ABN":                        "financial",
    "AU_ACN":                        "financial",
    "AU_TFN":                        "financial",
    "AU_MEDICARE":                   "health",
    # India
    "IN_PAN":                        "financial",
    "IN_AADHAAR":                    "identity",
    "IN_VEHICLE_REGISTRATION":       "identity",
    "IN_VOTER":                      "identity",
    "IN_PASSPORT":                   "identity",
    "IN_GSTIN":                      "financial",
    # Finland
    "FI_PERSONAL_IDENTITY_CODE":     "identity",
    # Korea
    "KR_DRIVER_LICENSE":             "identity",
    "KR_FRN":                        "financial",
    "KR_PASSPORT":                   "identity",
    "KR_BRN":                        "financial",
    "KR_RRN":                        "identity",
    # Nigeria
    "NG_NIN":                        "identity",
    "NG_VEHICLE_REGISTRATION":       "identity",
    # Thailand
    "TH_TNIN":                       "identity",
    # Medical
    "MEDICAL_DISEASE_DISORDER":      "health",
    "MEDICAL_MEDICATION":            "health",
    "MEDICAL_THERAPEUTIC_PROCEDURE": "health",
    "MEDICAL_CLINICAL_EVENT":        "health",
    "MEDICAL_BIOLOGICAL_ATTRIBUTE":  "health",
    "MEDICAL_BIOLOGICAL_STRUCTURE":  "health",
    "MEDICAL_FAMILY_HISTORY":        "health",
    "MEDICAL_HISTORY":               "health",
}

RISK_ORDER = {"critical": 0, "high": 1, "medium": 2, "low": 3}


# ── Analyzer (cached) ──────────────────────────────────────────────────────────

_analyzer = None

def _get_analyzer():
    global _analyzer
    if _analyzer is None:
        try:
            from presidio_analyzer import AnalyzerEngine
            from presidio_analyzer.nlp_engine import NlpEngineProvider
        except ImportError:
            raise RuntimeError(
                "presidio-analyzer not installed. Run: pip install presidio-analyzer"
            )
        try:
            import spacy
            if not spacy.util.is_package("en_core_web_lg"):
                raise RuntimeError(
                    "en_core_web_lg missing. Run: python -m spacy download en_core_web_lg"
                )
        except ImportError:
            raise RuntimeError("spacy not installed. Run: pip install spacy")

        provider = NlpEngineProvider(nlp_configuration={
            "nlp_engine_name": "spacy",
            "models": [{"lang_code": "en", "model_name": "en_core_web_lg"}],
        })
        _analyzer = AnalyzerEngine(
            nlp_engine=provider.create_engine(),
            supported_languages=["en"],
        )
        print("[pii] presidio loaded with en_core_web_lg", flush=True)
    return _analyzer


# ── Token helper ───────────────────────────────────────────────────────────────

def _tokens(text: str) -> int:
    return max(len(text.split()), 1)


# ── Core ───────────────────────────────────────────────────────────────────────

def run_pii_analysis(
    rows: list[dict],
    text_columns: list[str],
    max_rows: int = 300,
    min_score: float = 0.7,
) -> dict:
    if not text_columns:
        return {"error": "No text columns for PII analysis"}

    try:
        analyzer = _get_analyzer()
    except RuntimeError as e:
        return {"error": str(e)}

    sample = rows[:max_rows]

    # Per-column accumulators
    # col -> {entity_type -> count}
    col_entity_counts: dict[str, Counter] = {c: Counter() for c in text_columns}
    col_token_totals:  dict[str, int]     = {c: 0 for c in text_columns}
    col_pii_rows:      dict[str, int]     = {c: 0 for c in text_columns}

    # Global accumulators
    global_entity_counts: Counter = Counter()
    total_rows_with_pii = 0

    for row in sample:
        row_had_pii = False
        for col in text_columns:
            val = row.get(col)
            if not val or not isinstance(val, str) or not val.strip():
                continue

            text = val.strip()
            col_token_totals[col] += _tokens(text)

            try:
                results = analyzer.analyze(
                    text=text,
                    entities=PII_ENTITIES,
                    language="en",
                    score_threshold=min_score,
                )
            except Exception:
                continue

            if results:
                col_pii_rows[col] += 1
                row_had_pii = True
                for r in results:
                    col_entity_counts[col][r.entity_type] += 1
                    global_entity_counts[r.entity_type] += 1

        if row_had_pii:
            total_rows_with_pii += 1

    total_rows = len(sample)

    # ── Per-column density summary ─────────────────────────────────────────────
    column_summaries = []
    for col in text_columns:
        token_total  = col_token_totals[col]
        entity_total = sum(col_entity_counts[col].values())
        pii_rows     = col_pii_rows[col]

        if token_total == 0:
            continue

        density_per_100 = round(entity_total / token_total * 100, 2)
        pii_row_pct     = round(pii_rows / total_rows * 100, 1)

        # Highest risk found in this column
        risks_found = {PII_RISK.get(e, "low") for e in col_entity_counts[col]}
        peak_risk   = min(risks_found, key=lambda r: RISK_ORDER.get(r, 99)) if risks_found else None

        top_entities = [
            {
                "entity":   etype,
                "count":    cnt,
                "risk":     PII_RISK.get(etype, "low"),
                "category": PII_CATEGORY.get(etype, "other"),
            }
            for etype, cnt in col_entity_counts[col].most_common(8)
        ]

        column_summaries.append({
            "column":           col,
            "pii_rows":         pii_rows,
            "pii_row_pct":      pii_row_pct,
            "total_entities":   entity_total,
            "density_per_100":  density_per_100,
            "peak_risk":        peak_risk,
            "top_entities":     top_entities,
        })

    # Sort hottest column first
    column_summaries.sort(key=lambda x: x["density_per_100"], reverse=True)

    # ── Global entity breakdown ────────────────────────────────────────────────
    global_entities = [
        {
            "entity":   etype,
            "count":    cnt,
            "risk":     PII_RISK.get(etype, "low"),
            "category": PII_CATEGORY.get(etype, "other"),
        }
        for etype, cnt in global_entity_counts.most_common()
    ]

    # ── Category rollup ───────────────────────────────────────────────────────
    category_counts: Counter = Counter()
    for etype, cnt in global_entity_counts.items():
        category_counts[PII_CATEGORY.get(etype, "other")] += cnt
    category_total = sum(category_counts.values()) or 1
    category_breakdown = [
        {
            "category": cat,
            "count":    cnt,
            "pct":      round(cnt / category_total * 100, 1),
        }
        for cat, cnt in category_counts.most_common()
    ]

    # ── Overall risk level ────────────────────────────────────────────────────
    all_risks = {PII_RISK.get(e, "low") for e in global_entity_counts}
    overall_risk = min(all_risks, key=lambda r: RISK_ORDER.get(r, 99)) if all_risks else "none"

    pii_row_pct = round(total_rows_with_pii / total_rows * 100, 1) if total_rows else 0

    return {
        "rows_analyzed":      total_rows,
        "rows_with_pii":      total_rows_with_pii,
        "pii_row_pct":        pii_row_pct,
        "overall_risk":       overall_risk,
        "total_entities":     sum(global_entity_counts.values()),
        "global_entities":    global_entities,
        "category_breakdown": category_breakdown,
        "column_summaries":   column_summaries,
    }


# ── Request model ──────────────────────────────────────────────────────────────

class PIIRequest(BaseModel):
    rows:         list[dict]
    text_columns: list[str]
    max_rows:     int   = 300
    min_score:    float = 0.7


# ── Route ──────────────────────────────────────────────────────────────────────

@router.post("/pii")
async def pii_detection(req: PIIRequest):
    if not req.rows:
        raise HTTPException(status_code=400, detail="No rows provided")
    return run_pii_analysis(
        req.rows,
        req.text_columns,
        max_rows=req.max_rows,
        min_score=req.min_score,
    )