STT-Assistant/tests/test_assistant_service.py

"""Tests for the assistant service (Phase 2) using synthetic event streams."""

import itertools
import json
import re
import threading
import time
from typing import Any

from stt_daemon import events
from stt_daemon.assistant.inference import InferenceQueue
from stt_daemon.assistant.service import EXECUTE_BUSY, AssistantService
from stt_daemon.assistant.triage_dump import TriageDumper
from stt_daemon.config import Config

_sequence = itertools.count(1)


class Clock:
    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


def _service():
    cfg = Config()
    clock = Clock()
    emitted: list[dict] = []

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

    service = AssistantService(cfg, FakeSubscriber(), emit=emit, time_fn=clock)
    return service, emitted, clock


def _ev(type_, clock, **fields):
    event = {"v": 1, "type": type_, "seq": next(_sequence), "ts": clock()}
    event.update(fields)
    return event


def _transcript(clock, text, ts=None):
    return {
        "v": 1,
        "type": events.TRANSCRIPT,
        "seq": next(_sequence),
        "ts": clock() if ts is None else ts,
        "text": text,
    }


def test_ready_requires_matching_segment_done():
    service, _, clock = _service()

    service.handle_event(_ev(events.SPEECH_START, clock))
    assert service.ready() is False

    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    assert service.ready() is False  # awaiting segment 1

    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))
    assert service.ready() is True


def test_stale_segment_done_does_not_satisfy_later_turn():
    service, _, clock = _service()

    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    service.handle_event(_ev(events.SPEECH_START, clock))
    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=2))

    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))
    assert service.ready() is False  # still awaiting segment 2

    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=2, emitted=True))
    assert service.ready() is True


def test_gated_segment_done_clears_readiness():
    service, _, clock = _service()

    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=3))
    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=3, emitted=False))

    assert service.ready() is True


def test_speech_stop_without_seg_id_clears_readiness():
    """A discarded turn closes with speech_stop but no seg_id; readiness must recover."""
    service, _, clock = _service()

    service.handle_event(_ev(events.SPEECH_START, clock))
    assert service.ready() is False

    service.handle_event(_ev(events.SPEECH_STOP, clock))  # discarded: no seg_id
    assert service.speech_open is False
    assert service.awaited_seg_id is None
    assert service.ready() is True


def test_status_resets_stale_speech_state_on_reconnect():
    service, _, clock = _service()

    service.handle_event(_ev(events.SPEECH_START, clock))
    assert service.speech_open is True

    service.handle_event(_ev(events.STATUS, clock, capture="x", segments_done=0))
    assert service.speech_open is False
    assert service.awaited_seg_id is None
    assert service.ready() is True


def test_command_buffer_uses_arm_timestamp():
    service, _, clock = _service()
    clock.advance(10)
    assert service.arm() is True
    arm_ts = service.arm_ts
    assert arm_ts is not None

    # Replayed transcript from before arming is ignored.
    service.handle_event(_transcript(clock, "old", ts=arm_ts - 5))
    assert service.command_text == ""

    service.handle_event(_transcript(clock, "what time is it"))
    assert service.command_text == "what time is it"


def test_mutual_exclusion_between_command_and_chat():
    service, _, _ = _service()

    assert service.toggle_chat() is True
    assert service.arm() is False  # cannot arm in chat

    assert service.toggle_chat() is True  # back to idle
    assert service.arm() is True
    assert service.toggle_chat() is False  # cannot chat while armed


def test_arm_timeout_disarms_and_emits_error():
    service, emitted, clock = _service()
    service.arm()

    clock.advance(service.cfg.assistant.command_arm_timeout_s + 0.1)
    service.tick()

    assert service.modes.mode.value == "idle"
    assert service.command_text == ""
    assert any(e["type"] == "error" and e["code"] == "arm_timeout" for e in emitted)


def test_long_speech_does_not_disarm():
    service, emitted, clock = _service()
    service.arm()
    service.handle_event(_ev(events.SPEECH_START, clock))

    clock.advance(service.cfg.assistant.command_arm_timeout_s + 30.0)
    service.tick()

    assert service.modes.mode.value == "command_armed"
    assert not any(e["type"] == "error" for e in emitted)


def test_captured_command_survives_arm_timeout():
    service, emitted, clock = _service()
    service.arm()
    service.handle_event(_transcript(clock, "what time is it"))
    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))

    clock.advance(service.cfg.assistant.command_arm_timeout_s + 30.0)
    service.tick()

    assert service.modes.mode.value == "command_armed"
    assert service.command_text == "what time is it"
    assert not any(e["type"] == "error" for e in emitted)


def test_stt_stall_cancels_armed_command():
    service, emitted, clock = _service()
    service.handle_event(_ev(events.HEARTBEAT, clock, worker_stall_s=0.0))
    service.arm()

    clock.advance(service.cfg.assistant.stall_timeout_ms / 1000.0 + 0.5)
    service.tick()

    assert service.modes.mode.value == "idle"
    assert any(e["type"] == "error" and e["code"] == "stt_stall" for e in emitted)


def test_worker_stall_marks_stt_degraded():
    service, _, clock = _service()
    long_stall = service.cfg.assistant.worker_stall_ms / 1000.0 + 1.0

    service.handle_event(_ev(events.HEARTBEAT, clock, worker_stall_s=long_stall))
    assert service.stt_ok() is False
    assert service.ready() is False


