schema.py

from __future__ import annotations

import re
import unicodedata
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator


SourceType = Literal["free_text", "structured_input", "example_cache", "model_output"]
EscalationTier = Literal[
    "emergency_now",
    "urgent_medical",
    "same_day_urgent",
    "emergency_or_same_day",
    "urgent_dental",
    "dentist_discussion",
    "none",
]

BLOCKED_QUESTION_TERMS = (
    "you have",
    "you need a root canal",
    "failed root canal",
    "definitely",
)

_MODEL_CLAIM_PATTERNS = tuple(
    re.compile(pattern, re.IGNORECASE)
    for pattern in (
        r"\b(?:you|the patient)\s+(?:have|has|need|needs|must|should)\b",
        r"\b(?:diagnosis is|diagnosed with|definitely|confirmed|proven)\b",
        r"\b(?:is|are|was|were|could be|can be|might be|may be)\s+(?:consistent with|diagnostic of|indicative of|caused by|due to)\b",
        # Diagnosis asserted about a named / third-person subject ("Ahmed has
        # pulpitis", "presents with an abscess", "developed necrosis") — the
        # you/the-patient pattern misses these. Over-blocking a quoted prior
        # diagnosis is fail-safe: the field blanks to the deterministic fallback.
        # "reports/mentions/thinks/possible" are NOT here on purpose (patient history).
        r"\b(?:has|have|got|presents? with|presenting with|developed|suffers? from|showing signs of)\s+"
        r"(?:a|an|the|some)?\s*\w*(?:abscess|infection|pulpitis|periodontitis|gingivitis|necrosis|"
        r"cellulitis|cavity|caries|decay|fracture|cracked tooth|cyst|tumou?r|cancer|\w+itis)\b",
        # Treatment-need directive. The negative lookbehinds excuse HEDGED
        # patient-reported history ("the dentist mentioned she may need an
        # extraction") while still catching model-authored assertions ("you need
        # an extraction", "requires a root canal", bare "needs an extraction").
        # The guard is documented as model-output-only and must tolerate quoted
        # patient treatment history.
        r"\b(?<!may )(?<!might )(?:requires?|must undergo|should undergo|needs?|prescribe[ds]?"
        r"|start(?:s|ed)? (?:taking|on)|put on)\s+(?:an?\s+)?"
        r"(?:root canal|extraction|surgery|antibiotics?|crown replacement|implant|treatment"
        # Specific drugs as a treatment OBJECT (verb-gated, so the meds field's own
        # drug names — listed without a recommend/requires verb — are NOT blanked).
        r"|amoxicillin|amoxiclav|augmentin|clindamycin|metronidazole|flagyl|penicillin"
        r"|azithromycin|cephalexin|ibuprofen|brufen|naproxen|paracetamol|acetaminophen"
        r"|panadol|codeine|tramadol|cataflam|voltaren|diclofenac)\b",
        r"\b(?:prescribe|start taking|increase the dose|stop taking)\b",
        # Hedged diagnosis ("looks like an abscess", "appears to be pulpitis",
        # "sounds like an abscess" — added when the interview History-Taker guard
        # surfaced the gap).
        r"(?:looks?|appears?|seems?|sounds?)\s+(?:like|to be)\s+(?:a|an|the)?\s*"
        r"\w*(?:abscess|infection|pulpitis|cavity|decay|caries|fracture|cracked|cyst|"
        r"tumor|cancer|gingivitis|periodontitis|\w+itis)",
        r"(?:likely|probably|possibly|most likely)\s+(?:a|an|the)?\s*"
        r"(?:abscess|infection|pulpitis|cavity|fracture|cracked tooth|\w+itis)",
        # Bare "consistent with <diagnosis>" (no preceding is/are/was/were) — a
        # diagnostic claim the model must not author.
        r"consistent with\s+(?:a|an|the)?\s*"
        r"\w*(?:abscess|infection|pulpitis|cavity|decay|caries|fracture|cracked|cyst|"
        r"tumor|cancer|gingivitis|periodontitis|\w+itis)",
        # Soft treatment directives ("I'd recommend...", "it would be wise to get
        # that tooth pulled", "consider a crown replacement"). Word-bounded so
        # benign descriptive words ("suggestive of clenching", "suggestion") do
        # not trip the guard — that false positive blanked the whole model
        # contribution on live TMJ/bruxism stories (probe round 3).
        r"\b(?:recommend|advise|suggest)(?:s|ed|ing)?\b",
        r"would be wise to",
        r"consider (?:getting|seeing|a|an)",
        r"(?:get|have|having)\s+(?:that|the|your)?\s*"
        r"(?:tooth|crown|bridge|implant)?\s*(?:pulled|extracted|removed|replaced)",
        r"(?:root canal|extraction|crown|filling|implant)\s+would help",
    )
)

