STT-Assistant/tests/test_triage.py

"""Tests for the triage analyzer (todos 3, 9, 10, 11). Pure, hardware-free."""

import json
import logging
from pathlib import Path
from types import SimpleNamespace

import pytest

from stt_daemon import events
from stt_daemon.assistant import triage
from stt_daemon.assistant.service import AssistantService
from stt_daemon.assistant.triage import TriageAnalyzer, load_questions
from stt_daemon.config import Config
from stt_daemon.tui.state import ViewState

EXAMPLE_PATH = str(Path(__file__).resolve().parent.parent / "questions.example.json")


def _cfg(path=EXAMPLE_PATH, **overrides):
    base = {
        "questions_file": path,
        "state_char_budget": 1200,
        "confidence_field": "answer_confidence",
        "escalate_threshold": 0.6,
        "choice_min_confidence": 0.85,
    }
    base.update(overrides)
    return SimpleNamespace(**base)


def _write_questions(tmp_path, questions):
    path = tmp_path / "questions.json"
    path.write_text(json.dumps(questions), encoding="utf-8")
    return str(path)


def _analyzer(tmp_path, questions=None, **overrides):
    if questions is None:
        return TriageAnalyzer(_cfg(**overrides))
    return TriageAnalyzer(_cfg(_write_questions(tmp_path, questions), **overrides))


# --- todo 3: load_questions -------------------------------------------------


def test_load_example_order_preserved():
    raw = json.loads(Path(EXAMPLE_PATH).read_text(encoding="utf-8"))
    questions = load_questions(EXAMPLE_PATH)
    assert questions is not None
    assert questions
    assert list(questions.keys()) == list(raw.keys())
    for definition in questions.values():
        assert definition["type"] in {"choice", "score", "noul"}
        assert isinstance(definition["instructions"], str)
        assert definition["instructions"].strip()


def test_load_missing_file_returns_none_with_warning(caplog):
    with caplog.at_level(logging.WARNING, logger="stt_daemon.assistant.triage"):
        assert load_questions("/nonexistent/questions.json") is None
    assert caplog.records, "expected a warning for the missing file"


def test_load_bad_json_returns_none_with_warning(tmp_path, caplog):
    path = tmp_path / "questions.json"
    path.write_text("{not valid json", encoding="utf-8")
    with caplog.at_level(logging.WARNING, logger="stt_daemon.assistant.triage"):
        assert load_questions(str(path)) is None
    assert caplog.records


@pytest.mark.parametrize(
    "bad_entry",
    [
        {"type": "mystery", "instructions": "do something"},  # unknown type
        {"type": "score", "instructions": "rate it", "criteria": {"0": "low"}},  # score dict
        {  # noul criteria with a bad key
            "type": "noul",
            "instructions": "holds?",
            "criteria": {"true": "yes", "maybe": "hmm"},
        },
        "just a string",  # non-object entry
        {"type": "choice", "instructions": "pick one", "criteria": []},  # empty choice
        {"type": "choice", "instructions": "pick one", "criteria": {}},  # empty choice dict
        {"type": "choice", "criteria": ["a", "b"]},  # missing instructions
        {"type": "score", "instructions": "rate it", "criteria": ["a", None]},  # None level
        {"type": "noul", "instructions": "holds?", "criteria": ["true"]},  # noul list
    ],
)
def test_load_malformed_returns_none_with_warning(tmp_path, caplog, bad_entry):
    path = _write_questions(
        tmp_path, {"good": {"type": "noul", "instructions": "ok?"}, "bad": bad_entry}
    )
    with caplog.at_level(logging.WARNING, logger="stt_daemon.assistant.triage"):
        assert load_questions(str(path)) is None
    assert caplog.records


# --- todo 9: normalize ------------------------------------------------------


def _score_questions():
    return {
        "severity": {
            "type": "score",
            "instructions": "How severe is the issue?",
            "criteria": ["none", "low", "high", "critical"],
        },
        "category": {
            "type": "choice",
            "instructions": "What is this about?",
            "criteria": {"support": "help needed", "sales": "buying"},
        },
        "churn_risk": {"type": "noul", "instructions": "User may leave."},
    }