def test_chat_idle_timeout_exits_chat():
    service, _, clock = _service()
    assert service.toggle_chat() is True

    clock.advance(service.cfg.assistant.chat_idle_timeout_s + 0.1)
    service.tick()

    assert service.modes.mode.value == "idle"


def test_chat_exit_phrase_exits_chat():
    service, _, clock = _service()
    assert service.toggle_chat() is True

    service.handle_event(_transcript(clock, "That's all."))

    assert service.modes.mode.value == "idle"
    assert service.chat_history == []


def test_viewer_detached_exits_chat_and_forgets_history():
    service, _, clock = _service()
    assert service.toggle_chat() is True
    service.handle_event(_transcript(clock, "hello"))

    assert service.viewer_detached() is True
    assert service.modes.mode.value == "idle"
    assert service.chat_history == []


def test_viewer_detached_is_noop_outside_chat():
    service, _, _ = _service()
    assert service.viewer_detached() is False
    assert service.modes.mode.value == "idle"


def test_exit_chat_invokes_chat_reset():
    cfg = Config()
    resets: list[int] = []
    service = AssistantService(
        cfg, FakeSubscriber(), time_fn=Clock(), chat_reset=lambda: resets.append(1)
    )
    service.toggle_chat()

    service.viewer_detached()

    assert resets == [1]


def test_chat_accumulates_turns():
    service, _, clock = _service()
    service.toggle_chat()

    for text in ("hello there", "how are you"):
        service.handle_event(_transcript(clock, text))

    assert service.chat_history == ["Unknown: hello there", "Unknown: how are you"]


def test_execute_requires_readiness():
    service, _, clock = _service()
    service.arm()
    service.handle_event(_transcript(clock, "do it"))

    service.handle_event(_ev(events.SPEECH_START, clock))
    assert service.execute() is None  # speech still open
    assert service.modes.mode.value == "command_armed"

    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    assert service.execute() is None  # awaiting segment 1

    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))
    assert service.execute() == "do it"
    assert service.modes.mode.value == "executing"

    assert service.finish() is True
    assert service.modes.mode.value == "idle"


def test_execute_reports_busy_when_queue_rejects():
    service, emitted, clock = _service()

    # A zero-capacity lane rejects the priority command job on submit.
    service._inference = InferenceQueue(size=0)
    service.executor = object()
    service.arm()
    service.handle_event(_transcript(clock, "do it"))
    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))

    assert service.execute() is EXECUTE_BUSY
    assert service.modes.mode.value == "idle"
    assert not any(e.get("code") == "busy" for e in emitted)


def test_turn_ended_requires_quiet_gap():
    service, _, clock = _service()
    service.handle_event(_ev(events.SPEECH_STOP, clock, seg_id=1))
    service.handle_event(_ev(events.SEGMENT_DONE, clock, seg_id=1, emitted=True))

    assert service.ready() is True
    assert service.turn_ended() is False

    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    assert service.turn_ended() is True


def test_emits_mode_and_command_text():
    service, emitted, clock = _service()
    service.arm()
    assert {"type": "mode", "mode": "command_armed"} in emitted

    service.handle_event(_transcript(clock, "hi"))
    assert any(e["type"] == "command_text" and e["text"] == "hi" for e in emitted)


def test_cancel_clears_command_buffer():
    service, _, clock = _service()
    service.arm()
    service.handle_event(_transcript(clock, "hi"))
    assert service.cancel() is True
    assert service.modes.mode.value == "idle"
    assert service.command_text == ""
    assert service.arm_ts is None


def test_pending_chat_is_not_lost_across_threads():
    """Transcripts appended on the reader thread while the ticker dispatches
    must all reach the model: the join-and-clear in _dispatch_chat is atomic."""
    cfg = Config()
    cfg.assistant.turn_gap_s = 0.01
    dispatched: list[str] = []

    def emit(type_, **fields):
        if type_ == "chat_message" and fields.get("role") == "user":
            dispatched.append(fields["text"])

    service = AssistantService(cfg, FakeSubscriber(), emit=emit, chat_fn=lambda text: "ok")
    service.toggle_chat()

    total = 200
    producer = threading.Thread(
        target=lambda: [
            service.handle_event(
                {
                    "v": 1,
                    "type": events.TRANSCRIPT,
                    "seq": i,
                    "ts": time.time(),
                    "text": f"w{i}",
                }
            )
            for i in range(total)
        ],
        daemon=True,
    )
    producer.start()

    deadline = time.time() + 10.0
    while time.time() < deadline:
        service.tick()
        words = sum(len(text.split()) for text in dispatched)
        if not producer.is_alive() and words == 2 * total:
            break
        time.sleep(0.002)
    producer.join(timeout=5)

    assert sum(len(text.split()) for text in dispatched) == 2 * total


def test_chat_notify_emits_compaction_note():
    service, emitted, _ = _service()

    service.chat_notify("compacted", turns=4)

    assert {
        "type": "chat_message",
        "role": "note",
        "text": "Compacted 4 earlier turns into a summary.",
    } in emitted


def test_chat_notify_emits_truncated_error():
    service, emitted, _ = _service()

    service.chat_notify("truncated", message="cut off")

    assert {"type": "error", "code": "truncated", "message": "cut off"} in emitted