# Unicode format characters (category Cf) that can be injected between letters to
# defeat the substring/word-boundary guard, e.g. "You​have a failed root canal".
# NFKC normalization alone does NOT remove U+200B, so we handle them explicitly.
# Map each to a SPACE (not deletion) so injected chars between words restore the
# word boundary the spaced guard patterns rely on.
_GUARD_STRIP_CHARS = {ord(c): " " for c in ("​", "‌", "‍", "⁠", "")}


def _normalize_for_guard(text: str) -> str:
    """Normalize text for safety-guard matching only (not for storage).

    Folds curly quotes/apostrophes to ASCII and replaces zero-width / format
    characters with a single space so adversarial Unicode injected *between* words
    (e.g. "You​have") restores a natural word boundary instead of fusing the
    words into one token that the spaced guard patterns would miss. The text is
    NOT stored — this normalization is for matching only.
    """

    # Fold curly quotes/apostrophes to ASCII first so NFKC does not split them oddly.
    text = text.replace("’", "'").replace("‘", "'").replace("“", '"').replace("”", '"')
    text = unicodedata.normalize("NFKC", text)
    # Replace remaining Cf-category format chars (NFKC keeps U+200B et al.) with a
    # space. Using a space — not deletion — so "You​have" becomes "You have"
    # and still trips the guard, rather than the unmatchable token "Youhave".
    text = text.translate(_GUARD_STRIP_CHARS)
    text = "".join(" " if unicodedata.category(ch) == "Cf" else ch for ch in text)
    return text


def model_text_is_safe(text: str) -> bool:
    """Reject model-authored diagnosis or treatment assertions.

    This guard is intentionally applied to model output only. Patient-entered text
    may contain quoted diagnoses or treatment history and must remain representable
    in the deterministic fallback handoff.
    """

    normalized = _normalize_for_guard(text)
    return not any(pattern.search(normalized) for pattern in _MODEL_CLAIM_PATTERNS)


class EvidenceSpan(BaseModel):
    source: SourceType
    quote: str = Field(min_length=1, max_length=280)


class RedFlagFinding(BaseModel):
    rule_id: str
    title: str
    tier: EscalationTier
    patient_message: str
    clinician_question: str
    evidence: list[EvidenceSpan] = Field(default_factory=list)

    @field_validator("evidence")
    @classmethod
    def require_evidence(cls, value: list[EvidenceSpan]) -> list[EvidenceSpan]:
        if not value:
            raise ValueError("red flags must include at least one evidence span")
        return value


class PatientProfile(BaseModel):
    name: str = ""
    age: int | None = Field(default=None, ge=0, le=120)
    language: Literal["English", "Arabic", "Bilingual"] = "English"
    meds: str = ""
    allergies: str = ""
    goals: str = ""


class StructuredIntake(BaseModel):
    chief_concern: str = ""
    tooth_or_area: str = ""
    recent_dental_work: str = ""
    symptom_duration: str = ""
    pain_score: int = Field(default=0, ge=0, le=10)
    pain_prevents_sleep: bool = False
    biting_pain: bool = False
    hot_cold_sensitivity: bool = False
    swelling: bool = False
    rapidly_spreading_swelling: bool = False
    fever_or_unwell: bool = False
    breathing_or_swallowing_issue: bool = False
    limited_opening_or_locked_jaw: bool = False
    loose_crown_or_bridge: bool = False
    trauma_or_sudden_bite_change: bool = False
    numbness_or_neuro_symptoms: bool = False
    chest_pain_or_jaw_pain_with_exertion: bool = False
    jaw_pain_with_chewing_relieved_by_rest: bool = False
    vision_scalp_or_new_headache: bool = False
    gum_pimple_or_drainage: bool = False
    bruising_or_burning_after_root_canal: bool = False


