tests/test_interview_machine.py

"""Interview state machine: phase flow, turn caps, question guarding, immutability."""
from __future__ import annotations

import dataclasses

import pytest

import interview
from interview import (
    MAX_TOTAL_TURNS,
    InterviewState,
    finalize,
    opening_question,
    step,
)


def _fake_question_model(question: str = "When did this start?", axis: str = "onset"):
    """call_model stub returning a fixed History-Taker question."""

    def call_model(messages, schema):
        return {
            "question": question,
            "axis": axis,
            "covered_axes": [],
            "covered_details": [],
        }

    return call_model


BENIGN_ANSWER = "My crown feels a bit high when I bite on that side."


def test_opening_question_is_deterministic_and_safe():
    q = opening_question()
    assert q.endswith("?")
    assert "diagnos" not in q.lower()


def test_first_step_advances_from_demographics():
    state = InterviewState()
    result = step(state, BENIGN_ANSWER, call_model=_fake_question_model())
    assert result.state.phase != "demographics"
    assert result.state.user_turns == (BENIGN_ANSWER,)
    assert result.assistant_message.endswith("?")
    assert not result.state.done


def test_state_is_immutable():
    state = InterviewState()
    with pytest.raises(dataclasses.FrozenInstanceError):
        state.phase = "goals"  # type: ignore[misc]


def test_step_returns_new_state_object():
    state = InterviewState()
    result = step(state, BENIGN_ANSWER, call_model=_fake_question_model())
    assert result.state is not state
    assert state.user_turns == ()  # original untouched


def test_total_turn_cap_always_terminates():
    state = InterviewState()
    answers = 0
    for _ in range(MAX_TOTAL_TURNS + 3):
        result = step(state, BENIGN_ANSWER, call_model=_fake_question_model())
        state = result.state
        answers += 1
        if state.done:
            break
    assert state.done
    assert answers <= MAX_TOTAL_TURNS + 1


def test_interview_completes_through_all_phases():
    """A benign interview must reach done with phases traversed in order."""
    state = InterviewState()
    seen_phases = [state.phase]
    for _ in range(MAX_TOTAL_TURNS + 3):
        result = step(state, BENIGN_ANSWER, call_model=_fake_question_model())
        state = result.state
        if state.phase not in seen_phases:
            seen_phases.append(state.phase)
        if state.done:
            break
    assert state.done
    assert not state.early_exit
    assert seen_phases.index("demographics") < seen_phases.index("chief_concern")
    assert seen_phases.index("chief_concern") < seen_phases.index("dental_qualifiers")
    assert seen_phases.index("dental_qualifiers") < seen_phases.index("odipara")
    assert seen_phases.index("odipara") < seen_phases.index("background")
    assert seen_phases.index("background") < seen_phases.index("red_flag_screen")
    assert seen_phases.index("red_flag_screen") < seen_phases.index("goals")


def test_unsafe_model_question_falls_back_to_bank():
    """A diagnosis-flavored question from the model must never reach the patient."""
    bad = _fake_question_model(
        question="Sounds like an abscess — is it swollen?", axis="provocation"
    )
    result = step(InterviewState(), BENIGN_ANSWER, call_model=bad)
    msg = result.assistant_message.lower()
    assert "abscess" not in msg
    assert result.assistant_message.endswith("?")


def test_model_failure_falls_back_to_bank():
    def broken(messages, schema):
        raise RuntimeError("modal cold")

    result = step(InterviewState(), BENIGN_ANSWER, call_model=broken)
    assert result.assistant_message.endswith("?")
    assert not result.state.done


def test_non_question_model_output_falls_back():
    flat = _fake_question_model(question="Tell me more about the pain.", axis="quality")
    result = step(InterviewState(), BENIGN_ANSWER, call_model=flat)
    assert result.assistant_message.endswith("?")


def test_odipara_axes_are_not_repeated():
    """The same ODIPARA axis must not be asked twice in one interview."""
    state = InterviewState()
    for _ in range(MAX_TOTAL_TURNS + 3):
        result = step(state, BENIGN_ANSWER, call_model=_fake_question_model())
        state = result.state
        if state.done:
            break
    duplicated = [a for a in state.axes_asked if state.axes_asked.count(a) > 1]
    assert duplicated == [], f"axes repeated: {duplicated}"
    assert set(interview.ODIPARA_AXES).issubset(state.axes_covered)


def test_model_marks_volunteered_odipara_axes_as_covered():
    state = InterviewState(
        phase="odipara",
        user_turns=("Sam, 35.", "Upper molar pain."),
        questions_asked=("What brought you in today?",),
        axes_asked=("chief_concern",),
    )

    def coverage_model(messages, schema):
        return {
            "question": "Since it began, is it getting better, worse, or staying the same?",
            "axis": "progression",
            "covered_axes": ["onset", "duration", "intensity"],
            "covered_details": [
                "treatment_or_trauma_context",
                "constant_or_episodic",
                "episode_or_lingering_duration",
                "day_or_night_pattern",
                "current_and_worst_intensity",
                "functional_impact",
            ],
        }

    result = step(
        state,
        "It began yesterday, lasts ten minutes, and reaches 7 out of 10.",
        call_model=coverage_model,
    )
    assert {"onset", "duration", "intensity"}.issubset(result.state.axes_covered)
    assert result.state.axes_asked[-1] == "progression"


