"""Coordinate a voice conversation independently of its delivery channel."""

import asyncio
import logging
from collections.abc import Awaitable, Callable
from contextlib import aclosing

from voice_assistant.application.events import (
    AgentReply,
    AgentTextDelta,
    SessionFailure,
    SpeechChunk,
    TranscriptFinal,
    TranscriptPartial,
    TurnCompleted,
    VoiceEvent,
)
from voice_assistant.stt.adapter import STTAdapter
from voice_assistant.voice_pipeline import (
    PipelineAgentResponse,
    PipelineAudioChunk,
    PipelineTextDelta,
    VoicePipeline,
)

logger = logging.getLogger(__name__)

type EventSink = Callable[[VoiceEvent], Awaitable[None]]
type STTFactory = Callable[[], STTAdapter]


class VoiceSessionError(Exception):
    """An application operation failed; the channel decides how to report it."""

    def __init__(
        self,
        code: str,
        message: str,
        *,
        turn_id: int | None = None,
        retryable: bool = False,
    ) -> None:
        super().__init__(message)
        self.code = code
        self.turn_id = turn_id
        self.retryable = retryable


class VoiceSession:
    """Own STT execution, finalization, interruption, and provider cleanup.

    VoicePipeline retains its existing turn state and Agent/TTS behavior.
    An awaited event sink preserves flow control and cancellation without
    introducing a second, unbounded output queue.
    """

    def __init__(
        self,
        *,
        pipeline: VoicePipeline,
        create_stt: STTFactory,
        emit: EventSink,
        stt_stop_timeout: float = 10.0,
        finalization_cancel_timeout: float = 5.0,
    ) -> None:
        if stt_stop_timeout <= 0 or finalization_cancel_timeout <= 0:
            raise ValueError("Session timeouts must be positive.")
        self._pipeline = pipeline
        self._create_stt = create_stt
        self._emit = emit
        self._stt_stop_timeout = stt_stop_timeout
        self._finalization_cancel_timeout = finalization_cancel_timeout
        self._stt: STTAdapter | None = None
        self._stt_events_task: asyncio.Task[None] | None = None
        self._finalization_task: asyncio.Task[None] | None = None
        self._audio_input_closed = False
        self._turn_transition_lock = asyncio.Lock()
        self._interrupting_turns: set[int] = set()
        self._completed_turns: set[int] = set()

    @property
    def session_id(self) -> str:
        return self._pipeline.session.session_id

    @property
    def active_turn_id(self) -> int | None:
        return self._pipeline.session.active_turn_id

    @property
    def expected_audio_sequence(self) -> int:
        return self._pipeline.session.expected_audio_sequence

    @property
    def closed(self) -> bool:
        return self._pipeline.session.closed

    def validate_audio_input(self, turn_id: int) -> None:
        self._require_open()
        if self._audio_input_closed:
            raise ValueError("Audio input is closed for this turn.")
        if self.active_turn_id != turn_id:
            raise ValueError(
                f"Unexpected turn_id={turn_id}; "
                f"active_turn_id={self.active_turn_id}"
            )

    def _require_open(self) -> None:
        if self.closed:
            raise VoiceSessionError(
                "SESSION_ALREADY_CLOSED", "This session is already closed."
            )

    async def start_turn(self) -> int:
        self._require_open()
        if self.active_turn_id is not None:
            raise VoiceSessionError(
                "INVALID_TURN", "A turn is already active.",
                turn_id=self.active_turn_id,
            )
        turn_id = None
        try:
            turn_id = self._pipeline.start_turn().turn_id
            self._stt = self._create_stt()
            await self._stt.start_session()
            self._stt_events_task = asyncio.create_task(
                self._consume_stt_events(self._stt, turn_id)
            )
            self._pipeline.register_turn_task(self._stt_events_task)
        except (RuntimeError, ValueError) as exc:
            cleanup_complete = await self._rollback_start(turn_id)
            raise VoiceSessionError(
                "STT_START_FAILED", str(exc), turn_id=turn_id,
                retryable=cleanup_complete,
            ) from exc
        self._audio_input_closed = False
        return turn_id

    async def _rollback_start(self, turn_id: int | None) -> bool:
        cleanup_complete = True
        try:
            await self._cancel_stt_task()
        except Exception:
            cleanup_complete = False
            logger.exception(
                "Failed to cancel task during audio.start rollback. "
                "session_id=%s turn_id=%s", self.session_id, turn_id,
            )
        if self._stt is not None:
            try:
                await self._stt.abort_session()
            except Exception:
                cleanup_complete = False
                logger.exception(
                    "Failed to abort STT during audio.start rollback. "
                    "session_id=%s turn_id=%s", self.session_id, turn_id,
                )
        if cleanup_complete:
            if turn_id is not None:
                self._pipeline.interrupt_turn(turn_id)
            self._stt = None
        return cleanup_complete

    async def send_audio(
        self, turn_id: int, sequence_number: int, audio: bytes,
    ) -> None:
        self.validate_audio_input(turn_id)
        self._pipeline.session.accept_audio_chunk(
            turn_id, sequence_number, len(audio),
        )
        if self._stt is None:
            raise VoiceSessionError(
                "STT_NOT_STARTED", "audio.start is required before audio.",
                turn_id=turn_id,
            )
        try:
            await self._stt.send_audio(audio)
        except (RuntimeError, ValueError) as exc:
            raise VoiceSessionError(
                "STT_AUDIO_ERROR", str(exc), turn_id=turn_id, retryable=True,
            ) from exc

    def stop_turn(self, turn_id: int) -> None:
        self._require_open()
        if self._stt is None:
            raise RuntimeError("STT session has not been started.")
        if not self._pipeline.is_active_turn(turn_id):
            raise ValueError(f"Unexpected turn_id={turn_id}.")
        if self._finalization_task is not None and not self._finalization_task.done():
            raise RuntimeError("Turn finalization is already running.")
        self._audio_input_closed = True
        self._finalization_task = asyncio.create_task(
            self._finalize_turn(turn_id, self._stt)
        )

    async def interrupt(self, turn_id: int) -> None:
        self._require_open()
        if not self._pipeline.is_active_turn(turn_id):
            raise VoiceSessionError(
                "INVALID_TURN", f"Turn {turn_id} is no longer active.",
                turn_id=turn_id,
            )
        try:
            async with self._turn_transition_lock:
                self._interrupting_turns.add(turn_id)
            await self._cancel_finalization_task()
            await self._pipeline.cancel_turn_task()
            if self._stt is not None:
                await self._stt.abort_session()
            self._pipeline.interrupt_if_active(turn_id)
            self._stt_events_task = None
        except Exception as exc:
            raise VoiceSessionError(
                "TURN_CANCEL_FAILED", str(exc), turn_id=turn_id,
            ) from exc

    async def end(self) -> None:
        self._require_open()
        try:
            async with self._turn_transition_lock:
                if self.active_turn_id is not None:
                    self._interrupting_turns.add(self.active_turn_id)
            await self._cancel_finalization_task()
            await self._pipeline.cancel_turn_task()
            if self._stt is not None:
                await self._stt.abort_session()
            self._pipeline.close_session()
            self._stt_events_task = None
            self._stt = None
        except (RuntimeError, ValueError, TimeoutError) as exc:
            raise VoiceSessionError(
                "SESSION_END_FAILED", str(exc), turn_id=self.active_turn_id,
                retryable=True,
            ) from exc

    async def _consume_stt_events(self, stt: STTAdapter, turn_id: int) -> None:
        async for event in stt.events():
            if not self._pipeline.is_active_turn(turn_id):
                break
            if event.type == "partial":
                self._pipeline.on_partial(turn_id, event.text)
                await self._emit(TranscriptPartial(self.session_id, turn_id, event.text))
            elif event.type == "final":
                self._pipeline.on_final(turn_id, event.text)
                await self._emit(TranscriptFinal(self.session_id, turn_id, event.text))
                await self._respond(turn_id, event.text)

    async def _respond(self, turn_id: int, text: str) -> None:
        if callable(getattr(self._pipeline.agent, "stream", None)):
            async with aclosing(self._pipeline.stream_turn_response(
                session_id=self.session_id, turn_id=turn_id, text=text,
            )) as response_events:
                async for event in response_events:
                    if not self._pipeline.is_active_turn(turn_id):
                        break
                    if isinstance(event, PipelineTextDelta):
                        await self._emit(AgentTextDelta(self.session_id, turn_id, event.text))
                    elif isinstance(event, PipelineAgentResponse):
                        await self._emit(AgentReply(self.session_id, turn_id, event.text))
                    elif isinstance(event, PipelineAudioChunk):
                        await self._emit(SpeechChunk(self.session_id, turn_id, event.audio))
        else:
            result = await self._pipeline.process_transcript(
                session_id=self.session_id, turn_id=turn_id, text=text,
            )
            await self._emit(AgentReply(self.session_id, turn_id, result.response_text))
            async for audio in self._pipeline.synthesize_response(result.response_text):
                await self._emit(SpeechChunk(self.session_id, turn_id, audio))
            self._pipeline.on_tts_complete(turn_id)

    async def _finalize_turn(self, turn_id: int, stt: STTAdapter) -> None:
        try:
            try:
                async with asyncio.timeout(self._stt_stop_timeout):
                    await stt.stop_session()
            except TimeoutError:
                await self._handle_stt_stop_failure(
                    turn_id, stt, "STT_STOP_TIMEOUT",
                    "STT finalization exceeded the time limit.",
                )
                return
            except Exception:
                logger.exception(
                    "STT finalization failed. session_id=%s turn_id=%s",
                    self.session_id, turn_id,
                )
                await self._handle_stt_stop_failure(
                    turn_id, stt, "STT_STOP_FAILED",
                    "STT could not finalize audio input.",
                )
                return
            await self._pipeline.wait_turn_task()
            async with self._turn_transition_lock:
                if (
                    turn_id in self._interrupting_turns
                    or turn_id in self._completed_turns
                    or not self._pipeline.is_active_turn(turn_id)
                ):
                    return
                self._pipeline.complete_turn(turn_id)
                self._completed_turns.add(turn_id)
                await self._emit(TurnCompleted(self.session_id, turn_id))
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.exception(
                "Turn finalization failed. session_id=%s turn_id=%s",
                self.session_id, turn_id,
            )
            if turn_id not in self._interrupting_turns:
                try:
                    await self._emit(SessionFailure(
                        self.session_id, turn_id, "TURN_EXECUTION_FAILED", str(exc),
                    ))
                except Exception:
                    logger.exception("Failed to report turn execution failure.")

    async def _handle_stt_stop_failure(
        self, turn_id: int, stt: STTAdapter, code: str, message: str,
    ) -> None:
        if not self._pipeline.is_active_turn(turn_id):
            return
        cleanup_complete = True
        for operation_name, operation in (
            ("turn task cancellation", self._pipeline.cancel_turn_task),
            ("STT abort", stt.abort_session),
        ):
            try:
                await operation()
            except Exception:
                cleanup_complete = False
                logger.exception(
                    "Failed during %s. session_id=%s turn_id=%s",
                    operation_name, self.session_id, turn_id,
                )
        if cleanup_complete:
            self._pipeline.interrupt_if_active(turn_id)
        await self._emit(SessionFailure(
            self.session_id, turn_id, code, message, cleanup_complete,
        ))

    async def _cancel_finalization_task(self) -> None:
        task = self._finalization_task
        if task is None:
            return
        try:
            await self._cancel_task(task, self._finalization_cancel_timeout)
        except TimeoutError as exc:
            raise TimeoutError("Turn finalization cancellation timed out.") from exc
        finally:
            if task.done() and self._finalization_task is task:
                self._finalization_task = None

    @staticmethod
    async def _cancel_task(task: asyncio.Task[None], timeout: float) -> None:
        if task is asyncio.current_task():
            raise RuntimeError("A task cannot cancel itself.")
        if not task.done():
            task.cancel()
        try:
            await asyncio.wait_for(asyncio.shield(task), timeout=timeout)
        except asyncio.CancelledError:
            current = asyncio.current_task()
            if current is not None and current.cancelling():
                raise
            if not task.cancelled():
                raise

    async def _cancel_stt_task(self) -> None:
        task = self._stt_events_task
        if task is None:
            return
        try:
            if self._pipeline.turn_task is task:
                await self._pipeline.cancel_turn_task()
            else:
                await self._cancel_task(task, 5.0)
        finally:
            if task.done():
                self._stt_events_task = None

    async def aclose(self) -> None:
        """Attempt every cleanup operation, including after partial startup.

        The caller must shield this awaited cleanup from transport-scope
        cancellation. Provider lifecycle contracts are unchanged here.
        """
        for description, operation in (
            ("cancel finalization task", self._cancel_finalization_task),
            ("cancel voice turn task", self._pipeline.cancel_turn_task),
            ("clean up unregistered STT task", self._cancel_stt_task),
        ):
            try:
                await operation()
            except Exception:
                logger.exception(
                    f"Failed to {description}. session_id=%s", self.session_id,
                )
        if self._stt is not None:
            try:
                await self._stt.abort_session()
            except Exception:
                logger.exception("Failed to abort STT session. session_id=%s", self.session_id)
        try:
            self._pipeline.close_session()
        except Exception:
            logger.exception("Failed to close voice session. session_id=%s", self.session_id)
        close_agent = getattr(self._pipeline.agent, "aclose", None)
        if callable(close_agent):
            try:
                await close_agent()
            except Exception:
                logger.exception("Failed to close Agent provider. session_id=%s", self.session_id)