class HandoffOutput(BaseModel):
    product_name: str = "Dental SOAP"
    artifact_label: str = "Dentist Visit Handoff"
    patient_name: str = ""
    patient_age: int | None = None
    chief_concern: str
    concise_summary: str
    timeline: list[str] = Field(default_factory=list)
    current_symptoms: list[str] = Field(default_factory=list)
    dental_history: list[str] = Field(default_factory=list)
    medical_safety_notes: list[str] = Field(default_factory=list)
    patient_goals: list[str] = Field(default_factory=list)
    dentist_questions: list[str] = Field(default_factory=list)
    after_visit_tracker: list[str] = Field(default_factory=list)
    # Deterministic "bring to the visit" checklist — built by rules from the
    # intake/profile, never authored by the model (not a ModelHandoffDraft field).
    bring_checklist: list[str] = Field(default_factory=list)
    evidence: list[EvidenceSpan] = Field(default_factory=list)
    red_flags: list[RedFlagFinding] = Field(default_factory=list)
    limitations: list[str] = Field(
        default_factory=lambda: [
            "This is not a diagnosis.",
            "This does not interpret X-rays, CBCT, photos, or scans.",
            "This organizes patient-reported history to bring to a licensed dentist.",
        ]
    )

    @field_validator("dentist_questions")
    @classmethod
    def block_assessment_plan_language(cls, value: list[str]) -> list[str]:
        lowered = " ".join(value).lower()
        if any(term in lowered for term in BLOCKED_QUESTION_TERMS):
            raise ValueError("questions must not contain diagnostic or treatment claims")
        return value


class ModelHandoffDraft(BaseModel):
    """The narrow, validated surface the LLM may contribute to a handoff."""

    model_config = ConfigDict(extra="ignore")

    chief_concern: str | None = Field(default=None, max_length=240)
    concise_summary: str | None = Field(default=None, max_length=1200)
    timeline: list[str] | None = Field(default=None, max_length=8)
    current_symptoms: list[str] | None = Field(default=None, max_length=8)
    dental_history: list[str] | None = Field(default=None, max_length=8)
    dentist_questions: list[str] | None = Field(default=None, max_length=8)

    @field_validator("chief_concern", "concise_summary")
    @classmethod
    def validate_model_text(cls, value: str | None) -> str | None:
        if value is None:
            return None
        cleaned = value.strip()
        if not cleaned:
            return None
        if not model_text_is_safe(cleaned):
            raise ValueError("model narrative contains diagnosis or treatment language")
        return cleaned

    @field_validator("timeline", "current_symptoms", "dental_history")
    @classmethod
    def validate_model_lists(cls, value: list[str] | None) -> list[str] | None:
        if value is None:
            return None
        cleaned: list[str] = []
        for item in value:
            if not isinstance(item, str):
                raise ValueError("model list entries must be strings")
            item = item.strip()
            if not item:
                continue
            if len(item) > 500:
                raise ValueError("model list entry is too long")
            if not model_text_is_safe(item):
                raise ValueError("model list entry contains diagnosis or treatment language")
            cleaned.append(item)
        return cleaned

    @field_validator("dentist_questions")
    @classmethod
    def validate_model_questions(cls, value: list[str] | None) -> list[str] | None:
        if value is None:
            return None
        cleaned: list[str] = []
        for item in value:
            if not isinstance(item, str):
                raise ValueError("model questions must be strings")
            item = item.strip()
            if not item:
                continue
            lowered = item.lower()
            if any(term in lowered for term in BLOCKED_QUESTION_TERMS):
                raise ValueError("model question contains diagnostic or treatment claims")
            if len(item) > 300:
                raise ValueError("model question is too long")
            if not item.endswith("?"):
                raise ValueError("model questions must be phrased as questions")
            cleaned.append(item)
        return cleaned


class CaseInput(BaseModel):
    profile: PatientProfile
    intake: StructuredIntake
    story: str = Field(min_length=1)