src/time_machine/application/encounter_service.py

from __future__ import annotations

from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import contextmanager
from collections.abc import Callable, Iterator
from time import perf_counter
from typing import Any

from time_machine.application.session_state import InMemorySessionRepository
from time_machine.domain.events import (
    ConversationTextEvent,
    DestinationEvent,
    MachineStateEvent,
    NarrationEvent,
    PersonaEvent,
    PortraitEvent,
    SceneEvent,
    SouvenirEvent,
    TemporalEvent,
    TranscriptEvent,
    VisualArtifactEvent,
)
from time_machine.domain.models import (
    AudioResult,
    ConversationResponse,
    ConversationTurn,
    EncounterSession,
    SavedEncounter,
    Transcript,
    VoiceProfile,
)
from time_machine.ports.conversation import ConversationEngine
from time_machine.ports.immersion import ImmersiveExperienceGenerator
from time_machine.ports.persona import PersonaGenerator
from time_machine.ports.souvenir import SouvenirGenerator
from time_machine.ports.speech import TTSAdapter
from time_machine.ports.storage import EncounterStore
from time_machine.ports.trace import TraceSink
from time_machine.ports.world import DestinationGenerator


class EncounterService:
    def __init__(
        self,
        sessions: InMemorySessionRepository,
        destination_generator: DestinationGenerator,
        persona_generator: PersonaGenerator,
        conversation_engine: ConversationEngine,
        tts: TTSAdapter,
        souvenir_generator: SouvenirGenerator,
        store: EncounterStore,
        trace_sink: TraceSink,
        immersive_generator: ImmersiveExperienceGenerator | None = None,
    ) -> None:
        self.sessions = sessions
        self.destination_generator = destination_generator
        self.persona_generator = persona_generator
        self.conversation_engine = conversation_engine
        self.tts = tts
        self.souvenir_generator = souvenir_generator
        self.store = store
        self.trace_sink = trace_sink
        self.immersive_generator = immersive_generator

    def start_encounter(
        self,
        mode: str,
        coordinate_prompt: str | None,
        session_id: str,
        seed: int | None = None,
    ) -> Iterator[TemporalEvent]:
        encounter = EncounterSession(session_id=session_id)
        yield from self._emit(
            MachineStateEvent(
                state="launching",
                message="Temporal engines charging.",
                metadata=self._metadata(encounter),
            )
        )

        with _stage_timer("destination.generate") as destination_timing:
            encounter.destination = self.destination_generator.generate_destination(
                mode=mode,
                coordinate_prompt=coordinate_prompt,
                seed=seed,
            )
        encounter.touch()
        yield from self._emit(
            DestinationEvent(
                destination=encounter.destination,
                message="Destination resolved.",
                metadata=self._metadata(encounter, destination_timing),
            )
        )

        with _stage_timer("persona.generate") as persona_timing:
            encounter.persona = self.persona_generator.generate_persona(
                destination=encounter.destination,
                seed=seed,
            )
        with _stage_timer("tts.prepare_voice") as voice_timing:
            encounter.voice_profile = self.tts.prepare_voice(encounter.persona)
            warm_up = getattr(self.tts, "warm_up", None)
            if callable(warm_up):
                warm_up(encounter.voice_profile)
        immersive_timing = {"stage": "immersion.generate", "elapsed_ms": None}
        if self.immersive_generator is not None:
            with _stage_timer("immersion.generate") as immersive_timing:
                (
                    encounter.scene,
                    encounter.portrait,
                    encounter.narration,
                ) = self._generate_launch_immersion(encounter)
                if encounter.narration is not None:
                    encounter.narration = encounter.narration.model_copy(
                        update={
                            "voice_profile": _narrator_voice_profile(),
                            "audio": self._generate_narration_audio(encounter.narration.text),
                        }
                    )
        encounter.history.append(
            ConversationTurn(
                role="machine",
                text=(
                    "Signal locked onto "
                    f"{encounter.persona.name}: {encounter.persona.immediate_situation}"
                ),
            )
        )
        encounter.touch()
        self.sessions.save(encounter)

        yield from self._emit(
            PersonaEvent(
                persona=encounter.persona,
                message="Ordinary-person signal locked.",
                metadata=self._metadata(encounter, persona_timing, voice_timing),
            )
        )
        if encounter.scene is not None:
            yield from self._emit(
                SceneEvent(
                    scene=encounter.scene,
                    message="World image stabilized.",
                    metadata=self._metadata(encounter, immersive_timing),
                )
            )
        if encounter.portrait is not None:
            yield from self._emit(
                PortraitEvent(
                    portrait=encounter.portrait,
                    message="Character image resolved.",
                    metadata=self._metadata(encounter, immersive_timing),
                )
            )
        if encounter.narration is not None:
            yield from self._emit(
                NarrationEvent(
                    narration=encounter.narration,
                    message="Narrator signal prepared.",
                    metadata=self._metadata(encounter, immersive_timing),
                )
            )
        yield from self._emit(
            MachineStateEvent(
                state="conversation_ready",
                message="The radio channel is open.",
                metadata=self._metadata(encounter),
            )
        )

    def handle_text_turn(
        self,
        encounter_id: str,
        text: str,
    ) -> Iterator[TemporalEvent]:
        encounter = self.sessions.get(encounter_id)
        if encounter.destination is None or encounter.persona is None:
            raise ValueError("Encounter must have a destination and persona before turns.")

        user_turn = ConversationTurn(role="user", text=text)
        encounter.history.append(user_turn)
        yield from self._emit(
            TranscriptEvent(
                transcript=Transcript(text=text, confidence=1.0, language="en"),
                message="User signal decoded.",
                metadata=self._metadata(encounter),
            )
        )

        conversation_started = perf_counter()
        tts_started = perf_counter()
        response = ConversationResponse(text="")
        conversation_timing = {"stage": "conversation.respond", "elapsed_ms": None}
        spoken_audio_until = 0
        pending_tts = _PendingTTS()
        try:
            for text_event in self.conversation_engine.stream_response(
                destination=encounter.destination,
                persona=encounter.persona,
                history=encounter.history,
                latest_user_text=text,
            ):
                conversation_timing = {
                    "stage": "conversation.respond",
                    "elapsed_ms": round((perf_counter() - conversation_started) * 1000, 1),
                }
                response = ConversationResponse(
                    text=text_event.text,
                    prosody_hint=_optional_string(text_event.metadata.get("prosody_hint")),
                    ui_signal_event=text_event.message,
                )
                yield from self._emit(
                    text_event.model_copy(
                        update={
                            "metadata": self._metadata(encounter, conversation_timing),
                        }
                    )
                )
                yield from self._emit_pending_tts(pending_tts.drain_ready())
                if encounter.voice_profile is not None:
                    segment, spoken_audio_until = _next_tts_segment(
                        text_event.text,
                        consumed_chars=spoken_audio_until,
                        is_final=text_event.is_final,
                    )
                    if segment:
                        pending_tts.submit(
                            lambda segment=segment, prosody_hint=response.prosody_hint, conversation_timing=conversation_timing: self._build_tts_segment_events(
                                encounter=encounter,
                                text=segment,
                                prosody_hint=prosody_hint,
                                conversation_timing=conversation_timing,
                                tts_started=tts_started,
                            )
                        )
                        yield from self._emit_pending_tts(pending_tts.drain_ready())

            yield from self._emit_pending_tts(pending_tts.drain_ready())
            if not response.text.strip():
                with _stage_timer("conversation.respond") as conversation_timing:
                    response = self.conversation_engine.respond(
                        destination=encounter.destination,
                        persona=encounter.persona,
                        history=encounter.history,
                        latest_user_text=text,
                    )
                yield from self._emit(
                    ConversationTextEvent(
                        text=response.text,
                        is_final=True,
                        message=response.ui_signal_event,
                        metadata=self._metadata(encounter, conversation_timing),
                    )
                )
                if encounter.voice_profile is not None:
                    segment, spoken_audio_until = _next_tts_segment(
                        response.text,
                        consumed_chars=spoken_audio_until,
                        is_final=True,
                    )
                    if segment:
                        pending_tts.submit(
                            lambda segment=segment, prosody_hint=response.prosody_hint, conversation_timing=conversation_timing: self._build_tts_segment_events(
                                encounter=encounter,
                                text=segment,
                                prosody_hint=prosody_hint,
                                conversation_timing=conversation_timing,
                                tts_started=tts_started,
                            )
                        )

            yield from self._emit_pending_tts(pending_tts.drain_all())
        except Exception:
            pending_tts.close(cancel=True)
            raise
        else:
            pending_tts.close()

        response = ConversationResponse(
            text=response.text.strip(),
            prosody_hint=response.prosody_hint,
            ui_signal_event=response.ui_signal_event,
        )
        encounter.history.append(ConversationTurn(role="character", text=response.text))
        encounter.touch()
        self.sessions.save(encounter)

    def _build_tts_segment_events(
        self,
        encounter: EncounterSession,
        text: str,
        prosody_hint: str | None,
        conversation_timing: dict[str, Any],
        tts_started: float,
    ) -> list[TemporalEvent]:
        if encounter.voice_profile is None:
            return []
        events: list[TemporalEvent] = []
        for audio_event in self.tts.stream_audio(
            text=text,
            voice_profile=encounter.voice_profile,
            prosody_hint=prosody_hint,
        ):
            tts_timing = {
                "stage": "tts.synthesize",
                "elapsed_ms": round((perf_counter() - tts_started) * 1000, 1),
            }
            events.append(
                audio_event.model_copy(
                    update={
                        "message": audio_event.message or "Character audio prepared.",
                        "metadata": self._metadata(encounter, conversation_timing, tts_timing),
                    }
                )
            )
        return events

    def _stream_tts_segment(
        self,
        encounter: EncounterSession,
        text: str,
        prosody_hint: str | None,
        conversation_timing: dict[str, Any],
        tts_started: float,
    ) -> Iterator[TemporalEvent]:
        for event in self._build_tts_segment_events(
            encounter=encounter,
            text=text,
            prosody_hint=prosody_hint,
            conversation_timing=conversation_timing,
            tts_started=tts_started,
        ):
            yield from self._emit(event)

    def _emit_pending_tts(self, events: Iterator[TemporalEvent]) -> Iterator[TemporalEvent]:
        for event in events:
            yield from self._emit(event)

    def generate_souvenir(self, encounter_id: str) -> Iterator[TemporalEvent]:
        encounter = self.sessions.get(encounter_id)
        if encounter.destination is None or encounter.persona is None:
            raise ValueError("Encounter must have a destination and persona.")

        encounter.souvenir = self.souvenir_generator.generate_souvenir(
            destination=encounter.destination,
            persona=encounter.persona,
            history=encounter.history,
        )
        if self.immersive_generator is not None:
            encounter.artifact = self.immersive_generator.generate_artifact(
                destination=encounter.destination,
                persona=encounter.persona,
                history=encounter.history,
                souvenir=encounter.souvenir,
            )
        encounter.touch()
        self.sessions.save(encounter)

        if encounter.artifact is not None:
            yield from self._emit(
                VisualArtifactEvent(
                    artifact=encounter.artifact,
                    message="Temporal artifact materialized.",
                    metadata=self._metadata(encounter),
                )
            )
        yield from self._emit(
            SouvenirEvent(
                souvenir=encounter.souvenir,
                message="Temporal souvenir stabilized.",
                metadata=self._metadata(encounter),
            )
        )

    def save_encounter(self, encounter_id: str) -> SavedEncounter:
        encounter = self.sessions.get(encounter_id)
        return self.store.save_encounter(encounter)

    def _emit(self, event: TemporalEvent) -> Iterator[TemporalEvent]:
        self.trace_sink.record_event(event)
        yield event

    def _generate_launch_immersion(
        self,
        encounter: EncounterSession,
    ) -> tuple[Any, Any, Any]:
        if self.immersive_generator is None or encounter.destination is None or encounter.persona is None:
            return None, None, None
        with ThreadPoolExecutor(max_workers=3) as executor:
            scene_future = executor.submit(
                self.immersive_generator.generate_scene,
                encounter.destination,
                encounter.persona,
            )
            portrait_future = executor.submit(
                self.immersive_generator.generate_portrait,
                encounter.destination,
                encounter.persona,
            )
            narration_future = executor.submit(
                self.immersive_generator.generate_narration,
                encounter.destination,
                encounter.persona,
            )
            return scene_future.result(), portrait_future.result(), narration_future.result()

    def _generate_narration_audio(self, text: str) -> AudioResult | None:
        if not getattr(self.tts, "supports_narration_audio", False):
            return None
        try:
            return self.tts.synthesize(
                text=text,
                voice_profile=_narrator_voice_profile(),
                prosody_hint="cinematic narrator, calm, spacious, not the character voice",
            )
        except Exception:
            return None

    def _metadata(self, encounter: EncounterSession, *timings: dict[str, Any]) -> dict[str, Any]:
        metadata: dict[str, Any] = {
            "session_id": encounter.session_id,
            "encounter_id": encounter.encounter_id,
        }
        completed_timings = [timing for timing in timings if timing.get("elapsed_ms") is not None]
        if completed_timings:
            metadata["timings"] = completed_timings
            metadata["elapsed_ms"] = round(sum(timing["elapsed_ms"] for timing in completed_timings), 1)
        return metadata


