"""Anonymized approval/feedback store + nearest-neighbour learning loop.

When a clinician clicks "Approve & Log", we record what the AI recommended and
what the clinician actually chose. Disagreements (the clinician overriding the
AI's #1, or picking a code outside the Top-3) are the strongest learning signal
— they teach the system the clinic's "peculiar" coding preferences.

How the learning works (no model retraining needed):
  * Every approval is embedded (the anonymized case summary) and stored in an
    OpenSearch k-NN index `ohip_feedback`.
  * On a new case, we k-NN the current case vector against past approvals and
    build a per-code PRIOR, weighted by similarity and up-weighted when the code
    was a clinician override. These priors are surfaced to the LLM and used to
    re-rank retrieval, so future recommendations drift toward what clinicians
    actually pick for similar cases.

PHIPA: stored summaries are de-identified before persistence and keyed by an
anonymous case id (hash). De-id has two layers:
  1. Structured regex scrubbing (health-card numbers, calendar dates/DOB, phone,
     email, long digit runs, "Name:" fields) — always on.
  2. NER scrubbing via a local spaCy model (PERSON -> [NAME], locations ->
     [LOCATION]) when spaCy + the model are installed; degrades gracefully to
     regex-only otherwise. Runs fully offline (no external calls).
Relative ages/durations ("18 months old") are deliberately preserved because
they carry clinical signal for the learning loop and are not identifiers.
"""
from __future__ import annotations

import datetime as dt
import hashlib
import logging
import re

from opensearchpy import helpers

from .config import settings
from .embeddings import embed_text
from .opensearch_client import get_client

logger = logging.getLogger(__name__)

FEEDBACK_INDEX = "ohip_feedback"

# --- PHIPA scrubbing -------------------------------------------------------
# Layer 1: structured identifiers via regex.
_HEALTH_CARD = re.compile(r"\b\d{4}[-\s]?\d{3}[-\s]?\d{3}[-\s]?[A-Z]{0,2}\b")
_EMAIL = re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]+\b")
_PHONE = re.compile(r"\b(?:\+?1[-.\s]?)?\(?\d{3}\)?[-.\s]\d{3}[-.\s]\d{4}\b")
# Calendar dates / DOB: ISO (2024-11-22), slashed (11/22/2024), or written
# ("November 22, 2024" / "Nov 22 2024"). Relative ages ("18 months") are NOT
# matched here — they are clinical signal, not identifiers.
_ISO_DATE = re.compile(r"\b\d{4}[-/]\d{1,2}[-/]\d{1,2}\b")
_SLASH_DATE = re.compile(r"\b\d{1,2}[-/]\d{1,2}[-/]\d{2,4}\b")
_WRITTEN_DATE = re.compile(
    r"(?i)\b(?:jan|feb|mar|apr|may|jun|jul|aug|sep|oct|nov|dec)[a-z]*\.?\s+\d{1,2}(?:st|nd|rd|th)?,?\s*\d{0,4}\b"
)
_LONG_DIGITS = re.compile(r"\b\d{6,}\b")
_NAME_FIELD = re.compile(
    r"(?i)\b(?:name|patient|pt|dob|mrn)\s*[:=]\s*[A-Z0-9][\w'’\-]*(?:\s+[A-Z0-9][\w'’\-]*)*"
)

# Layer 2: NER — labels we redact and their placeholders. DATE/TIME are handled
# by the date regexes above so we don't nuke ages like "18 months".
_NER_LABELS = {
    "PERSON": "[NAME]",
    "GPE": "[LOCATION]",
    "LOC": "[LOCATION]",
    "FAC": "[LOCATION]",
    "ORG": "[ORG]",
}

# Lazy singleton for the spaCy pipeline: None = untried, False = unavailable.
_NLP: object | None | bool = None


