src/endopath/schema.py

"""CAP/ICCR endometrium checklist schema and case/field state model.

Implements docs/prd.md section 6 (the nine-to-fourteen-field CAP/ICCR endometrium
protocol, FIGO staging-edition tagging) and section 8.2 (case and field state
model). This module is the single source of truth for field definitions --
extraction, backend, and frontend code should all import from here rather
than re-deriving the checklist.
"""

from __future__ import annotations

from datetime import date
from enum import Enum
from typing import Any, Generic, Optional, TypeVar

from pydantic import BaseModel, Field, model_validator

# The field list and its value enums live in fields.py, which is stdlib-only so
# scripts/modal_precompute.py can share it (issue #37). Re-exported here because
# `from endopath.schema import Histotype` is used throughout.
from endopath.fields import (
    Histotype,
    LvsiStatus,
    MolecularClassification,
)

__all__ = [
    "Case",
    "CaseStatus",
    "ChecklistField",
    "EndometrialChecklist",
    "EvidenceSpan",
    "FieldStatus",
    "FigoStagingEdition",
    "Histotype",
    "LvsiStatus",
    "MolecularClassification",
    "Report",
    "ReportType",
    "Specimen",
    "UndesignatedReportError",
]

T = TypeVar("T")


class FieldStatus(str, Enum):
    """docs/prd.md section 8.2: "Confirmed, Needs Review... or Flagged, meaning
    actively searched for and not found, distinct from simply absent."

    NOT_APPLICABLE is a state in its own right: a gate removed the field's
    meaning here (FIGO grade on a non-endometrioid histotype, molecular
    classification pre-2023). Distinct from CONFIRMED so a rule-suppressed field
    never masquerades as one a reviewer accepted (issue #43).
    """

    CONFIRMED = "confirmed"
    NEEDS_REVIEW = "needs_review"
    FLAGGED = "flagged"
    NOT_APPLICABLE = "not_applicable"


class CaseStatus(str, Enum):
    """docs/prd.md section 8.2: "Queued to Processing to Ready for Review to In
    Review to Confirmed to Exported, with a separate Flagged state."
    """

    QUEUED = "queued"
    PROCESSING = "processing"
    READY_FOR_REVIEW = "ready_for_review"
    IN_REVIEW = "in_review"
    CONFIRMED = "confirmed"
    EXPORTED = "exported"
    FLAGGED = "flagged"


_VALID_CASE_TRANSITIONS: dict[CaseStatus, set[CaseStatus]] = {
    CaseStatus.QUEUED: {CaseStatus.PROCESSING, CaseStatus.FLAGGED},
    CaseStatus.PROCESSING: {CaseStatus.READY_FOR_REVIEW, CaseStatus.FLAGGED},
    CaseStatus.READY_FOR_REVIEW: {CaseStatus.IN_REVIEW, CaseStatus.FLAGGED},
    CaseStatus.IN_REVIEW: {CaseStatus.CONFIRMED, CaseStatus.FLAGGED},
    CaseStatus.CONFIRMED: {CaseStatus.EXPORTED, CaseStatus.FLAGGED},
    CaseStatus.EXPORTED: set(),
    # Flagged cases (amended-report / template-clustering misses, per section
    # 7.2) return to whichever stage a human resolution sends them back to.
    CaseStatus.FLAGGED: {
        CaseStatus.QUEUED,
        CaseStatus.PROCESSING,
        CaseStatus.READY_FOR_REVIEW,
        CaseStatus.IN_REVIEW,
    },
}


class FigoStagingEdition(str, Enum):
    """docs/prd.md section 6: FIGO staging changed editions in 2023, adding
    molecular classification. Every stage/grade field needs this tag at
    ingestion or cross-year exports silently mix incompatible values.
    """

    FIGO_2009 = "figo_2009"
    FIGO_2023 = "figo_2023"

    @classmethod
    def for_report_date(cls, report_date: Optional[date]) -> "FigoStagingEdition":
        # The 2023 FIGO revision was published mid-2023; a report dated
        # before that cannot have used it regardless of calendar year.
        if report_date is not None and report_date >= date(2023, 6, 1):
            return cls.FIGO_2023
        return cls.FIGO_2009