def test_clear_empties_chat_history():
    service, emitted, clock = _service()
    service.toggle_chat()

    for text in ("hello", "how are you"):
        service.handle_event(_transcript(clock, text))

    assert service.chat_history == ["Unknown: hello", "Unknown: how are you"]
    assert len(service._pending_chat) == 2

    assert service.clear() is True
    assert service.chat_history == []
    assert service._pending_chat == []
    assert {"type": "cleared"} in emitted


def test_clear_invokes_chat_reset():
    cfg = Config()
    resets: list[int] = []
    service = AssistantService(
        cfg, FakeSubscriber(), time_fn=Clock(), chat_reset=lambda: resets.append(1)
    )
    service.toggle_chat()

    service.clear()

    assert resets == [1]


def test_clear_does_not_change_mode():
    service, _, clock = _service()
    service.toggle_chat()
    assert service.modes.mode.value == "chat"

    service.clear()

    assert service.modes.mode.value == "chat"


def test_clear_always_returns_true():
    service, _, _ = _service()
    assert service.clear() is True

    service.toggle_chat()
    service.handle_event(_ev(events.TRANSCRIPT, Clock(), text="test"))
    assert service.clear() is True


def test_heartbeat_emits_inference_depth_and_in_flight():
    """Calling tick() emits heartbeat with inference_depth (int) and inference_in_flight (bool)."""
    service, emitted, clock = _service()
    service.tick()

    heartbeats = [e for e in emitted if e["type"] == "heartbeat"]
    assert len(heartbeats) == 1

    heartbeat = heartbeats[0]
    assert "inference_depth" in heartbeat
    assert isinstance(heartbeat["inference_depth"], int)
    assert "inference_in_flight" in heartbeat
    assert isinstance(heartbeat["inference_in_flight"], bool)


def test_status_contains_inference_depth_and_in_flight():
    """status() snapshot contains inference_depth (int) and inference_in_flight (bool)."""
    service, _, _ = _service()

    status = service.status()
    assert "inference_depth" in status
    assert isinstance(status["inference_depth"], int)
    assert "inference_in_flight" in status
    assert isinstance(status["inference_in_flight"], bool)


_AUTO = object()


class FakeTriage:
    def __init__(
        self,
        questions=_AUTO,
        state_pair=("ok", ""),
        escalate=False,
        parsed=_AUTO,
        normalized=_AUTO,
    ):
        if questions is _AUTO:
            questions = {"urgency": {"type": "noul", "instructions": "how urgent"}}
        self._questions = questions
        self._state_pair = state_pair
        self._escalate = escalate
        self._parsed = parsed
        self._normalized = normalized
        self.build_state_calls: list[tuple] = []
        self.normalize_calls: list = []
        self.prompt_calls: list[tuple] = []

    @property
    def questions(self):
        return self._questions

    def state(self):
        return self._state_pair

    def build_state(self, window):
        self.build_state_calls.append(tuple(window))
        return "\n".join("- " + str(item) for item in (window or []))

    def normalize(self, payload):
        self.normalize_calls.append(payload)
        if self._normalized is not _AUTO:
            return self._normalized
        if not isinstance(payload, dict):
            return {}
        signals = {"order": [], "items": {}, "escalated": False, "input_tokens": 0}
        signals.update(payload)
        if self._escalate:
            signals["escalated"] = True
        return signals

    def build_extraction_prompt(self, window, signals):
        self.prompt_calls.append(
            (tuple(window), dict(signals) if isinstance(signals, dict) else signals)
        )
        return "EXTRACT"

    def parse_extraction(self, text):
        if self._parsed is _AUTO:
            return {
                "summary": "s",
                "opportunities": [],
                "doubts": [],
                "action_items": [],
                "risks": [],
            }
        return self._parsed


class FakeLane:
    def __init__(self, ready=True):
        self._ready = ready
        self.submits: list[tuple] = []
        self.drops = 0
        self.predict_gate = threading.Event()

    def ready(self):
        return self._ready

    def state(self):
        return ("ok", "") if self._ready else ("loading", "")

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

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


def _ok_reply(text):
    return "ok"


def _triage_service(triage=None, lane=None, extract_fn=None, chat_fn: Any = _AUTO):
    cfg = Config()
    cfg.assistant.turn_gap_s = 0.5
    clock = Clock()
    emitted: list[dict] = []

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

    if triage is None:
        triage = FakeTriage()
    if lane is None:
        lane = FakeLane()
    if chat_fn is _AUTO:
        chat_fn = _ok_reply
    service = AssistantService(
        cfg,
        FakeSubscriber(),
        emit=emit,
        time_fn=clock,
        triage=triage,
        laya_lane=lane,
        triage_extract_fn=extract_fn,
        chat_fn=chat_fn,
    )
    return service, emitted, clock, triage, lane


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


def _escalated_payload():
    return {"order": [], "items": {}, "escalated": True, "input_tokens": 3}


def test_triage_status_disabled_when_no_analyzer():
    service, _, _ = _service()

    status = service.status()

    assert status["triage_state"] == "disabled"
    assert status["triage_reason"] == ""
    assert status["triage"] is None


def test_triage_status_reflects_analyzer():
    service, _, _, _, _ = _triage_service(
        triage=FakeTriage(state_pair=("loading", "warming up"))
    )

    status = service.status()

    assert status["triage_state"] == "loading"
    assert status["triage_reason"] == "warming up"
    assert status["triage"] is None


def test_triage_captures_transcript_lang():
    service, _, clock, _, _ = _triage_service()
    assert service._last_transcript_lang is None

    event = _transcript(clock, "hola")
    event["lang"] = "es"
    service.handle_event(event)
    assert service._last_transcript_lang == "es"

    service.handle_event(_transcript(clock, "hello"))
    assert service._last_transcript_lang is None