def _get_nlp():
    """Load the spaCy NER pipeline once; return None if unavailable."""
    global _NLP
    if _NLP is not None:
        return _NLP or None
    if not settings.deid_ner:
        _NLP = False
        return None
    try:
        import spacy

        # Only NER is needed; drop the rest for speed.
        _NLP = spacy.load(
            settings.deid_model,
            disable=["parser", "lemmatizer", "tagger", "attribute_ruler"],
        )
        logger.info("De-id NER model '%s' loaded", settings.deid_model)
    except Exception as exc:  # noqa: BLE001
        logger.warning(
            "De-id NER unavailable (%s); falling back to regex-only scrubbing", exc
        )
        _NLP = False
    return _NLP or None


def _regex_scrub(summary: str) -> str:
    s = _HEALTH_CARD.sub("[ID]", summary)
    s = _EMAIL.sub("[EMAIL]", s)
    s = _PHONE.sub("[PHONE]", s)
    s = _ISO_DATE.sub("[DATE]", s)
    s = _SLASH_DATE.sub("[DATE]", s)
    s = _WRITTEN_DATE.sub("[DATE]", s)
    s = _LONG_DIGITS.sub("[NUM]", s)
    s = _NAME_FIELD.sub(lambda m: m.group(0).split(":")[0].split("=")[0] + ": [REDACTED]", s)
    return s


def _ner_scrub(summary: str) -> str:
    nlp = _get_nlp()
    if nlp is None:
        return summary
    doc = nlp(summary)
    # Replace from the end so earlier offsets stay valid.
    out = summary
    for ent in sorted(doc.ents, key=lambda e: e.start_char, reverse=True):
        placeholder = _NER_LABELS.get(ent.label_)
        if placeholder:
            out = out[: ent.start_char] + placeholder + out[ent.end_char :]
    return out


def anonymize(summary: str) -> str:
    """De-identify a clinical summary (regex + NER) before persistence."""
    s = _regex_scrub(summary)
    s = _ner_scrub(s)
    return s.strip()


def case_id(anon_summary: str) -> str:
    seed = anon_summary + dt.datetime.utcnow().isoformat()
    return hashlib.sha256(seed.encode()).hexdigest()[:16]


def feedback_mapping() -> dict:
    return {
        "settings": {
            "index.knn": True,
            "number_of_shards": 1,
            "number_of_replicas": 0,
        },
        "mappings": {
            "properties": {
                "case_id": {"type": "keyword"},
                "anon_summary": {"type": "text"},
                "encounter_type": {"type": "keyword"},
                "province_code": {"type": "keyword"},
                "provider_specialty_code": {"type": "keyword"},
                "ai_top_codes": {"type": "keyword"},
                "ai_rank1": {"type": "keyword"},
                "approved_codes": {"type": "keyword"},
                "override_codes": {"type": "keyword"},
                "agreed": {"type": "boolean"},
                "note": {"type": "text"},
                "selected_claim_cad": {"type": "float"},
                "optimized_claim_cad": {"type": "float"},
                "difference_cad": {"type": "float"},
                "risk_level": {"type": "keyword"},
                "optimized_codes": {"type": "keyword"},
                "created_at": {"type": "date"},
                "case_vector": {
                    "type": "knn_vector",
                    "dimension": settings.embedding_dim,
                    "method": {
                        "name": "hnsw",
                        "space_type": "cosinesimil",
                        "engine": "lucene",
                    },
                },
            }
        },
    }


def ensure_feedback_index(client=None) -> None:
    client = client or get_client()
    if client.indices.exists(index=FEEDBACK_INDEX):
        # Recreate when embedding dim changes (e.g. local gte-768 → remote 1024).
        try:
            mapping = client.indices.get_mapping(index=FEEDBACK_INDEX)
            props = (
                mapping.get(FEEDBACK_INDEX, {})
                .get("mappings", {})
                .get("properties", {})
            )
            existing_dim = (
                props.get("case_vector", {}) or {}
            ).get("dimension")
            if existing_dim is not None and int(existing_dim) != int(
                settings.embedding_dim
            ):
                logger.warning(
                    "Feedback index dim %s != configured %s — recreating '%s'",
                    existing_dim,
                    settings.embedding_dim,
                    FEEDBACK_INDEX,
                )
                client.indices.delete(index=FEEDBACK_INDEX)
        except Exception as exc:  # noqa: BLE001
            logger.warning("Could not inspect feedback index mapping: %s", exc)
    if not client.indices.exists(index=FEEDBACK_INDEX):
        client.indices.create(index=FEEDBACK_INDEX, body=feedback_mapping())
        logger.info("Created feedback index '%s'", FEEDBACK_INDEX)


