"""
Sanitizer service — lazy-loads NER pipelines per model and caches them.
The default pipeline is pre-loaded at startup; additional models are loaded
on first use and kept in memory for the lifetime of the process.
"""

from __future__ import annotations

from typing import Any, Optional

from transformers import pipeline as hf_pipeline

from app.schemas.sanitize import EntityDetail, SanitizeResponse, SanitizeStats

# Replacement map: NER label → placeholder
LABEL_REPLACEMENTS: dict[str, str] = {
    "PER": "[PRIVATE_PERSON]",
    "PERSON": "[PRIVATE_PERSON]",
    "EMAIL": "[PRIVATE_EMAIL]",
    "PHONE": "[PRIVATE_PHONE]",
    "PHONE_NUM": "[PRIVATE_PHONE]",
    "ADDRESS": "[PRIVATE_ADDRESS]",
    "ADDR": "[PRIVATE_ADDRESS]",
    "LOC": "[PRIVATE_LOCATION]",
    "LOCATION": "[PRIVATE_LOCATION]",
    "GPE": "[PRIVATE_LOCATION]",
    "ORG": "[PRIVATE_ORGANIZATION]",
    "ORGANIZATION": "[PRIVATE_ORGANIZATION]",
    "DATE": "[PRIVATE_DATE]",
    "TIME": "[PRIVATE_DATE]",
    "ID": "[PRIVATE_ID]",
    "ID_NUM": "[PRIVATE_ID]",
    "ID_NUMBER": "[PRIVATE_ID]",
}

DEFAULT_REPLACEMENT = "[PRIVATE_DATA]"

# Canonical display label per replacement (for stats grouping)
REPLACEMENT_TO_TYPE: dict[str, str] = {
    "[PRIVATE_PERSON]": "PERSON",
    "[PRIVATE_EMAIL]": "EMAIL",
    "[PRIVATE_PHONE]": "PHONE",
    "[PRIVATE_ADDRESS]": "ADDRESS",
    "[PRIVATE_LOCATION]": "LOCATION",
    "[PRIVATE_ORGANIZATION]": "ORGANIZATION",
    "[PRIVATE_DATE]": "DATE",
    "[PRIVATE_ID]": "ID_NUMBER",
    "[PRIVATE_DATA]": "OTHER",
}

# Pipeline cache: model_id → loaded HuggingFace pipeline
_pipeline_cache: dict[str, Any] = {}
_default_model: Optional[str] = None


def preload_pipeline(model_id: str) -> None:
    """Pre-load and cache a pipeline (called at startup for the default model)."""
    global _default_model
    _get_or_load(model_id)
    _default_model = model_id


def _get_or_load(model_id: str) -> Any:
    """Return cached pipeline or load it on first use."""
    if model_id not in _pipeline_cache:
        _pipeline_cache[model_id] = hf_pipeline(
            "token-classification",
            model=model_id,
            aggregation_strategy="simple",
        )
    return _pipeline_cache[model_id]


# Keep backward-compatible set_pipeline for tests
def set_pipeline(pipeline: Any, model_id: str = "__test__") -> None:
    global _default_model
    _pipeline_cache[model_id] = pipeline
    if _default_model is None:
        _default_model = model_id


def sanitize_text(text: str, model_id: Optional[str] = None) -> SanitizeResponse:
    if not text or not text.strip():
        return _empty_response(text)

    resolved_model = model_id or _default_model
    if resolved_model is None:
        raise RuntimeError("NER pipeline has not been initialised")

    active_pipeline = _get_or_load(resolved_model)
    raw_entities: list[dict] = active_pipeline(text)

    # Sort by start position descending so replacements don't shift indices
    raw_entities.sort(key=lambda e: e["start"], reverse=True)

    sanitized = text
    entity_details: list[EntityDetail] = []

    for ent in raw_entities:
        label: str = ent.get("entity_group", ent.get("entity", "OTHER")).upper()
        replacement = LABEL_REPLACEMENTS.get(label, DEFAULT_REPLACEMENT)
        start: int = ent["start"]
        end: int = ent["end"]
        original_value: str = text[start:end]
        score: Optional[float] = ent.get("score")

        sanitized = sanitized[:start] + replacement + sanitized[end:]

        entity_details.append(
            EntityDetail(
                type=REPLACEMENT_TO_TYPE.get(replacement, "OTHER"),
                original_value=original_value,
                replacement=replacement,
                start=start,
                end=end,
                score=round(score, 4) if score is not None else None,
            )
        )

    # Re-sort entities by original position (ascending) for the response
    entity_details.sort(key=lambda e: e.start)

    by_type: dict[str, int] = {}
    for ent in entity_details:
        by_type[ent.type] = by_type.get(ent.type, 0) + 1

    original_length = len(text)
    sanitized_length = len(sanitized)
    sensitive_chars = sum(e.end - e.start for e in entity_details)
    sensitive_ratio = round(sensitive_chars / original_length, 4) if original_length > 0 else 0.0

    return SanitizeResponse(
        original_text=text,
        sanitized_text=sanitized,
        entities=entity_details,
        stats=SanitizeStats(
            total_entities=len(entity_details),
            by_type=by_type,
            original_length=original_length,
            sanitized_length=sanitized_length,
            sensitive_ratio=sensitive_ratio,
        ),
    )


def _empty_response(text: str) -> SanitizeResponse:
    return SanitizeResponse(
        original_text=text,
        sanitized_text=text,
        entities=[],
        stats=SanitizeStats(
            total_entities=0,
            by_type={},
            original_length=len(text),
            sanitized_length=len(text),
            sensitive_ratio=0.0,
        ),
    )