def test_maybe_triage_passes_lang_to_lane():
    service, _, clock, _, lane = _triage_service()
    event = _transcript(clock, "hola")
    event["lang"] = "es"
    service.handle_event(event)
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()

    assert len(lane.submits) == 1
    _, _, lang, _ = lane.submits[0]
    assert lang == "es"


def test_maybe_triage_runs_in_idle_mode():
    service, _, clock, triage, lane = _triage_service()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()

    assert len(lane.submits) == 1
    state, questions, lang, gen = lane.submits[0]
    assert state == "- hello"
    assert questions is triage.questions
    assert lang is None
    assert gen == 0
    assert triage.build_state_calls == [("hello",)]
    assert service.triage_history == ["hello"]


def test_maybe_triage_uses_transcript_buffer_not_chat_history():
    service, _, clock, triage, lane = _triage_service(chat_fn=None)
    service.chat_history = ["stale chat only"]
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()

    assert lane.submits == []
    assert triage.build_state_calls == []

    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()

    assert len(lane.submits) == 1
    assert triage.build_state_calls == [("hello",)]


def test_maybe_triage_requires_settled_turn():
    service, _, clock, _, lane = _triage_service()
    service.handle_event(_transcript(clock, "hello"))
    service.tick()

    assert lane.submits == []


def test_maybe_triage_skips_unready_lane():
    service, _, clock, _, lane = _triage_service(lane=FakeLane(ready=False))
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()

    assert lane.submits == []


def test_maybe_triage_skips_without_questions():
    service, _, clock, triage, lane = _triage_service(
        triage=FakeTriage(questions=None)
    )
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()

    assert lane.submits == []
    assert triage.build_state_calls == []


def test_maybe_triage_dedupes_identical_window():
    service, _, clock, triage, lane = _triage_service()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()

    assert len(lane.submits) == 1
    state, questions, lang, gen = lane.submits[0]
    assert state == "- hello"
    assert questions is triage.questions
    assert lang is None
    assert gen == 0

    service.tick()
    service.tick()
    assert len(lane.submits) == 1

    service.handle_event(_transcript(clock, "second"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    assert len(lane.submits) == 2
    assert lane.submits[1][0] == "- hello\n- second"


def test_second_transcript_produces_second_submit():
    service, _, clock, _, lane = _triage_service()
    service.handle_event(_transcript(clock, "first"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    assert len(lane.submits) == 1

    service.handle_event(_transcript(clock, "second"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()

    assert len(lane.submits) == 2
    assert lane.submits[1][0] == "- first\n- second"


def test_tick_under_50ms_with_blocked_lane():
    service, _, clock, _, lane = _triage_service()
    assert not lane.predict_gate.is_set()
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)

    start = time.perf_counter()
    service.tick()
    elapsed = time.perf_counter() - start

    assert len(lane.submits) == 1
    assert elapsed < 0.05


def test_on_triage_result_emits_signals():
    service, emitted, clock, _, lane = _triage_service()
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]
    payload = {
        "order": ["urgency"],
        "items": {"urgency": {"kind": "noul"}},
        "escalated": False,
        "input_tokens": 12,
    }

    service._on_triage_result(gen, payload)

    staged = _triage_events(emitted)
    assert len(staged) == 1
    assert staged[0]["stage"] == "signals"
    assert staged[0]["triage"] == service.last_triage
    assert staged[0]["triage"]["escalated"] is False
    assert staged[0]["triage"]["input_tokens"] == 12
    assert service.status()["triage"] == service.last_triage


def test_triage_events_carry_the_server_timestamp():
    service, emitted, clock, _, lane = _triage_service(extract_fn=lambda prompt: "RAW")
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]
    expected_ts = clock()

    service._on_triage_result(gen, _escalated_payload())

    signals = [e for e in _triage_events(emitted) if e.get("stage") == "signals"]
    assert len(signals) == 1
    assert isinstance(signals[0]["ts"], (int, float))
    assert signals[0]["ts"] == expected_ts

    job = [j for j in service._inference._queue if j.name == "triage-extract"][0]
    job.run()

    contents = [e for e in _triage_events(emitted) if e.get("stage") == "content"]
    assert len(contents) == 1
    assert isinstance(contents[0]["ts"], (int, float))
    assert contents[0]["ts"] == expected_ts  # content reuses the signals timestamp


def test_on_triage_result_suppresses_stale_generation():
    service, emitted, clock, _, lane = _triage_service()
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]
    service.clear()

    service._on_triage_result(gen, _escalated_payload())

    assert service.last_triage is None
    assert [e for e in _triage_events(emitted) if e.get("stage") == "signals"] == []

    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    assert len(lane.submits) == 2


def test_on_triage_result_accepts_non_chat_mode():
    service, emitted, clock, _, lane = _triage_service()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]
    window = service._last_triage_window

    service._on_triage_result(gen, _escalated_payload())

    assert service.last_triage is not None
    assert service._last_triage_window == window
    assert [e.get("stage") for e in _triage_events(emitted)] == ["signals"]


def test_on_triage_result_ignores_none_payload():
    service, emitted, clock, _, lane = _triage_service()
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]

    service._on_triage_result(gen, None)

    assert service.last_triage is None
    assert _triage_events(emitted) == []


