"""Psychometric scales: YAML schema, loading and deterministic scoring (issue #3).

Scale content (item wording, answer options, severity bands, flags) lives in `config/scales/*.yaml`
(ADR-008) and is approved by the team psychiatrist. This module only defines the schema and the
scoring algorithm, so thresholds can be changed without touching code.

Scoring is fully deterministic (ADR-007): the LLM never computes scores or decides severity.
"""

from __future__ import annotations

import operator
from collections.abc import Callable
from functools import lru_cache
from pathlib import Path
from typing import Literal

import yaml
from pydantic import BaseModel, ConfigDict, Field, model_validator

from app.config import get_settings

Sex = Literal["m", "f"]
Op = Literal[">=", ">", "<=", "<", "=="]

_OPS: dict[str, Callable[[int, int], bool]] = {
    ">=": operator.ge,
    ">": operator.gt,
    "<=": operator.le,
    "<": operator.lt,
    "==": operator.eq,
}


class _Strict(BaseModel):
    # Typos in YAML keys must fail loudly instead of silently dropping a threshold.
    model_config = ConfigDict(extra="forbid", frozen=True)


class Option(_Strict):
    value: int
    label: str


class Item(_Strict):
    id: str
    text: str
    # Empty list = inherit scale-level `options` (filled in by ScaleDef validation).
    options: list[Option] = Field(default_factory=list)
    # True = item must be asked with buttons only, never free text through the LLM (e.g. PHQ-9 item 9).
    buttons_only: bool = False
    # Reverse-keyed item: scored as (max_option + min_option - value). None of the current scales use it.
    reverse: bool = False

    @property
    def min_value(self) -> int:
        return min(o.value for o in self.options)

    @property
    def max_value(self) -> int:
        return max(o.value for o in self.options)


class Range(_Strict):
    min: int
    max: int


class SeverityLevel(_Strict):
    code: str
    label: str


class Band(SeverityLevel):
    """Severity band over the total score, inclusive on both ends."""

    min: int
    max: int


class Condition(_Strict):
    """Predicate over answers. Exactly one target must be set: item, sum_of, total, severity_in or any_of.

    Examples (YAML):
        {item: phq9_9, op: ">=", value: 1}            # single item
        {sum_of: [phq4_1, phq4_2], op: ">=", value: 3} # subscale (GAD-2 inside PHQ-4)
        {total: true, op: ">=", value: 33}            # total score
        {severity_in: [positive]}                     # computed severity code
        {any_of: [{item: q4, op: ">=", value: 1}, ...]}
    """

    item: str | None = None
    sum_of: list[str] | None = None
    total: bool = False
    severity_in: list[str] | None = None
    any_of: list[Condition] | None = None
    op: Op = ">="
    value: int | None = None

    @model_validator(mode="after")
    def _one_target(self) -> Condition:
        targets = [
            self.item is not None,
            self.sum_of is not None,
            self.total,
            self.severity_in is not None,
            self.any_of is not None,
        ]
        if sum(targets) != 1:
            raise ValueError("condition must have exactly one of: item, sum_of, total, severity_in, any_of")
        needs_value = self.item is not None or self.sum_of is not None or self.total
        if needs_value and self.value is None:
            raise ValueError("item/sum_of/total conditions require `value`")
        return self

    def item_refs(self) -> set[str]:
        refs = set(self.sum_of or [])
        if self.item is not None:
            refs.add(self.item)
        for sub in self.any_of or []:
            refs |= sub.item_refs()
        return refs


class Flag(_Strict):
    code: str
    when: Condition


class SeverityRule(_Strict):
    """Rule for `scoring: rules` scales (C-SSRS): the first matching rule gives the severity."""

    code: str
    label: str
    when: Condition


