from __future__ import annotations

from typing import Optional

from .models import CountyRisk, EpidemiologicalContext, PatientFacts


class EpidemiologicalContextEngine:
    """Computes context-sensitive risk modifiers for adaptive questioning."""

    def __init__(self, context: EpidemiologicalContext) -> None:
        self.context = context

    def proximity_risk_score(self, facts: PatientFacts) -> float:
        score = 0.0
        district = (facts.location or "").strip().lower()
        active = {item.lower() for item in self.context.active_outbreak_districts}
        neighboring = {item.lower() for item in self.context.neighboring_outbreak_districts}

        if district and district in active:
            score += 3.0
        elif district and district in neighboring:
            score += 1.5

        if self.context.community_transmission:
            score += 1.0

        if self.context.recent_confirmed_cases > 0:
            score += min(2.0, self.context.recent_confirmed_cases * 0.25)

        if self.context.cross_border_alerts:
            score += 0.5

        if self.context.health_facility_alerts:
            score += 0.5

        return score

    def get_kenya_county_risk(self, county_name: Optional[str]) -> Optional[CountyRisk]:
        """Get risk intelligence for a Kenya county using case-insensitive matching."""
        if not county_name or not self.context.county_risks:
            return None

        lookup = county_name.strip().lower()
        for name, risk in self.context.county_risks.items():
            if name.lower() == lookup:
                return risk
        return None

    def get_question_priority_hints(self, facts: PatientFacts) -> dict:
        """Generate county-informed hints that guide question priority for the LLM."""
        hints = {
            "county_name": facts.location,
            "county_risk": None,
            "is_border_county": False,
            "corridor_flags": [],
            "relevant_poes": [],
            "high_risk_profiles": [],
            "key_risk_factors": [],
            "question_priority_signal": "standard",
            "rationale": "",
        }

        county_risk = self.get_kenya_county_risk(facts.location)
        if not county_risk:
            return hints

        hints["county_risk"] = county_risk.risk_tier
        hints["is_border_county"] = county_risk.is_border_county
        hints["corridor_flags"] = county_risk.corridor_flags
        hints["relevant_poes"] = county_risk.relevant_poes
        hints["high_risk_profiles"] = county_risk.high_risk_profiles
        hints["key_risk_factors"] = county_risk.key_risk_factors

        if county_risk.risk_tier in ("very_high", "high"):
            if county_risk.is_border_county:
                hints["question_priority_signal"] = "high_exposure_priority"
                hints["rationale"] = (
                    f"Border county {county_risk.county_name} with {county_risk.risk_tier} risk tier; "
                    "prioritize cross-border exposure and travel history questions."
                )
            elif "healthcare_workers" in county_risk.high_risk_profiles:
                hints["question_priority_signal"] = "healthcare_amplification_priority"
                hints["rationale"] = (
                    f"County {county_risk.county_name} hosts major health facilities with potential "
                    "amplification risk; prioritize healthcare exposure and IPC-related questions."
                )
            else:
                hints["question_priority_signal"] = "high_exposure_priority"
                hints["rationale"] = (
                    f"County {county_risk.county_name} is {county_risk.risk_tier} risk; prioritize travel, "
                    "contact, and route-based exposure questions."
                )

        return hints

    def context_summary(self) -> str:
        active = ", ".join(self.context.active_outbreak_districts) or "None"
        neighboring = ", ".join(self.context.neighboring_outbreak_districts) or "None"
        alerts = ", ".join(self.context.cross_border_alerts) or "None"
        return (
            f"District in focus: {self.context.district}\n"
            f"Active outbreak districts: {active}\n"
            f"Neighboring outbreak districts: {neighboring}\n"
            f"Cross-border alerts: {alerts}\n"
            f"Recent confirmed cases: {self.context.recent_confirmed_cases}\n"
            f"Community transmission: {'Yes' if self.context.community_transmission else 'No'}"
        )