def test_none_payload_enqueues_state_only_extract():
    service, emitted, clock, triage, lane = _triage_service(
        extract_fn=lambda prompt: "RAW"
    )
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]

    service._on_triage_result(gen, None)

    assert service.last_triage is None
    assert [e.get("stage") for e in _triage_events(emitted)] == []
    assert [j.name for j in service._inference._queue] == ["chat", "triage-extract"]
    job = [j for j in service._inference._queue if j.name == "triage-extract"][0]
    job.run()

    contents = [e for e in _triage_events(emitted) if e.get("stage") == "content"]
    assert len(contents) == 1
    assert contents[0]["triage"] == {"content": triage.parse_extraction("RAW")}
    assert contents[0]["ts"] == clock()
    assert service.last_triage == contents[0]["triage"]
    assert triage.prompt_calls == [(service._last_triage_window, None)]


def test_none_payload_without_extract_fn_enqueues_nothing():
    service, emitted, clock, _, lane = _triage_service()
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]

    service._on_triage_result(gen, None)

    assert service.last_triage is None
    assert _triage_events(emitted) == []
    assert [j.name for j in service._inference._queue] == ["chat"]


def test_on_triage_result_ignores_empty_signals():
    service, emitted, clock, _, lane = _triage_service(
        triage=FakeTriage(normalized={})
    )
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]

    service._on_triage_result(gen, _escalated_payload())

    assert service.last_triage is None
    assert _triage_events(emitted) == []