# docs/prd.md section 6: FIGO grade "applies only to endometrioid and mucinous carcinoma."
# That is wrong. CAP v5.1 Note D grades endometrioid *only* (docs/taxonomy.md
# section 1), and `Histotype.MUCINOUS` is not in CAP's type list at all. Both the
# rule and the enum member are corrected in #36; this set is reproduced verbatim
# here so #37 stays a pure move with no behavior change.
_HISTOTYPES_WITH_FIGO_GRADE = {Histotype.ENDOMETRIOID, Histotype.MUCINOUS}


class EvidenceSpan(BaseModel):
    """Pointer back to whatever justified a field's value.

    docs/prd.md section 8.3: "No field is ever shown as a bare value. It's shown
    with what justified it."
    """

    quote: str
    char_start: Optional[int] = None
    char_end: Optional[int] = None
    page_number: Optional[int] = None
    # Which pass produced this evidence -- ties into the multi-modal
    # consensus and escalating not-found-search issues.
    source: str = "text"
    # Which document and specimen this value came from (issue #110). Optional and
    # unpopulated by today's extraction pipeline; a value without them still resolves to a
    # case's single report/specimen through Case's compatibility properties. Populating these
    # for every extracted value is #111's job, once the observation becomes the stored atom.
    report_id: Optional[str] = None
    specimen_id: Optional[str] = None


class ChecklistField(BaseModel, Generic[T]):
    """One CAP/ICCR checklist field, evidence-first per section 8.3."""

    value: Optional[T] = None
    confidence: Optional[float] = Field(default=None, ge=0.0, le=1.0)
    evidence: Optional[EvidenceSpan] = None
    status: FieldStatus = FieldStatus.NEEDS_REVIEW

    # A field can be genuinely inapplicable (e.g. FIGO grade for serous
    # carcinoma, molecular classification pre-2023) rather than missing.
    # Section 7.3: "Silence is never the response to a miss" -- inapplicable
    # fields must carry a stated reason, not just an empty value.
    applicable: bool = True
    not_applicable_reason: Optional[str] = None

    # Escalation steps actually tried before a field is marked Flagged.
    # Section 6: "not found is an earned result" -- this is the record that
    # earns it.
    search_log: list[str] = Field(default_factory=list)

    @model_validator(mode="after")
    def _not_applicable_is_internally_consistent(self) -> "ChecklistField[T]":
        if not self.applicable and not self.not_applicable_reason:
            raise ValueError(
                "A field marked not applicable must carry not_applicable_reason "
                "(section 8.3: silence is never the response to a miss)."
            )
        return self

    @classmethod
    def not_applicable(cls, reason: str) -> "ChecklistField[T]":
        return cls(
            value=None,
            status=FieldStatus.NOT_APPLICABLE,
            applicable=False,
            not_applicable_reason=reason,
        )

    def set_applicability(self, applicable: bool, reason: Optional[str] = None) -> None:
        """Gate this field on or off without destroying its drafted value.

        Applicability is a reviewed value's context, not the value itself
        (issue #43). Suppressing a field records `not_applicable` as its state
        while KEEPING `value`, `confidence`, `evidence`, and `search_log` as the
        suppressed draft, so the confirmation screen can show what was drafted
        and a reviewer can catch a wrong gate (invariant 2). Un-suppressing
        returns that preserved draft and sends the field back to NEEDS_REVIEW:
        the gate is itself a drafted value that can change, and the returned
        field was never actually reviewed while suppressed.

        Idempotent: a field already in the target applicability only has its
        reason refreshed (the gate condition can reword it, e.g. a different
        suppressing histotype), so calling this twice with the same input
        changes nothing.
        """
        if applicable:
            if not self.applicable:
                self.applicable = True
                self.not_applicable_reason = None
                self.status = FieldStatus.NEEDS_REVIEW
            return
        if reason is None:
            raise ValueError(
                "Suppressing a field requires a reason (section 8.3: silence is "
                "never the response to a miss)."
            )
        if self.applicable:
            self.applicable = False
            self.status = FieldStatus.NOT_APPLICABLE
        self.not_applicable_reason = reason

    def confirm(self) -> None:
        """The fast path (section 8.3): accept the drafted value as-is."""
        self.status = FieldStatus.CONFIRMED

    def edit(self, value: Optional[T], evidence: Optional[EvidenceSpan] = None) -> None:
        """The exception path (section 8.3): a reviewer overrides the value."""
        self.value = value
        if evidence is not None:
            self.evidence = evidence
        self.status = FieldStatus.CONFIRMED

    def flag_not_found(self, search_log: list[str]) -> None:
        """Only call once the full escalating search (LLM, embedding,
        combined) has actually been run and come back empty -- section 6.
        """
        if len(search_log) < 3:
            raise ValueError(
                "flag_not_found requires a search_log recording all three "
                "escalation passes (LLM, embedding, combined); a shorter "
                "log means the search wasn't actually earned (section 7)."
            )
        self.value = None
        self.status = FieldStatus.FLAGGED
        self.search_log = search_log


