witnessbox/engine.py

"""Turn-loop orchestrator — one exchange, end to end, UI-agnostic.

    examiner audio ─┬─► ASR ───────────► examiner_text
                    └─► stance (librosa) ─► CONFIDENT / NEUTRAL / HESITANT
                                              │ steers the witness
    examiner_text ─► ContradictionEngine ─► catch? (deterministic verdict)
    system prompt (persona + stance + tier + leak) ─► LLM ─► witness line
    state.apply_turn(...) ─► win / lose / continue
    witness line ─► VoxCPM2(style = game state) ─► audio   (break beat on win)

Kept free of Gradio so it can be driven from a test or a script.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np

import config
from witnessbox import script, stance as stance_mod
from witnessbox.backends import Backends
from witnessbox.backends.base import TTSResult
from witnessbox.contradictions import CatchResult, ContradictionEngine
from witnessbox.state import GameState, TurnEvents
from witnessbox.stance import StanceResult
from witnessbox.witness import build_system_prompt


@dataclass
class TurnResult:
    examiner_text: str
    stance: StanceResult
    witness_text: str
    witness_audio: np.ndarray | None
    audio_sr: int
    events: TurnEvents
    status: dict
    evidence: str = ""                 # the on-camera catch explanation (honest)
    epilogue_audio: np.ndarray | None = None  # win/lose sting, played after the line
    meta: dict = field(default_factory=dict)


class WitnessBoxEngine:
    def __init__(self, backends: Backends):
        self.b = backends
        self.detector = ContradictionEngine()
        self.state = GameState()

    # ---- intro --------------------------------------------------------- #
    def start(self) -> dict:
        self.state.begin()
        intro = self.b.tts.beat("intro")
        opening = self.b.tts.beat("opening")
        return {
            "narration": script.INTRO_NARRATION,
            "opening_text": script.WITNESS_OPENING,
            "intro_audio": _audio_tuple(intro),
            "opening_audio": _audio_tuple(opening),
            "status": self.state.status(),
            "backend": self.b.kind,
            "backend_note": self.b.note,
        }

    # ---- one turn ------------------------------------------------------ #
    def take_turn(
        self,
        *,
        audio: np.ndarray | None = None,
        sr: int | None = None,
        typed_text: str | None = None,
    ) -> TurnResult:
        if self.state.is_over:
            return self._terminal_result("The examination is already over.")

        # 1) Perceived delivery (always from audio if we have it).
        st = (
            stance_mod.analyze(audio, sr or config.VOICE_SR)
            if audio is not None
            else stance_mod._neutral("no audio (typed input)")
        )

        # 2) What did they say? Typed text wins (mock/accessibility); else ASR.
        if typed_text and typed_text.strip():
            examiner_text = typed_text.strip()
        else:
            examiner_text = self.b.asr.transcribe(audio, sr or config.ASR_SR).text if audio is not None else ""
        if not examiner_text:
            return self._terminal_result(
                "[no question heard]", witness_line="Counselor? I didn't catch that.", stance=st
            )

        # 3) Deterministic verdict on the examiner's words (before the witness reacts).
        catch: CatchResult | None = self.detector.detect(examiner_text, self.state.caught_ids)
        is_catch = bool(catch and catch.is_catch)

        # 4) Build the witness's situation and ask the model for his line.
        leak_target = self.state.choose_leak_target()
        system_prompt = build_system_prompt(
            stance_tier=st.tier,
            witness_tier=self.state.witness_tier(),
            caught_ids=self.state.caught_ids,
            leak_target=leak_target,
        )
        hints = {
            "turn": self.state.turn,
            "stance_tier": st.tier,
            "witness_tier": self.state.witness_tier(),
            "leak_text": leak_target.leak_when_hesitant if leak_target else "",
            "just_caught": is_catch,
            "caught_label": catch.lie.label if (catch and is_catch) else "",
            "near_miss": bool(catch and catch.matched_groups and not is_catch),
        }
        messages = self._messages(examiner_text)
        witness_text = self.b.llm.respond(system_prompt, messages, hints=hints).reply

        # 5) Fold into state -> may trigger win/lose.
        events = self.state.apply_turn(
            examiner_text=examiner_text,
            witness_text=witness_text,
            stance_tier=st.tier,
            catch=catch,
        )

        # 6) Voice. On the winning turn the witness's line is the cached break take.
        epilogue_audio = None
        if events.won:
            break_audio = self.b.tts.beat("break")
            witness_text = script.BREAK_LINE
            # keep the transcript consistent with what's actually spoken/shown
            self.state.transcript[-1].witness_text = witness_text
            witness_audio = _audio_arr(break_audio)
            audio_sr = _audio_sr(break_audio)
            epilogue_audio = _audio_arr(self.b.tts.beat("win"))
        elif events.lost:
            spoken = self.b.tts.speak(witness_text, self.state.voice_style())
            witness_audio, audio_sr = spoken.audio, spoken.sr
            epilogue_audio = _audio_arr(self.b.tts.beat("lose"))
        else:
            spoken = self.b.tts.speak(witness_text, self.state.voice_style())
            witness_audio, audio_sr = spoken.audio, spoken.sr

        return TurnResult(
            examiner_text=examiner_text,
            stance=st,
            witness_text=witness_text,
            witness_audio=witness_audio,
            audio_sr=audio_sr,
            events=events,
            status=self.state.status(),
            evidence=_evidence(catch) if is_catch else "",
            epilogue_audio=epilogue_audio,
            meta={"backend": self.b.kind, "stance_features": st.features},
        )

    # ---- helpers ------------------------------------------------------- #
    def _messages(self, examiner_text: str) -> list[dict]:
        msgs: list[dict] = []
        for rec in self.state.transcript:
            msgs.append({"role": "user", "content": rec.examiner_text})
            msgs.append({"role": "assistant", "content": rec.witness_text})
        msgs.append({"role": "user", "content": examiner_text})
        return msgs

    def _terminal_result(self, examiner_text, witness_line="", stance=None) -> TurnResult:
        st = stance or stance_mod._neutral("n/a")
        return TurnResult(
            examiner_text=examiner_text,
            stance=st,
            witness_text=witness_line,
            witness_audio=None,
            audio_sr=config.VOICE_SR,
            events=TurnEvents(),
            status=self.state.status(),
        )


def _audio_arr(t: TTSResult | None) -> np.ndarray | None:
    return t.audio if t else None


def _audio_sr(t: TTSResult | None) -> int:
    return t.sr if t else config.VOICE_SR


def _audio_tuple(t: TTSResult | None):
    if t is None or t.audio is None:
        return None
    return (t.sr, t.audio)


def _evidence(catch: CatchResult) -> str:
    """Plain, honest explanation of what the examiner surfaced and why it lands."""
    surfaced = ", ".join(f"“{v}”" for v in catch.matched_groups.values())
    return (
        f"CONTRADICTION CONFIRMED — {catch.lie.label}\n"
        f"You surfaced: {surfaced}\n"
        f"On the record: {catch.lie.truth}\n"
        f"(match score {catch.score:.2f} ≥ {config.CATCH_THRESHOLD:.2f})"
    )