def test_model_cannot_skip_an_axis_without_its_dental_details():
    state = InterviewState(
        phase="odipara",
        user_turns=("Sam, 35.", "Upper molar pain."),
        questions_asked=("What brought you in today?",),
        axes_asked=("chief_concern",),
    )

    def shallow_coverage_model(messages, schema):
        return {
            "question": "Since it began, is it getting better or worse?",
            "axis": "progression",
            "covered_axes": ["onset", "duration", "intensity"],
            "covered_details": [],
        }

    result = step(state, "It hurts.", call_model=shallow_coverage_model)
    assert not {"onset", "duration", "intensity"} & set(result.state.axes_covered)
    assert result.state.axes_asked[-1] == "progression"

    follow_up = step(result.state, "It is getting worse.", call_model=shallow_coverage_model)
    assert follow_up.state.axes_asked[-1] == "onset"
    assert follow_up.assistant_message == interview._QUESTION_BANK["onset"]


def test_model_cannot_leave_the_odipara_axis_contract():
    state = InterviewState(
        phase="odipara",
        user_turns=("Sam, 35.", "Upper molar pain."),
        questions_asked=("What brought you in today?",),
        axes_asked=("chief_concern",),
    )
    invalid = _fake_question_model(
        question="How would you describe the quality of the pain?",
        axis="quality",
    )
    result = step(state, "It hurts.", call_model=invalid)
    assert result.state.axes_asked[-1] == "onset"
    assert result.assistant_message == interview._QUESTION_BANK["onset"]


def test_fallback_reaches_both_red_flag_screens():
    def broken(messages, schema):
        raise RuntimeError("model unavailable")

    state = InterviewState()
    for _ in range(MAX_TOTAL_TURNS + 1):
        result = step(state, BENIGN_ANSWER, call_model=broken)
        state = result.state
        if state.done:
            break
    assert state.done
    assert "red_flag_infection" in state.axes_asked
    assert "red_flag_airway" in state.axes_asked
    assert set(interview.ODIPARA_AXES).issubset(state.axes_covered)
    assert "character_radiation" in state.axes_asked
    assert set(interview._QUALIFIER_DETAILS).issubset(state.details_covered)


def test_fallback_questions_cover_dental_specific_odipara_details():
    bank = interview._QUESTION_BANK
    assert "spread" in bank["character_radiation"]
    assert "dental work" in bank["onset"]
    assert "linger" in bank["duration"]
    assert "eating" in bank["intensity"]
    assert "biting down or release" in bank["aggravating"]
    assert "gum pimple" in bank["associated"]


def test_progress_summary_is_complete_only_when_interview_is_done():
    completed, total, stage = interview.progress_summary(InterviewState())
    assert completed == 0
    # Progress milestones are defined by the phase budgets, not the turn cap.
    assert total == sum(interview._PHASE_BUDGET.values())
    assert stage == "Getting acquainted"

    done = dataclasses.replace(InterviewState(), done=True)
    completed, total, stage = interview.progress_summary(done)
    assert completed == total
    assert stage == "Ready to build"


def test_step_after_done_is_a_noop():
    done_state = dataclasses.replace(InterviewState(), done=True)
    result = step(done_state, "anything", call_model=_fake_question_model())
    assert result.state.done
    assert result.state.user_turns == ()


def test_finalize_returns_validated_extracted_intake():
    def extractor(messages, schema):
        return {
            "chief_concern": "Crown feels high",
            "tooth_or_area": "Upper left molar",
            "recent_dental_work": "Crown placement",
            "symptom_duration": "Three weeks",
            "pain_score": 6,
            "biting_pain": True,
            "goals": "Understand whether the crown fit explains the pain",
        }

    state = InterviewState(user_turns=(BENIGN_ANSWER,))
    extracted = finalize(state, call_model=extractor)
    assert extracted.chief_concern == "Crown feels high"
    assert extracted.biting_pain is True
    assert extracted.pain_score == 6


def test_finalize_raises_on_invalid_payload():
    def extractor(messages, schema):
        return {"pain_score": 55}  # out of range

    with pytest.raises(Exception):
        finalize(InterviewState(user_turns=(BENIGN_ANSWER,)), call_model=extractor)


def test_story_so_far_joins_only_user_turns():
    state = InterviewState(user_turns=("First answer.", "Second answer."))
    story = interview.story_so_far(state)
    assert "First answer." in story
    assert "Second answer." in story