class EndometrialChecklist(BaseModel):
    """The nine-to-fourteen-field CAP/ICCR endometrium protocol (docs/prd.md
    section 6), not the original seven-field draft.
    """

    histologic_type: ChecklistField[Histotype] = Field(default_factory=ChecklistField)
    histologic_grade: ChecklistField[str] = Field(default_factory=ChecklistField)
    myometrial_invasion_percent: ChecklistField[float] = Field(default_factory=ChecklistField)
    cervical_stromal_invasion: ChecklistField[bool] = Field(default_factory=ChecklistField)
    lymphovascular_space_invasion: ChecklistField[LvsiStatus] = Field(default_factory=ChecklistField)
    regional_lymph_node_status: ChecklistField[str] = Field(default_factory=ChecklistField)
    pathologic_stage: ChecklistField[str] = Field(default_factory=ChecklistField)
    molecular_classification: ChecklistField[MolecularClassification] = Field(
        default_factory=ChecklistField
    )

    def apply_applicability_rules(self, staging_edition: FigoStagingEdition) -> None:
        """Recompute which gated fields apply. Call after histologic_type or
        staging_edition changes (extracted or reviewer-edited).

        Total and reversible (issue #43): every gated field is recomputed on
        every call, so a gate that opens (a reviewer corrects a wrong histotype)
        returns the previously suppressed field with its drafted value intact,
        just as a gate that closes suppresses it. Suppression never deletes the
        draft. The method is idempotent and order-independent.
        """
        histotype_value = self.histologic_type.value
        grade_suppressed = (
            histotype_value is not None and histotype_value not in _HISTOTYPES_WITH_FIGO_GRADE
        )
        self.histologic_grade.set_applicability(
            applicable=not grade_suppressed,
            reason=(
                f"FIGO grade only applies to endometrioid and mucinous carcinoma; "
                # A value edited through the API arrives as a raw string until the
                # case is next read back and coerced, so read .value defensively.
                f"histotype is {getattr(histotype_value, 'value', histotype_value)}."
                if grade_suppressed
                else None
            ),
        )

        molecular_suppressed = staging_edition != FigoStagingEdition.FIGO_2023
        self.molecular_classification.set_applicability(
            applicable=not molecular_suppressed,
            reason=(
                "Report predates the 2023 FIGO revision that added molecular "
                "classification."
                if molecular_suppressed
                else None
            ),
        )

    def all_fields(self) -> dict[str, ChecklistField]:
        return {name: getattr(self, name) for name in self.__class__.model_fields}


class ReportType(str, Enum):
    """What kind of document a report is (issue #110). A typical case carries an outpatient
    biopsy report and a hysterectomy report as two independently significant, non-superseding
    documents, not one report occasionally amended: neither is "the" report by default.
    """

    BIOPSY = "biopsy"
    HYSTERECTOMY = "hysterectomy"
    AMENDMENT = "amendment"
    EXTERNAL_CONSULT = "external_consult"
    ANCILLARY = "ancillary"


class UndesignatedReportError(ValueError):
    """A case carries more than one report and none, or more than one, is designated
    authoritative. Issue #110: the system never infers this from report type, recency, or list
    position -- a reviewer must designate exactly one before the case's checklist resolves.

    Carries `suggested_report_id` (issue #127): a report-type-precedence suggestion the reviewer
    can confirm, never applied automatically. `None` when no single report is unambiguously
    first.
    """

    def __init__(self, message: str, *, suggested_report_id: Optional[str] = None) -> None:
        super().__init__(message)
        self.suggested_report_id = suggested_report_id