@contextmanager
def _stage_timer(stage: str) -> Iterator[dict[str, Any]]:
    timing: dict[str, Any] = {"stage": stage, "elapsed_ms": None}
    started = perf_counter()
    try:
        yield timing
    finally:
        timing["elapsed_ms"] = round((perf_counter() - started) * 1000, 1)


def _optional_string(value: object) -> str | None:
    return value if isinstance(value, str) and value else None


def _narrator_voice_profile() -> VoiceProfile:
    return VoiceProfile(
        voice_id="cinematic-temporal-narrator",
        description=(
            "A distinct cinematic narrator voice: measured, spacious, resonant, "
            "documentary tone, clearly not the encountered character."
        ),
        pace="slow",
        emotion="awe",
        accent_hint="neutral English narrator",
    )


class _PendingTTS:
    def __init__(self) -> None:
        self.executor = ThreadPoolExecutor(max_workers=1)
        self.futures: list[Future[list[TemporalEvent]]] = []
        self.closed = False

    def submit(self, events: Callable[[], list[TemporalEvent]]) -> None:
        if self.closed:
            return
        self.futures.append(self.executor.submit(lambda: list(events())))

    def drain_ready(self) -> Iterator[TemporalEvent]:
        while self.futures and self.futures[0].done():
            future = self.futures.pop(0)
            yield from future.result()

    def drain_all(self) -> Iterator[TemporalEvent]:
        try:
            while self.futures:
                future = self.futures.pop(0)
                yield from future.result()
        finally:
            self.close()

    def close(self, *, cancel: bool = False) -> None:
        if self.closed:
            return
        self.closed = True
        self.executor.shutdown(wait=not cancel, cancel_futures=cancel)


def _next_tts_segment(
    text: str,
    consumed_chars: int,
    is_final: bool,
) -> tuple[str, int]:
    clean_text = text.strip()
    if not clean_text or consumed_chars >= len(clean_text):
        return "", consumed_chars

    start = consumed_chars
    while start < len(clean_text) and clean_text[start].isspace():
        start += 1
    if start >= len(clean_text):
        return "", start

    if is_final:
        return clean_text[start:].strip(), len(clean_text)

    boundary = _last_sentence_boundary(clean_text, start)
    if boundary <= start:
        return "", consumed_chars
    return clean_text[start:boundary].strip(), boundary


def _last_sentence_boundary(text: str, start: int) -> int:
    boundary = -1
    for mark in ".!?":
        boundary = max(boundary, text.rfind(mark))
    if boundary < start:
        return -1
    return boundary + 1