def test_normalize_choice_exact(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    payload = {
        "model": "multilingual",
        "answers": {
            "category": {
                "type": "choice",
                "choice": "support",
                "confidence": 0.7,
                "answer_confidence": 0.4,
            }
        },
        "usage": {"input_tokens": 123, "output_tokens": 0},
    }
    result = analyzer.normalize(payload)
    assert result["order"] == ["category"]
    assert result["items"]["category"] == {
        "kind": "choice",
        "choice": "support",
        "confidence": 0.7,
        "answer_confidence": 0.4,
        "gate": 0.4,
        "flagged": True,
    }
    assert result["input_tokens"] == 123
    assert result["escalated"] is True


def test_normalize_score_with_legend(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    payload = {
        "answers": {
            "severity": {
                "type": "score",
                "score": 2.6,
                "legend": {"0": "none", "1": "low", "2": "high", "3": "critical"},
                "confidence": 0.5,
                "answer_confidence": 0.6,
            }
        },
        "usage": {"input_tokens": 7, "output_tokens": 0},
    }
    item = analyzer.normalize(payload)["items"]["severity"]
    assert item["kind"] == "score"
    assert item["expected"] == 2.6
    assert item["level"] == 3
    assert item["label"] == "critical"
    assert item["confidence"] == 0.5
    assert item["answer_confidence"] == 0.6
    assert item["gate"] == 0.6
    assert item["flagged"] is True


def test_normalize_score_legend_fallback(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    payload = {
        "answers": {
            "severity": {"type": "score", "score": 1.2, "confidence": 0.9, "answer_confidence": 0.9}
        }
    }
    item = analyzer.normalize(payload)["items"]["severity"]
    assert item["level"] == 1
    assert item["label"] == "low"  # falls back to criteria[level]
    assert item["flagged"] is False  # 1.2 < (4 - 1) * 0.5


def test_normalize_noul_and_bare_payload(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    bare = {
        "churn_risk": {"type": "noul", "noul": 0.9, "confidence": 0.9, "answer_confidence": 0.85}
    }
    result = analyzer.normalize(bare)
    assert result["order"] == ["churn_risk"]
    assert result["items"]["churn_risk"] == {
        "kind": "noul",
        "p": 0.9,
        "confidence": 0.9,
        "answer_confidence": 0.85,
        "gate": 0.85,
        "flagged": True,
    }
    assert result["input_tokens"] == 0


def test_normalize_skips_unknown_qid_and_malformed(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    payload = {
        "answers": {
            "ghost": {"type": "noul", "noul": 0.99, "confidence": 0.9, "answer_confidence": 0.9},
            "category": {"type": "choice", "choice": 42},  # malformed: no confidences
            "churn_risk": "nonsense",  # malformed: not a dict
        },
        "usage": {"input_tokens": 3, "output_tokens": 0},
    }
    result = analyzer.normalize(payload)
    assert result["order"] == []
    assert result["items"] == {}
    assert result["escalated"] is False


def test_normalize_ignores_side_channel_fields(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    payload = {
        "answers": {
            "category": {
                "type": "choice",
                "choice": "sales",
                "confidence": 0.95,
                "answer_confidence": 0.9,
                "action": {"act" + "_probability": 0.99},
            }
        },
        "usage": {"input_tokens": 1, "output_tokens": 0},
    }
    result = analyzer.normalize(payload)
    assert result["items"]["category"]["flagged"] is False
    assert "act" + "_probability" not in json.dumps(result)
    source = Path(triage.__file__).read_text(encoding="utf-8")
    assert "act" + "_probability" not in source


def test_normalize_gating_field_selectable(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions(), confidence_field="confidence")
    payload = {
        "answers": {
            "category": {
                "type": "choice",
                "choice": "sales",
                "confidence": 0.5,
                "answer_confidence": 0.99,
            }
        }
    }
    item = analyzer.normalize(payload)["items"]["category"]
    assert item["gate"] == 0.5  # gates on confidence, keeps both
    assert item["answer_confidence"] == 0.99
    assert item["flagged"] is True


# --- todo 9: condensation ---------------------------------------------------


def test_condense_keeps_newest_chronological_no_numerals(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions(), state_char_budget=40)
    window = ["alpha first", "beta second", "gamma third newest here"]
    state = analyzer.build_state(window)
    assert "gamma third newest here" in state
    assert "alpha first" not in state  # oldest dropped under the budget
    lines = state.splitlines()
    assert all(line.startswith("- ") for line in lines)
    assert not any(ch.isdigit() for ch in state)  # no numerals injected
    # chronological restore: kept lines stay in original order
    kept = [u for u in window if u in state]
    assert [line[2:] for line in lines] == kept


def test_condense_empty_window(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    assert analyzer.build_state([]) == ""


# --- todo 10: escalation + extraction ---------------------------------------


def _signals(choice_gate=0.9, noul_p=0.1, score_expected=0.2):
    return {
        "order": ["category", "churn_risk", "severity"],
        "items": {
            "category": {
                "kind": "choice",
                "choice": "sales",
                "confidence": choice_gate,
                "answer_confidence": choice_gate,
                "gate": choice_gate,
                "flagged": choice_gate < 0.85,
            },
            "churn_risk": {
                "kind": "noul",
                "p": noul_p,
                "confidence": noul_p,
                "answer_confidence": noul_p,
                "gate": noul_p,
                "flagged": noul_p >= 0.6,
            },
            "severity": {
                "kind": "score",
                "expected": score_expected,
                "level": 0,
                "label": "none",
                "confidence": 0.9,
                "answer_confidence": 0.9,
                "gate": 0.9,
                "flagged": score_expected >= 1.5,
            },
        },
        "escalated": False,
        "input_tokens": 0,
    }


def test_escalation_arms_fire_independently(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    assert analyzer.should_escalate(_signals(choice_gate=0.5)) is True  # choice arm
    assert analyzer.should_escalate(_signals(noul_p=0.8)) is True  # noul arm
    assert analyzer.should_escalate(_signals(score_expected=2.0)) is True  # score arm
    assert analyzer.should_escalate(_signals()) is False


def test_normalize_escalated_matches_should_escalate(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    payload = {
        "answers": {
            "category": {
                "type": "choice",
                "choice": "sales",
                "confidence": 0.9,
                "answer_confidence": 0.9,
            },
            "churn_risk": {
                "type": "noul",
                "noul": 0.1,
                "confidence": 0.9,
                "answer_confidence": 0.9,
            },
            "severity": {
                "type": "score",
                "score": 0.1,
                "confidence": 0.9,
                "answer_confidence": 0.9,
            },
        }
    }
    result = analyzer.normalize(payload)
    assert result["escalated"] == analyzer.should_escalate(result) is False


def test_build_extraction_prompt_words_not_dump(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    prompt = analyzer.build_extraction_prompt(["please cancel my account"], _signals(noul_p=0.8))
    assert "please cancel my account" in prompt
    assert "churn_risk" in prompt  # flagged signal named in words
    assert "category" in prompt  # unflagged signals included too
    assert "severity" in prompt
    assert "Signals:" in prompt
    assert "Flagged signals:" not in prompt
    assert "{'kind'" not in prompt and '{"kind"' not in prompt  # no raw dict dump
    for key in ("summary", "opportunities", "doubts", "action_items", "risks"):
        assert key in prompt


def test_build_extraction_prompt_state_only_when_no_signals(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    for signals in (None, {}, {"order": [], "items": {}}):
        prompt = analyzer.build_extraction_prompt(["hello there"], signals)
        assert "hello there" in prompt
        assert "Signals:" not in prompt
        assert "Flagged signals:" not in prompt
        assert "summary" in prompt  # instruction keys still present
        assert "opportunities" in prompt
    prompt = analyzer.build_extraction_prompt(["hello there"], None)
    assert "Signals:" not in prompt
    assert "\nnone" not in prompt


def _extraction_doc(**overrides):
    doc = {
        "summary": "user wants to cancel",
        "opportunities": ["winback offer"],
        "doubts": ["reason unclear"],
        "action_items": [{"what": "call back", "why": "save account"}],
        "risks": ["churn"],
    }
    doc.update(overrides)
    return json.dumps(doc)


def test_parse_extraction_valid_fenced(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    text = "```json\n" + _extraction_doc() + "\n```"
    parsed = analyzer.parse_extraction(text)
    assert parsed is not None
    assert parsed["summary"] == "user wants to cancel"
    assert parsed["action_items"] == [{"what": "call back", "why": "save account"}]


def test_parse_extraction_invalid_returns_none(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    assert analyzer.parse_extraction("not json at all {{{") is None
    assert analyzer.parse_extraction(json.dumps({"summary": "only"})) is None  # bad shape
    assert analyzer.parse_extraction(None) is None
    assert analyzer.parse_extraction("") is None
    assert analyzer.parse_extraction("x" * 30000) is None  # oversized


def test_parse_extraction_caps_lists_and_strings(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    parsed = analyzer.parse_extraction(
        _extraction_doc(
            opportunities=["o%d" % i for i in range(9)],
            doubts=["d%d" % i for i in range(7)],
            risks=["r%d" % i for i in range(6)],
            summary="s" * 2000,
        )
    )
    assert parsed is not None
    assert len(parsed["opportunities"]) == 5
    assert len(parsed["doubts"]) == 5
    assert len(parsed["risks"]) == 5
    assert len(parsed["summary"]) <= 500


# --- todo 11: cached state() ------------------------------------------------


def test_state_initial_loading_with_good_questions(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    assert analyzer.state() == ("loading", "")


def test_state_initial_no_questions_with_bad_path(tmp_path):
    missing = str(tmp_path / "missing.json")
    analyzer = TriageAnalyzer(_cfg(missing))
    state, reason = analyzer.state()
    assert state == "no_questions"
    assert missing in reason


def test_state_vocabulary_has_no_extra_member():
    assert set(triage.TRIAGE_REASONS) == {
        "ok",
        "no_questions",
        "loading",
        "laya_missing",
        "disabled",
        "error",
    }
    assert "service" + "_down" not in triage.TRIAGE_REASONS


def test_refresh_transitions(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())
    analyzer.refresh("ok")
    assert analyzer.state() == ("ok", "")
    analyzer.refresh("laya_missing", "run: uv sync --extra triage")
    assert analyzer.state() == ("laya_missing", "run: uv sync --extra triage")
    analyzer.refresh("error", "boom")
    assert analyzer.state() == ("error", "boom")


def test_refresh_unknown_state_ignored(tmp_path, caplog):
    analyzer = _analyzer(tmp_path, _score_questions())
    analyzer.refresh("ok")
    with caplog.at_level(logging.WARNING, logger="stt_daemon.assistant.triage"):
        analyzer.refresh("bogus_state")
    assert analyzer.state() == ("ok", "")
    assert caplog.records


def test_no_questions_state_cannot_be_overridden(tmp_path):
    analyzer = TriageAnalyzer(_cfg(str(tmp_path / "missing.json")))
    assert analyzer.state()[0] == "no_questions"
    analyzer.refresh("ok")
    analyzer.refresh("error", "boom")
    assert analyzer.state()[0] == "no_questions"


def test_state_does_not_touch_lane(tmp_path):
    analyzer = _analyzer(tmp_path, _score_questions())

    class StrictLane:
        def __init__(self):
            self.calls = []

        def __getattr__(self, name):
            def record(*args, **kwargs):
                self.calls.append(name)

            return record

    lane = StrictLane()
    for _ in range(3):
        analyzer.state()
    analyzer.refresh("ok")
    analyzer.state()
    assert lane.calls == []


# --- F3: end-to-end triage flow on fakes ------------------------------------

EXTRACT_JSON = (
    '{"summary":"s","opportunities":[],"doubts":[],"action_items":[],"risks":[]}'
)


class _Clock:
    """Injected clock: no sleeps, fully deterministic."""

    def __init__(self, start=1000.0):
        self.t = float(start)

    def __call__(self):
        return self.t

    def advance(self, dt):
        self.t += float(dt)


class _FakeSubscriber:
    def events(self):
        return iter(())

    def close(self):
        pass


class _FakeLane:
    """Duck-typed LayaLane: canned-ready, records submits, counts drops."""

    def __init__(self):
        self.submits: list[tuple] = []
        self.drops = 0

    def ready(self) -> bool:
        return True

    def state(self):
        return ("ok", "")

    def submit(self, state, questions, lang, generation):
        self.submits.append((state, questions, lang, generation))
        return True

    def drop_pending(self):
        self.drops += 1
        return 1


def _answers(noul_p=0.1):
    """A full laya-shaped answers payload for the example questions."""
    answers = {
        "category": {
            "type": "choice",
            "choice": "sales",
            "confidence": 0.9,
            "answer_confidence": 0.9,
        },
        "urgency": {
            "type": "choice",
            "choice": "low",
            "confidence": 0.9,
            "answer_confidence": 0.9,
        },
    }
    for qid in (
        "frustration",
        "churn_risk",
        "needs_action",
        "needs_clarification",
        "contains_idea",
        "business_opportunity",
        "positive_sentiment",
    ):
        answers[qid] = {
            "type": "noul",
            "noul": noul_p,
            "confidence": 0.9,
            "answer_confidence": 0.9,
        }
    return {"model": "multilingual", "answers": answers, "usage": {"input_tokens": 42}}


def test_end_to_end_triage_flow():
    """signals -> content -> cleared across service, analyzer, lane, and ViewState.

    The lane result is wired by calling `service._on_triage_result(gen, payload)`
    directly (the runner does exactly this in production); the extraction job is
    drained with `job.run()` so no worker thread is started.
    """
    cfg = Config()
    cfg.assistant.turn_gap_s = 0.5
    cfg.triage.enabled = True
    cfg.triage.questions_file = EXAMPLE_PATH
    clock = _Clock(1000.0)
    lane = _FakeLane()
    emitted: list[dict] = []

    def emit(type_, **fields):
        emitted.append({"type": type_, **fields})

    analyzer = TriageAnalyzer(cfg.triage)
    service = AssistantService(
        cfg,
        _FakeSubscriber(),
        emit=emit,
        time_fn=clock,
        chat_fn=lambda text: "reply",
        triage=analyzer,
        laya_lane=lane,
        triage_extract_fn=lambda prompt: EXTRACT_JSON,
    )

    def triage_events():
        return [e for e in emitted if e["type"] == "triage"]

    def extract_jobs():
        return [j for j in service._inference._queue if j.name == "triage-extract"]

    # 1. settled turn -> exactly one lane submit, lang + generation recorded
    assert service.toggle_chat() is True
    service.handle_event(
        {"type": events.TRANSCRIPT, "text": "hola, necesito ayuda", "lang": "es"}
    )
    clock.advance(1.0)
    service.tick()
    assert len(lane.submits) == 1
    state_text, questions, lang, gen1 = lane.submits[0]
    assert "hola, necesito ayuda" in state_text
    assert questions is analyzer.questions
    assert lang == "es"
    assert gen1 == 0

    # 2. non-escalating result -> signals plus one state+signals extraction job
    service._on_triage_result(gen1, _answers(noul_p=0.1))
    assert [e["stage"] for e in triage_events()] == ["signals"]
    assert triage_events()[0]["triage"]["escalated"] is False
    assert service.last_triage is not None
    assert service.last_triage["escalated"] is False
    assert len(extract_jobs()) == 1

    # 3. escalating turn -> second submit, then a second extract job
    service.handle_event(
        {"type": events.TRANSCRIPT, "text": "quiero cancelar mi suscripcion", "lang": "es"}
    )
    clock.advance(1.0)
    service.tick()
    assert len(lane.submits) == 2
    gen2 = lane.submits[1][3]
    service._on_triage_result(gen2, _answers(noul_p=0.9))
    assert [e["stage"] for e in triage_events()] == ["signals", "signals"]
    assert triage_events()[1]["triage"]["escalated"] is True
    jobs = extract_jobs()
    assert len(jobs) == 2

    # 4. run the extraction jobs directly (deterministic, no worker thread)
    for job in jobs:
        job.run()
    content_events = [e for e in triage_events() if e["stage"] == "content"]
    assert len(content_events) == 2
    merged = content_events[1]["triage"]
    assert merged["content"]["summary"] == "s"
    assert merged["escalated"] is True  # signals preserved under the content
    assert merged["order"] == triage_events()[1]["triage"]["order"]

    # 5. drive the emitted stream through a fresh ViewState
    view = ViewState()
    for event in emitted:
        view.ingest(event)
    analyzer.refresh("ok")  # the lane callback does this in production
    view.ingest({"type": "status", **service.status()})
    assert view.triage is not None
    assert view.triage["content"]["summary"] == "s"
    lines = view.triage_lines()
    assert lines
    assert any("summary" in line and "s" in line for line in lines)

    # 6. clear -> cleared event empties the panel and bumps the epoch
    baseline_epoch = view.triage_epoch
    already = len(emitted)
    service.clear()
    for event in emitted[already:]:
        view.ingest(event)
    assert view.triage is None
    assert view.triage_epoch > baseline_epoch

    # a stale result after clear must not re-populate the panel
    service._on_triage_result(gen2, _answers(noul_p=0.9))
    assert [e["stage"] for e in triage_events()] == [
        "signals",
        "signals",
        "content",
        "content",
        "cleared",
    ]
    assert view.triage is None
    assert lane.drops == 1

    # 7. disabled control: no triage/lane -> status disabled, clear emits nothing
    cfg2 = Config()
    clock2 = _Clock(2000.0)
    emitted2: list[dict] = []

    def emit2(type_, **fields):
        emitted2.append({"type": type_, **fields})

    plain = AssistantService(cfg2, _FakeSubscriber(), emit=emit2, time_fn=clock2)
    assert plain.status()["triage_state"] == "disabled"
    assert plain.clear() is True
    assert [e for e in emitted2 if e["type"] == "triage"] == []