class Specimen(BaseModel):
    """One specimen within a report: what a checklist's observations describe (issue #110).
    The checklist lives here, not on Report or Case, because CAP's checklist is inherently
    specimen-scoped (a hysterectomy specimen's findings).
    """

    specimen_id: str
    procedure: Optional[str] = None
    checklist: EndometrialChecklist = Field(default_factory=EndometrialChecklist)


class Report(BaseModel):
    """One pathology document assembled into a case: a biopsy, hysterectomy, amendment,
    external consult, or ancillary/molecular report (issue #110).
    """

    report_id: str
    report_type: ReportType = ReportType.HYSTERECTOMY
    institution: Optional[str] = None
    report_date: Optional[date] = None
    staging_edition: FigoStagingEdition = FigoStagingEdition.FIGO_2009
    # Amendment lineage only: an addendum names the report it supersedes. Two concurrent
    # report types (e.g. biopsy and hysterectomy) are never in a supersedes relationship.
    supersedes_report_id: Optional[str] = None
    # None: undesignated. A reviewer sets True/False explicitly; the system never infers this
    # from report type, recency, or position (issue #110).
    is_authoritative: Optional[bool] = None
    specimens: list[Specimen] = Field(default_factory=list)
    # The recognized text this report was read from (issue #39), so a stored EvidenceSpan's
    # char_start/char_end resolves to its snippet without re-reading the source corpus file.
    report_text: Optional[str] = None

    @model_validator(mode="after")
    def _supersession_is_amendment_only(self) -> "Report":
        if self.supersedes_report_id is not None and self.report_type is not ReportType.AMENDMENT:
            raise ValueError(
                "supersedes_report_id is only valid on an AMENDMENT report; two concurrent "
                "report types (e.g. biopsy and hysterectomy) are never in a supersedes "
                "relationship with each other (issue #110)"
            )
        return self


def _default_report(case_barcode: str) -> Report:
    """The single auto-vivified report/specimen a legacy single-report case resolves to
    (issue #110). Deterministic ids so re-persisting the same case doesn't duplicate rows."""
    return Report(
        report_id=f"{case_barcode}-report-1",
        specimens=[Specimen(specimen_id=f"{case_barcode}-specimen-1")],
    )


# Highest to lowest (issue #127): "the latest amendment supersedes," then "the hysterectomy
# report supersedes the biopsy for a resection finding" -- today's checklist is entirely
# resection-scoped, so that reduces cleanly to a case-wide tier order. A consult or ancillary
# report is never the primary diagnostic authority, so it ranks lowest.
_AUTHORITY_PRECEDENCE: tuple[ReportType, ...] = (
    ReportType.AMENDMENT,
    ReportType.HYSTERECTOMY,
    ReportType.BIOPSY,
    ReportType.EXTERNAL_CONSULT,
    ReportType.ANCILLARY,
)


def _suggest_authoritative(reports: list[Report]) -> Optional[Report]:
    """A suggested authoritative report the reviewer can confirm, never applied automatically
    (issue #127). Un-superseded reports (nothing else names them via `supersedes_report_id`,
    transitively true by construction: a superseding report is itself un-superseded unless a
    further amendment names it) are ranked by `_AUTHORITY_PRECEDENCE`; the sole report in the
    highest non-empty tier is the suggestion. A tie within a tier, or a case with no reports,
    returns no suggestion -- a genuine ambiguity is not the tool's call to make. This is a plain
    set-membership pass, not a walk of the supersession chain, so a broken or cyclic
    `supersedes_report_id` (bad data, not a state the schema forbids) can't loop it: it can only
    ever leave a report out of, or in, the un-superseded set.
    """
    superseded_ids = {r.supersedes_report_id for r in reports if r.supersedes_report_id}
    current = [r for r in reports if r.report_id not in superseded_ids]
    for report_type in _AUTHORITY_PRECEDENCE:
        candidates = [r for r in current if r.report_type is report_type]
        if len(candidates) == 1:
            return candidates[0]
        if len(candidates) > 1:
            return None
    return None