def test_escalation_enqueues_extract_job():
    service, emitted, clock, _, lane = _triage_service(
        triage=FakeTriage(escalate=True), extract_fn=lambda prompt: "RAW"
    )
    service.toggle_chat()
    service.handle_event(_transcript(clock, "we might churn"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]

    service._on_triage_result(
        gen, {"order": [], "items": {}, "escalated": False, "input_tokens": 0}
    )

    assert service.last_triage is not None
    assert service.last_triage["escalated"] is True
    assert service._inference.depth() == 2
    assert [j.name for j in service._inference._queue] == ["chat", "triage-extract"]
    assert [e.get("stage") for e in _triage_events(emitted)] == ["signals"]


def test_extract_runs_without_escalation():
    service, _, clock, _, lane = _triage_service(extract_fn=lambda prompt: "RAW")
    service.toggle_chat()
    service.handle_event(_transcript(clock, "all calm"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]

    service._on_triage_result(
        gen, {"order": [], "items": {}, "escalated": False, "input_tokens": 0}
    )

    assert service.last_triage is not None
    assert service._inference.depth() == 2
    assert [j.name for j in service._inference._queue] == ["chat", "triage-extract"]


def test_no_extract_when_disabled_in_config():
    service, _, clock, _, lane = _triage_service(
        triage=FakeTriage(escalate=True), extract_fn=lambda prompt: "RAW"
    )
    service.cfg.triage.extract_enabled = False
    service.toggle_chat()
    service.handle_event(_transcript(clock, "churn risk"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]

    service._on_triage_result(gen, _escalated_payload())

    assert service.last_triage is not None
    assert service.last_triage["escalated"] is True
    assert [j.name for j in service._inference._queue] == ["chat"]


def test_extract_refused_when_queue_full():
    service, emitted, clock, _, lane = _triage_service(
        triage=FakeTriage(escalate=True), extract_fn=lambda prompt: "RAW"
    )
    service._inference = InferenceQueue(size=2)
    service.toggle_chat()
    service.handle_event(_transcript(clock, "churn risk"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    assert service._inference.depth() == 1
    _, _, _, gen = lane.submits[0]

    service._on_triage_result(gen, _escalated_payload())

    assert service._inference.depth() == 1
    assert service._inference._queue[0].name == "chat"
    assert service.last_triage is not None
    assert service.last_triage["escalated"] is True
    assert [e.get("stage") for e in _triage_events(emitted)] == ["signals"]


def test_extract_merges_content():
    service, emitted, _, triage, _ = _triage_service(extract_fn=lambda prompt: "RAW")
    service.toggle_chat()
    window = ("hello",)
    signals = _escalated_payload()

    service._run_triage_extract(0, window, signals)

    contents = [e for e in _triage_events(emitted) if e.get("stage") == "content"]
    assert len(contents) == 1
    assert contents[0]["triage"] == {**signals, "content": triage.parse_extraction("RAW")}
    assert service.last_triage == contents[0]["triage"]
    assert triage.prompt_calls == [(window, signals)]


def test_extract_none_parse_keeps_signals():
    service, emitted, _, _, _ = _triage_service(
        triage=FakeTriage(parsed=None), extract_fn=lambda prompt: "GARBAGE"
    )
    service.toggle_chat()
    signals = _escalated_payload()
    service.last_triage = signals

    service._run_triage_extract(0, ("hello",), signals)

    assert [e for e in _triage_events(emitted) if e.get("stage") == "content"] == []
    assert service.last_triage == signals
    assert "content" not in service.last_triage


def test_extract_error_is_swallowed():
    def boom(prompt):
        raise RuntimeError("llm down")

    service, emitted, _, _, _ = _triage_service(extract_fn=boom)
    service.toggle_chat()

    service._run_triage_extract(0, ("hello",), _escalated_payload())

    assert _triage_events(emitted) == []
    assert service.last_triage is None


def test_extract_suppressed_after_reset():
    service, emitted, _, _, _ = _triage_service(extract_fn=lambda prompt: "RAW")
    service.toggle_chat()
    service._reset_triage()
    emitted.clear()

    service._run_triage_extract(0, ("hello",), _escalated_payload())

    assert _triage_events(emitted) == []
    assert service.last_triage is None


def test_extract_works_after_leaving_chat():
    seen = []
    service, emitted, _, _, _ = _triage_service(
        extract_fn=lambda prompt: seen.append(prompt) or "RAW"
    )
    service.toggle_chat()
    assert service.modes.exit_chat() is True

    service._run_triage_extract(0, ("hello",), _escalated_payload())

    assert seen != []
    assert [e.get("stage") for e in _triage_events(emitted)] == ["content"]
    assert service.last_triage is not None
    assert "content" in service.last_triage


def test_clear_resets_triage():
    service, emitted, clock, _, lane = _triage_service(
        triage=FakeTriage(escalate=True), extract_fn=lambda prompt: "RAW"
    )
    service.toggle_chat()
    service.handle_event(_transcript(clock, "churn"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]
    service._on_triage_result(gen, _escalated_payload())
    assert service._inference.depth() == 2

    assert service.clear() is True

    assert service._triage_generation == 1
    assert service.last_triage is None
    assert service._last_triage_window == ()
    assert service.triage_history == []
    assert lane.drops == 1
    assert service._inference.depth() == 0
    assert list(service._inference._queue) == []
    cleared = [e for e in _triage_events(emitted) if e.get("stage") == "cleared"]
    assert len(cleared) == 1
    assert cleared[0]["triage"] is None
    assert {"type": "cleared"} in emitted

    signals_before = len(
        [e for e in _triage_events(emitted) if e.get("stage") == "signals"]
    )
    service._on_triage_result(gen, _escalated_payload())
    assert service.last_triage is None
    assert (
        len([e for e in _triage_events(emitted) if e.get("stage") == "signals"])
        == signals_before
    )

    submits_before = len(lane.submits)
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    assert len(lane.submits) == submits_before


def test_exit_chat_does_not_reset_triage():
    service, emitted, clock, _, lane = _triage_service(
        extract_fn=lambda prompt: "RAW"
    )
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    assert len(lane.submits) == 1
    _, _, _, gen = lane.submits[0]
    service._on_triage_result(gen, _escalated_payload())
    assert service.last_triage is not None
    assert service._triage_generation == 0

    assert service.toggle_chat() is True

    assert service.modes.mode.value == "idle"
    assert service._triage_generation == 0
    assert service.last_triage is not None
    assert service._last_triage_window == ("hello",)
    assert service.triage_history == ["hello"]
    assert lane.drops == 0
    cleared = [e for e in _triage_events(emitted) if e.get("stage") == "cleared"]
    assert cleared == []


def test_clear_without_triage_emits_no_triage_event():
    service, emitted, _ = _service()

    service.clear()

    assert {"type": "cleared"} in emitted
    assert _triage_events(emitted) == []


def test_triage_end_to_end_signals_content_cleared():
    prompts: list[str] = []
    service, emitted, clock, triage, lane = _triage_service(
        extract_fn=lambda prompt: prompts.append(prompt) or '{"summary": "s"}'
    )
    service.toggle_chat()

    service.handle_event(_transcript(clock, "first thought"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    assert len(lane.submits) == 1
    _, _, _, gen0 = lane.submits[0]
    service._on_triage_result(
        gen0, {"order": ["urgency"], "items": {}, "escalated": False, "input_tokens": 1}
    )

    assert [e.get("stage") for e in _triage_events(emitted)] == ["signals"]
    assert prompts == []
    assert service.last_triage is not None
    assert service.last_triage["escalated"] is False
    assert [j.name for j in service._inference._queue] == ["chat", "triage-extract"]

    triage._escalate = True
    service.handle_event(_transcript(clock, "we might churn"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    assert len(lane.submits) == 2
    _, _, _, gen1 = lane.submits[1]
    service._on_triage_result(gen1, _escalated_payload())

    assert [e.get("stage") for e in _triage_events(emitted)] == ["signals", "signals"]
    assert [j.name for j in service._inference._queue] == [
        "chat",
        "triage-extract",
        "chat",
        "triage-extract",
    ]
    for job in [j for j in service._inference._queue if j.name == "triage-extract"]:
        job.run()

    contents = [e for e in _triage_events(emitted) if e.get("stage") == "content"]
    assert len(contents) == 2
    assert len(prompts) == 2
    assert contents[0]["triage"]["escalated"] is False
    assert contents[1]["triage"]["escalated"] is True
    assert contents[1]["triage"]["content"] == triage.parse_extraction("x")
    assert service.last_triage == contents[1]["triage"]
    assert service.status()["triage"] == contents[1]["triage"]

    service.clear()

    cleared = [e for e in _triage_events(emitted) if e.get("stage") == "cleared"]
    assert len(cleared) == 1
    assert cleared[0]["triage"] is None
    assert service.last_triage is None
    service._on_triage_result(gen1, _escalated_payload())
    assert [e.get("stage") for e in _triage_events(emitted)] == [
        "signals",
        "signals",
        "content",
        "content",
        "cleared",
    ]

    assert service.toggle_chat() is True
    service.handle_event(_transcript(clock, "idle thought"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    assert len(lane.submits) == 3
    assert lane.submits[2][0].endswith("- idle thought")


def test_tick_emits_status_on_triage_state_change():
    service, emitted, _, triage, _ = _triage_service(
        triage=FakeTriage(state_pair=("loading", "warming up"))
    )

    service.tick()
    statuses = [e for e in emitted if e["type"] == "status"]
    assert len(statuses) == 1
    assert statuses[0]["triage_state"] == "loading"
    assert statuses[0]["triage_reason"] == "warming up"
    assert "triage" in statuses[0]

    triage._state_pair = ("ok", "")
    service.tick()
    statuses = [e for e in emitted if e["type"] == "status"]
    assert len(statuses) == 2
    assert statuses[1]["triage_state"] == "ok"
    assert statuses[1]["triage_reason"] == ""
    assert "triage" in statuses[1]

    service.tick()
    assert len([e for e in emitted if e["type"] == "status"]) == 2


def test_no_status_emitted_when_triage_disabled():
    service, emitted, _ = _service()

    service.tick()
    service.tick()

    assert [e for e in emitted if e["type"] == "status"] == []


def test_tick_heals_triage_panel_after_lane_ready():
    import time as _time
    from pathlib import Path

    from stt_daemon.assistant.laya_lane import LayaLane
    from stt_daemon.assistant.triage import TriageAnalyzer
    from stt_daemon.tui.state import ViewState

    class _FakeRouter:
        def preload(self, names):
            return self

        def predict(self, state, questions, model=None, lang=None, **kwargs):
            if "_warmup" in questions:
                return {"model": model, "answers": {}, "usage": {}, "routing": {}}
            return {
                "model": model,
                "answers": {
                    "frustration": {
                        "noul": 0.1,
                        "confidence": 0.9,
                        "answer_confidence": 0.9,
                    },
                    "urgency": {
                        "choice": "low",
                        "confidence": 0.9,
                        "answer_confidence": 0.9,
                    },
                },
                "usage": {"input_tokens": 5},
            }

    def _factory(**kwargs):
        return _FakeRouter()

    cfg = Config()
    cfg.assistant.turn_gap_s = 0.5
    cfg.triage.questions_file = str(
        Path(__file__).resolve().parent.parent / "questions.example.json"
    )
    clock = Clock()
    emitted: list[dict] = []

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

    analyzer = TriageAnalyzer(cfg.triage)
    assert analyzer.questions, "example questions must load"
    holder: dict = {}
    lane = LayaLane(
        cfg.triage,
        agent_factory=_factory,
        on_result=lambda gen, payload: holder["svc"]._on_triage_result(gen, payload),
        on_state=lambda state, reason: analyzer.refresh(state, reason),
    )
    service = AssistantService(
        cfg,
        FakeSubscriber(),
        emit=emit,
        time_fn=clock,
        triage=analyzer,
        laya_lane=lane,
        chat_fn=_ok_reply,
    )
    holder["svc"] = service
    try:
        lane.start()
        deadline = _time.monotonic() + 5.0
        while not lane.ready() and _time.monotonic() < deadline:
            _time.sleep(0.005)
        assert lane.ready(), "lane did not become ready"
        assert analyzer.state() == ("ok", "")

        service.tick()
        statuses = [e for e in emitted if e["type"] == "status"]
        assert statuses, "tick emitted no status after lane ready"
        assert statuses[-1]["triage_state"] == "ok"

        service.handle_event(_transcript(clock, "hello"))
        clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
        service.tick()

        deadline = _time.monotonic() + 5.0
        while service.last_triage is None and _time.monotonic() < deadline:
            _time.sleep(0.005)
        assert service.last_triage is not None, "no triage signals from lane"

        view = ViewState()
        for event in emitted:
            view.ingest(event)
        assert view.triage_state == "ok"
        lines = view.triage_lines()
        assert lines
        assert not any("unavailable" in line for line in lines)
        assert not any("loading" in line for line in lines)
    finally:
        lane.stop()


def _dump_extract(prompt, on_io=None):
    """Extract fake that mimics LlmClient.chat's on_io phases."""
    if on_io is not None:
        on_io({"phase": "request", "url": "http://llm/v1", "body": {"prompt": prompt}})
        on_io({"phase": "response", "payload": {"content": "parsed"}, "error": None})
    return "RAW"


def _run_triage_flow(service, emitted, clock, lane, payload):
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]
    service._on_triage_result(gen, payload)
    return gen


def _dump_files(tmp_path):
    return list(tmp_path.rglob("*.json"))


def test_dump_writes_full_laya_and_llm_io(tmp_path):
    service, emitted, clock, triage, lane = _triage_service(extract_fn=_dump_extract)
    service._triage_dump = TriageDumper(True, directory=str(tmp_path))
    service.cfg.triage.dump_io = True

    _run_triage_flow(service, emitted, clock, lane, _escalated_payload())
    jobs = [j for j in service._inference._queue if j.name == "triage-extract"]
    assert len(jobs) == 1
    jobs[0].run()

    files = _dump_files(tmp_path)
    assert len(files) == 1
    assert re.fullmatch(r"\d{6}-\d{3}\.json", files[0].name)
    data = json.loads(files[0].read_text(encoding="utf-8"))
    assert data["ts"] == clock()
    assert isinstance(data["iso"], str) and data["iso"]
    assert data["laya"]["request"]["state"] == "- hello"
    assert data["laya"]["request"]["questions"] == triage.questions
    assert data["laya"]["request"]["model"] == "multilingual"
    assert data["laya"]["request"]["max_len"] == 1024
    assert data["laya"]["response"] == _escalated_payload()
    assert data["laya"]["error"] is None
    assert data["llm"]["request"]["url"] == "http://llm/v1"
    assert data["llm"]["request"]["body"] == {"prompt": "EXTRACT"}
    assert data["llm"]["response"] == {"content": "parsed"}
    assert data["llm"]["error"] is None


def test_dump_laya_error_keeps_state_only_llm_half(tmp_path):
    service, emitted, clock, _, lane = _triage_service(extract_fn=_dump_extract)
    service._triage_dump = TriageDumper(True, directory=str(tmp_path))
    service.cfg.triage.dump_io = True

    _run_triage_flow(service, emitted, clock, lane, None)
    jobs = [j for j in service._inference._queue if j.name == "triage-extract"]
    assert len(jobs) == 1
    jobs[0].run()

    files = _dump_files(tmp_path)
    assert len(files) == 1
    data = json.loads(files[0].read_text(encoding="utf-8"))
    assert data["laya"]["response"] is None
    assert data["laya"]["error"] == "predict failed"
    assert data["llm"]["request"]["url"] == "http://llm/v1"
    assert data["llm"]["response"] == {"content": "parsed"}


def test_dump_without_extract_has_laya_only(tmp_path):
    service, emitted, clock, _, lane = _triage_service(extract_fn=_dump_extract)
    service._triage_dump = TriageDumper(True, directory=str(tmp_path))
    service.cfg.triage.dump_io = True
    service.cfg.triage.extract_enabled = False

    _run_triage_flow(service, emitted, clock, lane, _escalated_payload())

    files = _dump_files(tmp_path)
    assert len(files) == 1
    data = json.loads(files[0].read_text(encoding="utf-8"))
    assert data["laya"]["response"] == _escalated_payload()
    assert data["llm"] is None


def test_dump_queue_full_has_laya_only(tmp_path):
    service, emitted, clock, _, lane = _triage_service(extract_fn=_dump_extract)
    service._triage_dump = TriageDumper(True, directory=str(tmp_path))
    service.cfg.triage.dump_io = True
    service._inference = InferenceQueue(size=2)

    _run_triage_flow(service, emitted, clock, lane, _escalated_payload())

    assert service._inference.depth() == 1
    files = _dump_files(tmp_path)
    assert len(files) == 1
    data = json.loads(files[0].read_text(encoding="utf-8"))
    assert data["laya"]["response"] == _escalated_payload()
    assert data["llm"] is None


def test_dump_disabled_writes_nothing(tmp_path):
    service, emitted, clock, _, lane = _triage_service(
        extract_fn=lambda prompt: "RAW"
    )
    assert service._triage_dump.enabled is False
    service._triage_dump = TriageDumper(False, directory=str(tmp_path))

    _run_triage_flow(service, emitted, clock, lane, _escalated_payload())
    jobs = [j for j in service._inference._queue if j.name == "triage-extract"]
    assert len(jobs) == 1
    jobs[0].run()

    assert _dump_files(tmp_path) == []


def test_dump_stale_result_discards_without_file(tmp_path):
    service, emitted, clock, triage, lane = _triage_service(extract_fn=_dump_extract)
    service._triage_dump = TriageDumper(True, directory=str(tmp_path))
    service.cfg.triage.dump_io = True
    service.toggle_chat()
    service.handle_event(_transcript(clock, "hello"))
    clock.advance(service.cfg.assistant.turn_gap_s + 0.1)
    service.tick()
    _, _, _, gen = lane.submits[0]

    real_normalize = triage.normalize

    def stale_normalize(payload):
        # Simulate a clear landing between the first check/begin and the re-check.
        service._triage_generation += 1
        return real_normalize(payload)

    triage.normalize = stale_normalize
    service._on_triage_result(gen, _escalated_payload())

    assert _dump_files(tmp_path) == []
    assert service._triage_dump._pending == {}


def test_dump_stale_extract_discards_without_file(tmp_path):
    service, _, clock, _, _ = _triage_service(extract_fn=_dump_extract)
    service._triage_dump = TriageDumper(True, directory=str(tmp_path))
    service.cfg.triage.dump_io = True
    ts = clock()
    service._triage_dump.begin(ts, laya={"request": {}, "response": {}, "error": None})
    service._triage_generation += 1  # stale before the extract even starts

    service._run_triage_extract(0, ("hello",), _escalated_payload(), ts)

    assert _dump_files(tmp_path) == []
    assert service._triage_dump._pending == {}


def test_dump_extract_cleared_mid_flight_writes_nothing(tmp_path):
    service, _, clock, _, _ = _triage_service(extract_fn=_dump_extract)
    service._triage_dump = TriageDumper(True, directory=str(tmp_path))
    service.cfg.triage.dump_io = True

    def stale_extract(prompt, on_io=None):
        # Simulate a clear landing mid-LLM-call (after the request phase fired).
        service._triage_generation += 1
        if on_io is not None:
            on_io({"phase": "request", "url": "http://llm/v1", "body": {}})
            on_io({"phase": "response", "payload": {"c": 1}, "error": None})
        return "RAW"

    service._triage_extract_fn = stale_extract
    ts = clock()
    service._triage_dump.begin(ts, laya={"request": {}, "response": {}, "error": None})

    service._run_triage_extract(0, ("hello",), _escalated_payload(), ts)

    assert _dump_files(tmp_path) == []
    assert service._triage_dump._pending == {}