decision/engine/fhir_context.py

"""
FHIR Context Loader (Phase 3A).

Builds a PatientContext from FHIR R4 resources for context-aware risk rules.

Key features:
  - Minimum necessary principle: only fetch resources needed for the active wedge
  - Lookback periods: filter by date relevance
  - Freshness detection: warn when clinical data is stale
  - ICD-10 condition matching for risk rule evaluation
  - Medication count for polypharmacy detection
  - Pre-populated slot extraction for journey agendas

This module provides:
  1. PatientContext dataclass — normalized patient data for risk evaluation
  2. PatientContextBuilder — builds context from raw FHIR bundles
  3. FHIRClient (abstract) — interface for FHIR data fetching (Phase 3 full impl)

Safety invariants:
  - Missing context = fail closed (context-dependent rules don't fire)
  - Stale data is flagged but still used (better than no data)
  - PHI is never logged — only ICD codes and counts
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, Set

from decision.engine.config_loader import DecisionConfigLoader

logger = logging.getLogger("decision.fhir_context")


# ---------------------------------------------------------------------------
# Patient Context
# ---------------------------------------------------------------------------

@dataclass
class PatientContext:
    """
    Normalized patient context for clinical decision-making.

    Built from FHIR resources. Used by RiskClassifier for context-aware rules.
    """
    patient_id: str = ""

    # Conditions (ICD-10 codes of active problems)
    active_conditions: List[str] = field(default_factory=list)
    # Convenience sets for common condition groups
    has_chf: bool = False             # I50.*
    has_copd: bool = False            # J44.*
    has_diabetes: bool = False        # E10.*, E11.*
    has_ckd: bool = False             # N18.*
    has_cancer: bool = False          # C00-C97
    has_afib: bool = False            # I48.*
    has_depression: bool = False      # F32.*, F33.*
    has_anxiety: bool = False         # F41.*

    # Medications
    active_medication_count: int = 0
    active_medication_names: List[str] = field(default_factory=list)
    is_on_anticoagulant: bool = False
    is_on_insulin: bool = False
    is_on_opioid: bool = False

    # Demographics
    age: Optional[int] = None
    is_elderly: bool = False          # 65+
    is_very_elderly: bool = False     # 75+

    # Recent vitals
    last_bp_systolic: Optional[float] = None
    last_bp_diastolic: Optional[float] = None
    last_heart_rate: Optional[float] = None
    last_weight_kg: Optional[float] = None
    last_glucose: Optional[float] = None
    last_temperature: Optional[float] = None

    # Recent encounter
    days_since_discharge: Optional[int] = None
    discharge_diagnosis_codes: List[str] = field(default_factory=list)

    # Allergies
    allergy_count: int = 0
    drug_allergies: List[str] = field(default_factory=list)

    # Appointments
    has_upcoming_appointment: bool = False
    next_appointment_days: Optional[int] = None

    # Freshness / staleness warnings
    stale_resources: List[str] = field(default_factory=list)
    context_timestamp: Optional[datetime] = None

    def to_risk_context(self) -> Dict[str, Any]:
        """
        Convert to the dict format expected by RiskClassifier.assess(patient_context=...).
        """
        return {
            "patient_id": self.patient_id,
            "active_conditions": self.active_conditions,
            "active_medication_count": self.active_medication_count,
            "has_chf": self.has_chf,
            "has_copd": self.has_copd,
            "has_diabetes": self.has_diabetes,
            "has_ckd": self.has_ckd,
            "has_cancer": self.has_cancer,
            "has_afib": self.has_afib,
            "has_depression": self.has_depression,
            "has_anxiety": self.has_anxiety,
            "is_on_anticoagulant": self.is_on_anticoagulant,
            "is_on_insulin": self.is_on_insulin,
            "is_on_opioid": self.is_on_opioid,
            "age": self.age,
            "is_elderly": self.is_elderly,
            "is_very_elderly": self.is_very_elderly,
            "days_since_discharge": self.days_since_discharge,
            "stale_resources": self.stale_resources,
        }


# ---------------------------------------------------------------------------
# Patient Context Builder
# ---------------------------------------------------------------------------

# ICD-10 pattern groups for condition detection
_CONDITION_GROUPS = {
    "has_chf": ["I50"],
    "has_copd": ["J44"],
    "has_diabetes": ["E10", "E11"],
    "has_ckd": ["N18"],
    "has_cancer": [f"C{i:02d}" for i in range(98)],  # C00-C97
    "has_afib": ["I48"],
    "has_depression": ["F32", "F33"],
    "has_anxiety": ["F41"],
}

# Medication class detection (by name substring)
_ANTICOAGULANTS = {"warfarin", "coumadin", "heparin", "enoxaparin", "lovenox",
                   "apixaban", "eliquis", "rivaroxaban", "xarelto", "dabigatran", "pradaxa"}
_INSULINS = {"insulin", "humalog", "novolog", "lantus", "levemir", "tresiba",
             "basaglar", "admelog", "fiasp", "toujeo"}
_OPIOIDS = {"oxycodone", "hydrocodone", "morphine", "fentanyl", "codeine",
            "tramadol", "methadone", "hydromorphone", "oxymorphone", "dilaudid",
            "percocet", "vicodin", "norco"}


class PatientContextBuilder:
    """
    Builds a PatientContext from FHIR resource bundles.

    Usage:
        builder = PatientContextBuilder(config)
        context = builder.build_from_fhir(
            patient_id="P-123",
            fhir_bundle=bundle_dict,
            wedge_type="pde",
        )
    """

    def __init__(self, config: DecisionConfigLoader):
        self._config = config
        self._fhir_mappings = config.fhir_mappings
        self._freshness_thresholds = self._fhir_mappings.get("freshness_thresholds", {})

    def build_from_fhir(
        self,
        patient_id: str,
        fhir_bundle: Dict[str, Any],
        wedge_type: Optional[str] = None,
    ) -> PatientContext:
        """
        Build PatientContext from a FHIR bundle response.

        Args:
            patient_id: Patient identifier
            fhir_bundle: FHIR Bundle resource (or dict of resource lists)
            wedge_type: Active wedge for minimum necessary filtering

        Returns:
            PatientContext ready for risk rule evaluation
        """
        ctx = PatientContext(
            patient_id=patient_id,
            context_timestamp=datetime.now(timezone.utc),
        )

        # Extract resources from bundle
        resources = self._extract_resources(fhir_bundle)

        # Build context from each resource type
        self._process_patient(ctx, resources.get("Patient", []))
        self._process_conditions(ctx, resources.get("Condition", []))
        self._process_medications(ctx, resources.get("MedicationRequest", []))
        self._process_observations(ctx, resources.get("Observation", []))
        self._process_allergies(ctx, resources.get("AllergyIntolerance", []))
        self._process_encounters(ctx, resources.get("Encounter", []))
        self._process_appointments(ctx, resources.get("Appointment", []))

        # Check freshness
        self._check_freshness(ctx, resources)

        logger.info(
            "PatientContext built: patient=%s conditions=%d meds=%d age=%s chf=%s copd=%s dm=%s",
            patient_id,
            len(ctx.active_conditions),
            ctx.active_medication_count,
            ctx.age,
            ctx.has_chf,
            ctx.has_copd,
            ctx.has_diabetes,
        )

        return ctx

    def build_from_dict(
        self,
        patient_id: str,
        data: Dict[str, Any],
    ) -> PatientContext:
        """
        Build PatientContext from a pre-processed dict (e.g., from API request).

        This allows callers to pass patient context directly without FHIR.
        """
        ctx = PatientContext(
            patient_id=patient_id,
            context_timestamp=datetime.now(timezone.utc),
        )

        ctx.active_conditions = data.get("active_conditions", [])
        ctx.active_medication_count = data.get("active_medication_count", 0)
        ctx.active_medication_names = data.get("active_medication_names", [])
        ctx.age = data.get("age")
        ctx.days_since_discharge = data.get("days_since_discharge")

        # Derive flags from conditions
        self._set_condition_flags(ctx)

        # Derive flags from medications
        self._set_medication_flags(ctx)

        # Age flags
        if ctx.age:
            ctx.is_elderly = ctx.age >= 65
            ctx.is_very_elderly = ctx.age >= 75

        return ctx

    # ------------------------------------------------------------------
    # Resource processing
    # ------------------------------------------------------------------

    def _extract_resources(
        self, bundle: Dict[str, Any]
    ) -> Dict[str, List[Dict[str, Any]]]:
        """Extract resources from a FHIR Bundle, grouped by resourceType."""
        resources: Dict[str, List[Dict[str, Any]]] = {}

        # Handle standard FHIR Bundle format
        entries = bundle.get("entry", [])
        for entry in entries:
            resource = entry.get("resource", {})
            rtype = resource.get("resourceType", "Unknown")
            resources.setdefault(rtype, []).append(resource)

        # Also handle flat dict format {resourceType: [resources]}
        for key, value in bundle.items():
            if key != "entry" and isinstance(value, list):
                resources.setdefault(key, []).extend(value)

        return resources

    def _process_patient(
        self, ctx: PatientContext, patients: List[Dict[str, Any]]
    ) -> None:
        if not patients:
            return
        patient = patients[0]
        birth_date = patient.get("birthDate")
        if birth_date:
            try:
                dob = datetime.strptime(birth_date, "%Y-%m-%d")
                today = datetime.now()
                ctx.age = today.year - dob.year - ((today.month, today.day) < (dob.month, dob.day))
                ctx.is_elderly = ctx.age >= 65
                ctx.is_very_elderly = ctx.age >= 75
            except (ValueError, TypeError):
                pass

    def _process_conditions(
        self, ctx: PatientContext, conditions: List[Dict[str, Any]]
    ) -> None:
        for condition in conditions:
            # Only active conditions
            clinical_status = condition.get("clinicalStatus", {})
            if isinstance(clinical_status, dict):
                codings = clinical_status.get("coding", [])
                status_code = codings[0].get("code", "") if codings else ""
            else:
                status_code = str(clinical_status)

            if status_code and status_code not in ("active", "recurrence", "relapse"):
                continue

            # Extract ICD-10 codes
            code_concept = condition.get("code", {})
            for coding in code_concept.get("coding", []):
                system = coding.get("system", "")
                code = coding.get("code", "")
                if "icd" in system.lower() or code:
                    ctx.active_conditions.append(code)

        self._set_condition_flags(ctx)

    def _set_condition_flags(self, ctx: PatientContext) -> None:
        """Set boolean condition flags from ICD-10 codes."""
        for flag_name, prefixes in _CONDITION_GROUPS.items():
            has_condition = any(
                any(code.startswith(prefix) for prefix in prefixes)
                for code in ctx.active_conditions
            )
            setattr(ctx, flag_name, has_condition)

    def _process_medications(
        self, ctx: PatientContext, medications: List[Dict[str, Any]]
    ) -> None:
        active_meds = []
        for med in medications:
            status = med.get("status", "")
            if status not in ("active", "completed"):
                continue
            med_name = ""
            med_concept = med.get("medicationCodeableConcept", {})
            if med_concept:
                med_name = med_concept.get("text", "")
                if not med_name:
                    codings = med_concept.get("coding", [])
                    if codings:
                        med_name = codings[0].get("display", "")
            if med_name:
                active_meds.append(med_name)

        ctx.active_medication_names = active_meds
        ctx.active_medication_count = len(active_meds)
        self._set_medication_flags(ctx)

    def _set_medication_flags(self, ctx: PatientContext) -> None:
        """Set medication class flags from med names."""
        names_lower = {n.lower() for n in ctx.active_medication_names}
        ctx.is_on_anticoagulant = bool(names_lower & _ANTICOAGULANTS)
        ctx.is_on_insulin = bool(names_lower & _INSULINS)
        ctx.is_on_opioid = bool(names_lower & _OPIOIDS)

    def _process_observations(
        self, ctx: PatientContext, observations: List[Dict[str, Any]]
    ) -> None:
        # Sort by date descending to get most recent first
        observations.sort(
            key=lambda o: o.get("effectiveDateTime", ""),
            reverse=True,
        )

        for obs in observations:
            code_concept = obs.get("code", {})
            codings = code_concept.get("coding", [])
            loinc_code = ""
            for coding in codings:
                if "loinc" in coding.get("system", "").lower():
                    loinc_code = coding.get("code", "")
                    break

            value = obs.get("valueQuantity", {}).get("value")
            if value is None:
                continue

            # Map LOINC codes to context fields (most recent only)
            if loinc_code == "8480-6" and ctx.last_bp_systolic is None:  # Systolic BP
                ctx.last_bp_systolic = float(value)
            elif loinc_code == "8462-4" and ctx.last_bp_diastolic is None:  # Diastolic BP
                ctx.last_bp_diastolic = float(value)
            elif loinc_code == "8867-4" and ctx.last_heart_rate is None:  # Heart rate
                ctx.last_heart_rate = float(value)
            elif loinc_code == "29463-7" and ctx.last_weight_kg is None:  # Weight
                ctx.last_weight_kg = float(value)
            elif loinc_code in ("2339-0", "2345-7") and ctx.last_glucose is None:  # Glucose
                ctx.last_glucose = float(value)
            elif loinc_code == "8310-5" and ctx.last_temperature is None:  # Temperature
                ctx.last_temperature = float(value)

    def _process_allergies(
        self, ctx: PatientContext, allergies: List[Dict[str, Any]]
    ) -> None:
        active_allergies = [
            a for a in allergies
            if a.get("clinicalStatus", {}).get("coding", [{}])[0].get("code") == "active"
            or not a.get("clinicalStatus")
        ]
        ctx.allergy_count = len(active_allergies)

        for allergy in active_allergies:
            category = allergy.get("category", [])
            if "medication" in category:
                substance = allergy.get("code", {}).get("text", "")
                if substance:
                    ctx.drug_allergies.append(substance)

    def _process_encounters(
        self, ctx: PatientContext, encounters: List[Dict[str, Any]]
    ) -> None:
        # Find most recent discharge
        encounters.sort(
            key=lambda e: e.get("period", {}).get("end", ""),
            reverse=True,
        )

        for enc in encounters:
            period = enc.get("period", {})
            end_date = period.get("end")
            if not end_date:
                continue
            try:
                discharge_dt = datetime.fromisoformat(end_date.replace("Z", "+00:00"))
                now = datetime.now(timezone.utc)
                delta = now - discharge_dt
                ctx.days_since_discharge = delta.days

                # Extract discharge diagnosis codes
                diagnoses = enc.get("diagnosis", [])
                for diag in diagnoses:
                    code = diag.get("condition", {}).get("reference", "")
                    if code:
                        ctx.discharge_diagnosis_codes.append(code)
                break  # Most recent only
            except (ValueError, TypeError):
                continue

    def _process_appointments(
        self, ctx: PatientContext, appointments: List[Dict[str, Any]]
    ) -> None:
        now = datetime.now(timezone.utc)
        for appt in appointments:
            status = appt.get("status", "")
            if status in ("cancelled", "noshow", "entered-in-error"):
                continue
            start = appt.get("start")
            if not start:
                continue
            try:
                appt_dt = datetime.fromisoformat(start.replace("Z", "+00:00"))
                if appt_dt > now:
                    ctx.has_upcoming_appointment = True
                    delta = appt_dt - now
                    if ctx.next_appointment_days is None or delta.days < ctx.next_appointment_days:
                        ctx.next_appointment_days = delta.days
            except (ValueError, TypeError):
                continue

    def _check_freshness(
        self, ctx: PatientContext, resources: Dict[str, List[Dict[str, Any]]]
    ) -> None:
        """Check resource freshness against configured thresholds."""
        now = datetime.now(timezone.utc)

        for resource_type, threshold_hours in self._freshness_thresholds.items():
            if threshold_hours is None:
                continue

            # Map compound types (e.g., "Observation_laboratory") to base type
            base_type = resource_type.split("_")[0]
            entries = resources.get(base_type, [])
            if not entries:
                continue

            # Find most recent entry date
            latest_date = None
            for entry in entries:
                date_str = (
                    entry.get("effectiveDateTime")
                    or entry.get("authoredOn")
                    or entry.get("recordedDate")
                    or entry.get("meta", {}).get("lastUpdated")
                )
                if date_str:
                    try:
                        dt = datetime.fromisoformat(date_str.replace("Z", "+00:00"))
                        if latest_date is None or dt > latest_date:
                            latest_date = dt
                    except (ValueError, TypeError):
                        continue

            if latest_date:
                hours_old = (now - latest_date).total_seconds() / 3600
                if hours_old > threshold_hours:
                    ctx.stale_resources.append(
                        f"{resource_type} (last updated {hours_old:.0f}h ago, threshold {threshold_hours}h)"
                    )
                    logger.warning(
                        "Stale FHIR data: %s is %.0fh old (threshold: %dh)",
                        resource_type,
                        hours_old,
                        threshold_hours,
                    )