class Case(BaseModel):
    case_barcode: str
    status: CaseStatus = CaseStatus.QUEUED
    reports: list[Report] = Field(default_factory=list)

    @model_validator(mode="before")
    @classmethod
    def _normalize_legacy_constructor(cls, data: Any) -> Any:
        """Back-compat (issue #110): ~20 call sites and ~100 test assertions construct
        `Case(case_barcode=..., report_date=..., status=..., checklist=...)`, the pre-grain
        shape. When a caller doesn't pass `reports` directly, fold those scalar kwargs into a
        single default report/specimen so every existing call site, and `case.checklist` /
        `case.report_date` / `case.staging_edition`, keep working unchanged. `mode="before"`
        (not a custom `__init__`) so this also covers `Case.model_validate(...)`.
        """
        if not isinstance(data, dict) or "reports" in data:
            return data
        report_date = data.pop("report_date", None)
        # This runs before pydantic's own field coercion, unlike the old mode="after"
        # validator it replaces, so a raw ISO string (storage.get_case reads report_date
        # straight off a TEXT column) needs parsing before for_report_date can compare it.
        if isinstance(report_date, str):
            report_date = date.fromisoformat(report_date)
        staging_edition = data.pop("staging_edition", None)
        checklist = data.pop("checklist", None)
        case_barcode = data.get("case_barcode")
        data["reports"] = [
            Report(
                report_id=f"{case_barcode}-report-1",
                report_date=report_date,
                staging_edition=(
                    staging_edition
                    if staging_edition is not None
                    else FigoStagingEdition.for_report_date(report_date)
                ),
                specimens=[
                    Specimen(
                        specimen_id=f"{case_barcode}-specimen-1",
                        checklist=checklist if checklist is not None else EndometrialChecklist(),
                    )
                ],
            )
        ]
        return data

    def _resolved_report(self) -> Report:
        """The report `.checklist` / `.report_date` / `.staging_edition` read and write
        through. Never picks by position or type (issue #110): a single report resolves
        trivially; more than one requires a reviewer's `is_authoritative` designation."""
        if not self.reports:
            self.reports = [_default_report(self.case_barcode)]
        if len(self.reports) == 1:
            return self.reports[0]
        authoritative = [r for r in self.reports if r.is_authoritative is True]
        if len(authoritative) == 1:
            return authoritative[0]
        suggestion = _suggest_authoritative(self.reports)
        raise UndesignatedReportError(
            f"case {self.case_barcode} carries {len(self.reports)} reports and "
            f"{len(authoritative)} are designated authoritative; a reviewer must designate "
            "exactly one before .checklist / .report_date / .staging_edition can resolve "
            "(issue #110: no positional or type default"
            + (
                f"; {suggestion.report_id} is suggested by report-type precedence, issue #127)"
                if suggestion is not None
                else "; no suggestion, the candidates tie)"
            ),
            suggested_report_id=suggestion.report_id if suggestion is not None else None,
        )

    @property
    def checklist(self) -> EndometrialChecklist:
        report = self._resolved_report()
        if not report.specimens:
            report.specimens = [Specimen(specimen_id=f"{self.case_barcode}-specimen-1")]
        return report.specimens[0].checklist

    @property
    def report_date(self) -> Optional[date]:
        return self._resolved_report().report_date

    @report_date.setter
    def report_date(self, value: Optional[date]) -> None:
        report = self._resolved_report()
        report.report_date = value
        report.staging_edition = FigoStagingEdition.for_report_date(value)

    @property
    def staging_edition(self) -> FigoStagingEdition:
        return self._resolved_report().staging_edition

    @staging_edition.setter
    def staging_edition(self, value: FigoStagingEdition) -> None:
        self._resolved_report().staging_edition = value

    def transition_to(self, new_status: CaseStatus) -> None:
        allowed = _VALID_CASE_TRANSITIONS.get(self.status, set())
        if new_status not in allowed:
            raise ValueError(
                f"Invalid case state transition: {self.status.value} -> {new_status.value}"
            )
        self.status = new_status

    def is_ready_for_export(self) -> bool:
        # A field is settled when a reviewer confirmed it, or when a gate coded
        # its absence (NOT_APPLICABLE, a state in its own right). Both are named
        # explicitly rather than reading `not field.applicable`, so a field
        # un-suppressed by a corrected gate (now applicable and NEEDS_REVIEW)
        # correctly blocks export again (issue #43).
        settled = {FieldStatus.CONFIRMED, FieldStatus.NOT_APPLICABLE}
        return all(
            field.status in settled for field in self.checklist.all_fields().values()
        )