class ScaleDef(_Strict):
    id: str
    title: str
    recall_period: str
    instruction: str | None = None
    license_note: str | None = None
    translation_note: str | None = None
    notes: str | None = None
    # `sum`: severity from total-score bands. `rules`: severity from ordered item rules (C-SSRS triage).
    scoring: Literal["sum", "rules"] = "sum"
    # Default answer options for items that do not define their own.
    options: list[Option] = Field(default_factory=list)
    items: list[Item]
    range: Range
    severity: list[Band] = Field(default_factory=list)
    # Sex-specific bands (AUDIT-C). When sex is unknown, `severity` is used.
    by_sex: dict[Sex, list[Band]] | None = None
    severity_rules: list[SeverityRule] = Field(default_factory=list)
    default_severity: SeverityLevel | None = None
    flags: list[Flag] = Field(default_factory=list)

    @model_validator(mode="before")
    @classmethod
    def _inherit_options(cls, data: object) -> object:
        if isinstance(data, dict) and data.get("options"):
            for item in data.get("items", []):
                if isinstance(item, dict) and not item.get("options"):
                    item["options"] = data["options"]
        return data

    @model_validator(mode="after")
    def _validate(self) -> ScaleDef:
        ids = [i.id for i in self.items]
        if len(ids) != len(set(ids)):
            raise ValueError(f"{self.id}: duplicate item ids")
        for item in self.items:
            if not item.options:
                raise ValueError(f"{self.id}.{item.id}: no answer options")
            values = [o.value for o in item.options]
            if len(values) != len(set(values)):
                raise ValueError(f"{self.id}.{item.id}: duplicate option values")

        # Declared range must match what the items can actually produce.
        lo = sum(i.min_value for i in self.items)
        hi = sum(i.max_value for i in self.items)
        if (self.range.min, self.range.max) != (lo, hi):
            raise ValueError(f"{self.id}: range {self.range.min}..{self.range.max} != items {lo}..{hi}")

        if self.scoring == "sum":
            if not self.severity:
                raise ValueError(f"{self.id}: sum scoring requires `severity` bands")
            self._check_bands(self.severity, "severity")
            for sex, bands in (self.by_sex or {}).items():
                self._check_bands(bands, f"by_sex.{sex}")
        else:
            if not self.severity_rules or self.default_severity is None:
                raise ValueError(f"{self.id}: rules scoring requires `severity_rules` and `default_severity`")

        known = set(ids)
        conditions = [f.when for f in self.flags] + [r.when for r in self.severity_rules]
        for cond in conditions:
            unknown = cond.item_refs() - known
            if unknown:
                raise ValueError(f"{self.id}: condition references unknown items {sorted(unknown)}")
        return self

    def _check_bands(self, bands: list[Band], where: str) -> None:
        # Bands must tile the whole range without gaps or overlaps, otherwise some totals get no severity.
        expected = self.range.min
        for band in sorted(bands, key=lambda b: b.min):
            if band.min != expected or band.max < band.min:
                raise ValueError(f"{self.id}.{where}: bands must be contiguous from {self.range.min}")
            expected = band.max + 1
        if expected != self.range.max + 1:
            raise ValueError(f"{self.id}.{where}: bands must end at {self.range.max}")

    @property
    def item_ids(self) -> list[str]:
        return [i.id for i in self.items]

    def item(self, item_id: str) -> Item:
        for item in self.items:
            if item.id == item_id:
                return item
        raise KeyError(f"{self.id}: unknown item {item_id!r}")


class ScaleResult(BaseModel):
    model_config = ConfigDict(frozen=True)

    scale: str
    total: int
    max_total: int
    severity_code: str
    severity_label: str
    complete: bool
    answered: int
    flags: list[str] = Field(default_factory=list)


# ---------------------------------------------------------------- loading


def _scales_dir(config_dir: Path | None) -> Path:
    return (config_dir or get_settings().config_dir) / "scales"


@lru_cache
def _load_dir(scales_dir: Path) -> dict[str, ScaleDef]:
    scales: dict[str, ScaleDef] = {}
    for path in sorted(scales_dir.glob("*.yaml")):
        with path.open(encoding="utf-8") as fh:
            scale = ScaleDef.model_validate(yaml.safe_load(fh))
        if scale.id != path.stem:
            raise ValueError(f"{path.name}: scale id {scale.id!r} must match file name")
        scales[scale.id] = scale
    return scales


def load_all(config_dir: Path | None = None) -> dict[str, ScaleDef]:
    """All scales from `<config_dir>/scales/*.yaml`, keyed by id. Cached; see `clear_cache()`."""
    return dict(_load_dir(_scales_dir(config_dir)))


def load_scale(scale_id: str, config_dir: Path | None = None) -> ScaleDef:
    scales = _load_dir(_scales_dir(config_dir))
    if scale_id not in scales:
        raise KeyError(f"unknown scale {scale_id!r}")
    return scales[scale_id]


def clear_cache() -> None:
    """Drop cached scale definitions (after editing YAML at runtime or in tests)."""
    _load_dir.cache_clear()


# ---------------------------------------------------------------- scoring


def _validate_answers(scale: ScaleDef, answers: dict[str, int]) -> dict[str, int]:
    """Check every answer is a known item with an allowed option value; return item -> scored value."""
    scored: dict[str, int] = {}
    for item_id, value in answers.items():
        try:
            item = scale.item(item_id)
        except KeyError as exc:
            raise ValueError(str(exc)) from exc
        # bool is a subclass of int; reject it so `True` is not silently accepted as 1.
        if isinstance(value, bool) or not isinstance(value, int):
            raise ValueError(f"{scale.id}.{item_id}: answer must be int, got {type(value).__name__}")
        if value not in {o.value for o in item.options}:
            raise ValueError(f"{scale.id}.{item_id}: invalid value {value}")
        scored[item_id] = item.max_value + item.min_value - value if item.reverse else value
    return scored


