from __future__ import annotations

import json
import os
from pathlib import Path

from .models import CountyRisk, EpidemiologicalContext


DEFAULT_CONTEXT_PATH = Path(__file__).resolve().parent.parent / "data" / "epi_context.sample.json"
DEFAULT_COUNTY_RISK_PATH = Path(__file__).resolve().parent.parent / "data" / "kenya_county_risk.json"


def load_epidemiological_context(path: str | None = None) -> EpidemiologicalContext:
    context_path = Path(path) if path else Path(os.getenv("EVD_CONTEXT_PATH", str(DEFAULT_CONTEXT_PATH)))
    if not context_path.exists():
        context = EpidemiologicalContext()
    else:
        with context_path.open("r", encoding="utf-8") as handle:
            payload = json.load(handle)

        context = EpidemiologicalContext.model_validate(payload)

    county_risk_path = Path(os.getenv("EVD_COUNTY_RISK_PATH", str(DEFAULT_COUNTY_RISK_PATH)))
    if county_risk_path.exists():
        context.county_risks = load_kenya_county_risks(county_risk_path)

    return context


def load_kenya_county_risks(path: Path) -> dict[str, CountyRisk]:
    """Load Kenya county-level risk intelligence from JSON."""
    with path.open("r", encoding="utf-8") as handle:
        payload = json.load(handle)

    county_risks: dict[str, CountyRisk] = {}
    for county_name, county_data in payload.get("counties", {}).items():
        county_risks[county_name] = CountyRisk.model_validate(county_data)

    return county_risks