def log_approval(
    *,
    clinical_summary: str,
    encounter_type: str | None,
    ai_top_codes: list[str],
    approved_codes: list[str],
    note: str | None = None,
    case_vector: list[float] | None = None,
    selected_claim_cad: float | None = None,
    optimized_claim_cad: float | None = None,
    difference_cad: float | None = None,
    risk_level: str | None = None,
    optimized_codes: list[str] | None = None,
    provider_specialty_code: str | None = None,
    province_code: str | None = None,
) -> dict:
    """Persist one approval event; returns the recorded (anonymized) doc summary."""
    client = get_client()
    ensure_feedback_index(client)

    anon = anonymize(clinical_summary)
    cid = case_id(anon)
    ai_rank1 = ai_top_codes[0] if ai_top_codes else None
    # A disagreement = clinician did not (only) accept the AI's #1 pick.
    override_codes = [c for c in approved_codes if c not in (ai_top_codes[:1] or [])]
    agreed = bool(approved_codes) and approved_codes[0] == ai_rank1

    vector = case_vector or embed_text(anon)
    doc = {
        "case_id": cid,
        "anon_summary": anon,
        "encounter_type": encounter_type,
        "province_code": province_code or settings.default_province_code,
        "provider_specialty_code": provider_specialty_code
        or settings.default_specialty_code,
        "ai_top_codes": ai_top_codes,
        "ai_rank1": ai_rank1,
        "approved_codes": approved_codes,
        "override_codes": override_codes,
        "agreed": agreed,
        "note": note,
        "selected_claim_cad": selected_claim_cad,
        "optimized_claim_cad": optimized_claim_cad,
        "difference_cad": difference_cad,
        "risk_level": risk_level,
        "optimized_codes": optimized_codes or [],
        "created_at": dt.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S"),
        "case_vector": vector,
    }
    client.index(index=FEEDBACK_INDEX, id=cid, body=doc, refresh=True)
    logger.info(
        "Logged approval %s (agreed=%s, overrides=%s)", cid, agreed, override_codes
    )
    return {"case_id": cid, "agreed": agreed, "override_codes": override_codes}


def learned_priors(case_vector: list[float], k: int | None = None) -> dict[str, float]:
    """Return {code: prior_score} learned from similar past approvals.

    Similar past cases contribute their approved codes, weighted by vector
    similarity; clinician OVERRIDES are up-weighted so corrections dominate.
    Returns {} when no feedback has been collected yet.
    """
    k = k or settings.feedback_neighbours
    client = get_client()
    if not client.indices.exists(index=FEEDBACK_INDEX):
        return {}
    try:
        resp = client.search(
            index=FEEDBACK_INDEX,
            body={
                "size": k,
                "_source": ["approved_codes", "override_codes"],
                "query": {"knn": {"case_vector": {"vector": case_vector, "k": k}}},
            },
        )
    except Exception as exc:  # noqa: BLE001
        logger.warning("Feedback lookup failed: %s", exc)
        return {}

    priors: dict[str, float] = {}
    for hit in resp["hits"]["hits"]:
        sim = float(hit.get("_score", 0.0))  # ~similarity for cosine/lucene
        src = hit["_source"]
        overrides = set(src.get("override_codes") or [])
        for code in src.get("approved_codes") or []:
            weight = sim * (settings.feedback_override_boost if code in overrides else 1.0)
            priors[code] = priors.get(code, 0.0) + weight
    return priors