def _eval(cond: Condition, scored: dict[str, int], total: int, severity_code: str | None) -> bool:
    """Evaluate a condition on answered items only.

    Unanswered items never satisfy an item condition. Sums over partially answered item sets use
    the answered part, which is a lower bound of the true sum: a positive `>=` result is still valid,
    a negative one is not conclusive (the caller sees `complete=False`).
    """
    if cond.any_of is not None:
        return any(_eval(sub, scored, total, severity_code) for sub in cond.any_of)
    if cond.severity_in is not None:
        return severity_code in cond.severity_in
    assert cond.value is not None  # guaranteed by Condition validation
    if cond.item is not None:
        if cond.item not in scored:
            return False
        lhs = scored[cond.item]
    elif cond.sum_of is not None:
        if not any(i in scored for i in cond.sum_of):
            return False
        lhs = sum(scored.get(i, 0) for i in cond.sum_of)
    else:
        lhs = total
    return _OPS[cond.op](lhs, cond.value)


def _band_for(bands: list[Band], total: int) -> Band:
    for band in bands:
        if band.min <= total <= band.max:
            return band
    raise ValueError(f"total {total} outside severity bands")  # unreachable: bands tile the range


def score(scale_id: str, answers: dict[str, int], sex: Sex | None = None) -> ScaleResult:
    """Score a questionnaire deterministically.

    Algorithm:
    1. Validate: every answer must belong to a known item and be one of its option values (ValueError).
       Reverse-keyed items are flipped.
    2. Total = plain sum of item values (standard scoring for PHQ-9, GAD-7, PCL-5, AUDIT-C, PC-PTSD-5,
       PHQ-4; C-SSRS yes=1/no=0, total is just the count of "yes").
    3. Severity:
       - `scoring: sum` — the band that contains the total. AUDIT-C uses `by_sex[sex]` bands
         (positive >=4 men, >=3 women); with unknown sex the default bands use the lower, more
         sensitive female cut-off so a risky pattern is not missed.
       - `scoring: rules` (C-SSRS screen) — ordered rules, first match wins (high > moderate > low),
         otherwise `default_severity` (none). Triage: q1-2 ideation -> low, q3 method -> moderate,
         q4 intent / q5 plan / q6 behaviour in past 3 months -> high.
    4. Flags — every flag whose condition holds (e.g. PHQ-9 item 9 >= 1 -> `phq9_item9_positive`,
       which downstream triggers C-SSRS and the crisis branch).

    Incomplete answers: we do NOT impute missing items (no prorating / mean substitution). Prorating
    can invent a score the patient never reported, and for safety items (PHQ-9 item 9, C-SSRS) a
    missing answer must not look like "no risk". So the total is a lower bound over answered items,
    `complete=False` tells consumers the severity may be underestimated, and flags/rules still fire
    on what was answered (a positive item 9 is acted upon even in a half-finished PHQ-9).
    """
    scale = load_scale(scale_id)
    scored = _validate_answers(scale, answers)
    total = sum(scored.values())

    if scale.scoring == "rules":
        assert scale.default_severity is not None  # guaranteed by ScaleDef validation
        matched = next((r for r in scale.severity_rules if _eval(r.when, scored, total, None)), None)
        level = matched or scale.default_severity
        code, label = level.code, level.label
    else:
        # sex=None or a scale without sex-specific bands -> default bands.
        bands = (scale.by_sex or {}).get(sex, scale.severity) if sex else scale.severity
        band = _band_for(bands, total)
        code, label = band.code, band.label

    flags = [f.code for f in scale.flags if _eval(f.when, scored, total, code)]
    return ScaleResult(
        scale=scale.id,
        total=total,
        max_total=scale.range.max,
        severity_code=code,
        severity_label=label,
        complete=len(scored) == len(scale.items),
        answered=len(scored),
        flags=flags,
    )


def delta(prev: ScaleResult, cur: ScaleResult) -> int:
    """Change of total score between two measurements of the same scale (positive = worse).

    Used by rules like "PHQ-9 grew by >= 5" or "PCL-5 grew by >= 10". Both results must be complete:
    a lower-bound total from a partial questionnaire would produce a fake improvement or deterioration.
    """
    if prev.scale != cur.scale:
        raise ValueError(f"cannot compare different scales: {prev.scale} vs {cur.scale}")
    if not (prev.complete and cur.complete):
        raise ValueError("delta requires complete results")
    return cur.total - prev.total