"""Owns TTS payload preparation and synthesis for a Session.

Responsibilities:

- Apply output processors (markdown stripping, phonetic replacement,
  pauses) to raw agent text and produce a :class:`TTSInput` payload.
- Drive the underlying :class:`TTSSynthesizer` to produce audio chunks
  and feed them to the outbound audio queue owned by :class:`AudioRouter`.
- Provide the single-shot :meth:`synthesize_bypass` path used by
  greeting announcements.
- Track the in-flight synthesis task so cancellation can target it.
- Reserve the future :meth:`_synthesize_sentences` private hook for
  sentence-level pipelining.

Outbound queue ownership note: the :class:`BoundedAudioQueue` that
carries TTS audio out to the transport lives on :class:`Session`
because :class:`Session.start` may rebuild it when the previous queue
was closed.  :class:`AudioRouter` holds the live reference for draining
and :class:`TTSSynthesizer` (owned by this scheduler) holds the same
reference for writing.
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Callable
from typing import TYPE_CHECKING

from easycat._bounded_queue import BoundedAudioQueue
from easycat._tts_synthesizer import TTSSynthesizer, TTSSynthResult
from easycat.cancel import CancelToken
from easycat.events import EventBus
from easycat.llm_output_processing import (
    LLMOutputProcessor,
    apply_output_processors,
)
from easycat.runtime.context import RunContext
from easycat.session._journal_sink import SessionJournalSink
from easycat.stages.tts import TTSStage
from easycat.timeouts import TimeoutConfig
from easycat.tts.input import TTSInput, resolve_tts_input_policy, strip_ssml_tags
from easycat.turn_manager import TurnManager, TurnManagerState

if TYPE_CHECKING:
    from easycat._epoch import Lease
    from easycat._turn_context import TurnContext
    from easycat.session._audio_router import AudioRouter
    from easycat.session._wiring import SessionWiringContext

logger = logging.getLogger(__name__)


class TTSScheduler:
    """Prepares and synthesizes TTS payloads for a Session."""

    def __init__(
        self,
        *,
        wiring: SessionWiringContext,
        tts_stage: TTSStage,
        turn_manager: TurnManager,
        event_bus: EventBus,
        journal_sink: SessionJournalSink,
        run_ctx: RunContext,
        no_turn: TurnContext,
        audio_router: AudioRouter,
        outbound_queue: BoundedAudioQueue,
        timeout_config: TimeoutConfig | None,
        audio_gate: Callable[[], bool] | None,
        # Config
        output_processors: list[LLMOutputProcessor],
        strip_markdown_enabled: bool,
    ) -> None:
        self._tts_getter = wiring.tts
        self._tts_stage = tts_stage
        self._turn_manager = turn_manager
        self._event_bus = event_bus
        self._journal_sink = journal_sink
        self._run_ctx = run_ctx
        self._no_turn = no_turn
        self._audio_router = audio_router

        self._output_processors = output_processors
        self._strip_markdown = strip_markdown_enabled

        self._current_turn = wiring.current_turn
        self._capture_identity = wiring.capture_identity
        self._is_gated = wiring.is_gated
        self._session_actions = wiring.session_actions
        self._drain_session_actions = wiring.drain_session_actions
        self._clear_turn = wiring.clear_turn

        self._synth = TTSSynthesizer(
            tts=wiring.tts(),
            event_bus=event_bus,
            outbound_queue=outbound_queue,
            timeout_config=timeout_config,
            correlation_ids=wiring.correlation_ids,
            audio_gate=audio_gate,
            direct_first_audio=audio_router.try_send_first_audio_inline,
        )
        self._synth.bind_stage(
            tts_stage,
            run_ctx_getter=lambda: self._run_ctx,
            turn_getter=lambda: self._current_turn() or self._no_turn,
        )

        self._current_tts_task: asyncio.Task[None] | None = None
        self._playback_suppressed: bool = False

    # ── Properties ─────────────────────────────────────────────

    @property
    def is_playback_suppressed(self) -> bool:
        return self._playback_suppressed

    def set_playback_suppressed(self, value: bool) -> None:
        self._playback_suppressed = value

    @property
    def active_turn_task(self) -> asyncio.Task[None] | None:
        """The in-flight ``TurnRunner.on_turn_ended`` task for the active turn.

        This is the *whole* end-of-turn coroutine (STT finalize → agent stream →
        TTS), not just the TTS synthesis step. Cancelling it aborts the agent
        stream, so ``Session.cancel_tts_playback`` must leave it running.
        """
        return self._current_tts_task

    @active_turn_task.setter
    def active_turn_task(self, task: asyncio.Task[None] | None) -> None:
        self._current_tts_task = task

    @property
    def strip_markdown_enabled(self) -> bool:
        return self._strip_markdown

    @property
    def synthesizer(self) -> TTSSynthesizer:
        """Underlying :class:`TTSSynthesizer` for low-level access.

        Session ``start()`` uses this to swap the outbound queue when
        rebuilding the queue after a prior teardown.
        """
        return self._synth

    def replace_outbound_queue(self, queue: BoundedAudioQueue) -> None:
        """Point the underlying synthesizer at a rebuilt outbound queue.

        Session.start() re-creates the queue after a prior teardown; this
        keeps the producer side (the synthesizer) pointed at the same
        instance the :class:`AudioRouter` drains, without Session reaching
        into synthesizer internals.
        """
        self._synth.replace_outbound_queue(queue)

    # ── Payload preparation ────────────────────────────────────

    def prepare(self, text: str, *, is_streaming: bool, is_final: bool) -> TTSInput:
        original_payload = TTSInput(text=text, format="plain")
        payload = original_payload
        processed_payload = apply_output_processors(
            payload,
            self._output_processors,
            is_final=is_final,
            is_streaming=is_streaming,
        )
        payload = processed_payload
        input_policy = resolve_tts_input_policy(self._tts_getter())
        ssml_downgraded = processed_payload.format == "ssml" and not input_policy.accepts("ssml")
        if ssml_downgraded:
            payload = TTSInput(text=strip_ssml_tags(processed_payload.text), format="plain")
        self._record_tts_payload_prepared(
            original_text=original_payload.text,
            original_format=original_payload.format,
            prepared_payload=payload,
            ssml_downgraded=ssml_downgraded,
            is_streaming=is_streaming,
            is_final=is_final,
        )
        return payload

    # ── Synthesis ──────────────────────────────────────────────

    async def finalize_speaking_turn(
        self,
        turn: TurnContext | None,
        *,
        identity: Lease[TurnContext | None] | None = None,
        activity: Lease[TurnManagerState] | None = None,
    ) -> bool:
        """Run the end-of-turn drain → stop → drain → clear sequence.

        The single owner of the bot-speaking stop / turn-clear decision:
        the streaming agent path (``TurnRunner._process_tts``) calls this
        so the barge-in / clear semantics live in exactly one place.

        Drains pending session actions (end_call, transfer) *before*
        transitioning the turn manager to IDLE. Since that drain can yield
        to cancellation-resistant application cleanup, every later manager
        transition is guarded against a barge-in successor. When
        ``should_stop`` is signalled the outbound audio is drained before
        stopping; otherwise the manager stops first and the queued tail
        audio is drained afterwards so the router can still record sent
        bytes and emit playback marks.

        The turn pointer is only cleared while the captured identity remains
        current and the exact bot-speaking transition still owns the manager
        activity. Replacement or same-object/state republication fences the
        stale finalizer.

        Returns ``True`` if a drained session action signalled that the
        session should stop.
        """
        identity = identity or self._capture_identity()
        activity = activity or self._turn_manager.capture_activity()
        should_stop = await self._drain_session_actions()
        settled_activity: Lease[TurnManagerState] | None = None
        try:
            if not self._turn_is_current(turn, identity) or not self._activity_is_current(
                activity, TurnManagerState.BOT_SPEAKING
            ):
                return should_stop
            if should_stop:
                await self._audio_router.await_drain()
                if self._turn_is_current(turn, identity) and self._activity_is_current(
                    activity, TurnManagerState.BOT_SPEAKING
                ):
                    settled_activity = await self._turn_manager.bot_stopped_speaking()
            else:
                settled_activity = await self._turn_manager.bot_stopped_speaking()
                if (
                    self._turn_is_current(turn, identity)
                    and settled_activity is not None
                    and self._activity_is_current(settled_activity, TurnManagerState.IDLE)
                ):
                    await self._audio_router.await_drain()
        finally:
            actions = self._session_actions()
            if actions is not None:
                actions.clear_no_interrupt()
        turn_still_current = (
            self._turn_is_current(turn, identity)
            and settled_activity is not None
            and self._activity_is_current(settled_activity, TurnManagerState.IDLE)
        )
        if turn_still_current:
            assert settled_activity is not None
            await self._audio_router.flush_trailing_playback_mark(turn)
            turn_still_current = self._turn_is_current(
                turn, identity
            ) and self._activity_is_current(settled_activity, TurnManagerState.IDLE)
            if turn_still_current:
                self._clear_turn()
        return should_stop

    def _turn_is_current(
        self,
        turn: TurnContext | None,
        identity: Lease[TurnContext | None],
    ) -> bool:
        """Whether a finalizer's captured identity still owns ``turn``."""
        return identity.value is turn and identity.guard()

    @staticmethod
    def _activity_is_current(
        activity: Lease[TurnManagerState],
        state: TurnManagerState,
    ) -> bool:
        """Whether a finalizer still owns the expected manager activity."""
        return activity.value is state and activity.guard()

    async def synthesize_bypass(self, text: str) -> None:
        """Synthesize text via TTS, bypassing the classification gate.

        Used for hold audio and screening responses that must reach the
        transport even while the gate is closed.
        """
        await self._synth.synthesize(text, token=None, bypass_gate=True)

    async def begin_synthesis_with_bot_start(
        self,
        payload: TTSInput,
        token: CancelToken | None,
        *,
        is_active: Callable[[], bool] | None,
        lifecycle_ready: asyncio.Future[bool] | None = None,
        activity_started: Callable[[Lease[TurnManagerState]], None] | None = None,
    ) -> asyncio.Task[TTSSynthResult]:
        """Start provider work while dispatching ``BotStartedSpeaking``.

        The provider request starts first and may receive its first event while
        agent-delta and lifecycle handlers run. A synthesizer barrier keeps
        that event private until the optional first-payload admission future
        succeeds and every ``BotStartedSpeaking`` handler completes. If either
        dispatch fails or is cancelled, speculative synthesis is drained.
        """
        barrier = asyncio.Event()
        task = asyncio.create_task(
            self._synth.synthesize(
                payload,
                token,
                is_active=is_active,
                start_barrier=barrier,
            )
        )

        try:
            # Give the provider task one loop turn to issue its network request
            # and settle on either network I/O or the first-event barrier. Keep
            # the yield inside the cleanup guard because cancellation can land
            # at this first suspension point.
            await asyncio.sleep(0)
            if lifecycle_ready is not None:
                allowed = await asyncio.shield(lifecycle_ready)
                if not allowed:
                    task.cancel()
                    barrier.set()
                    await asyncio.gather(task, return_exceptions=True)
                    return task
            # The lifecycle wait above can run arbitrary event handlers. An
            # exact identity/activity guard must still admit this turn before
            # it publishes BOT_SPEAKING; cancellation is only cooperative and
            # cannot stand in for this commit boundary.
            if is_active is not None and not is_active():
                task.cancel()
                barrier.set()
                await asyncio.gather(task, return_exceptions=True)
                return task
            activity = await self._turn_manager.bot_started_speaking()
            if activity is not None and activity_started is not None:
                activity_started(activity)
        except BaseException:
            task.cancel()
            barrier.set()
            await asyncio.gather(task, return_exceptions=True)
            raise

        barrier.set()
        return task

    async def _synthesize_sentences(
        self,
        payloads: object,
        cancel_token: CancelToken | None,
        turn: TurnContext,
    ) -> object:
        """Reserved private placeholder for sentence-level pipelining.

        Not yet implemented — current behaviour is one-at-a-time synthesis
        driven by ``TurnRunner._process_tts`` via ``synthesizer.synthesize``.
        Kept private so the scheduler's public surface only advertises
        implemented methods; the hook exists so the pipelining change is a
        local one when it lands.
        """
        raise NotImplementedError("sentence-level TTS pipelining is not implemented yet")

    # ── Cancellation ───────────────────────────────────────────

    def request_turn_cancel(self) -> asyncio.Task[None] | None:
        """Synchronously cancel and capture the active turn task.

        Returns the exact task that was active at the cutoff boundary.  A
        detached barge-in cleanup can later drain that captured task without
        accidentally cancelling a successor turn installed in the meantime.
        """
        task = self._current_tts_task
        current_task = asyncio.current_task()
        if task is not None and task is not current_task and not task.done():
            task.cancel()
        return task

    async def finish_turn_cancel(self, task: asyncio.Task[None] | None) -> None:
        """Cancel provider work and drain the task captured at cutoff."""
        await self._synth.cancel()
        if task is None or task is asyncio.current_task() or task.done():
            return
        try:
            await task
        except (asyncio.CancelledError, Exception):  # noqa: BLE001, S110 intentional boundary or best-effort cleanup
            pass

    async def cancel(self) -> None:
        task = self.request_turn_cancel()
        await self.finish_turn_cancel(task)

    # ── Journal helpers ────────────────────────────────────────

    def _record_markdown_strip(
        self,
        *,
        phase: str,
        original_text: str,
        stripped_text: str,
        turn_id: str | None = None,
    ) -> None:
        """Append a journal record when final-response markdown stripping runs."""
        self._journal_sink.append_record(
            name="markdown_stripped",
            turn_id=turn_id,
            data={
                "phase": phase,
                "changed": original_text != stripped_text,
                "original_text": original_text,
                "stripped_text": stripped_text,
            },
        )

    def _record_tts_payload_prepared(
        self,
        *,
        original_text: str,
        original_format: str,
        prepared_payload: TTSInput,
        ssml_downgraded: bool,
        is_streaming: bool,
        is_final: bool,
        turn_id: str | None = None,
    ) -> None:
        self._journal_sink.append_record(
            name="tts_payload_prepared",
            turn_id=turn_id,
            data={
                "is_streaming": is_streaming,
                "is_final": is_final,
                "changed": (
                    original_text != prepared_payload.text
                    or original_format != prepared_payload.format
                ),
                "original_text": original_text,
                "original_format": original_format,
                "prepared_text": prepared_payload.text,
                "prepared_format": prepared_payload.format,
                "processors": [type(processor).__name__ for processor in self._output_processors],
                "ssml_downgraded": ssml_downgraded,
            },
        )
