"""
VoiceAI Pipeline Builder

High-level fluent API for wiring STT → LLM → TTS pipelines.
Handles the entire lifecycle: call events, audio streaming, AI processing,
and TTS injection — all with a single .start() call.

Architecture (2026 cascaded pipeline pattern):

  Inbound call
      ↓
  AudioStream (WebSocket, slin16 → Float32)
      ↓
  SttAdapter.start_stream() → on_transcript(is_final=True)
      ↓
  LlmAdapter.chat() → streaming tokens
      ↓ (first token triggers TTS immediately)
  TtsAdapter.synthesize() → PCM chunks
      ↓
  TtsInjector.inject_streaming()
      ↓
  Call audio (Asterisk)

Latency budget target: <500ms end-to-end (STT ~200ms, LLM TTFT ~80ms, TTS TTFA ~150ms)
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from typing import AsyncIterable, Awaitable, Callable

from dvgateway.streams.audio_stream import create_audio_stream
from dvgateway.streams.call_events import CallEventStream
from dvgateway.streams.tts_stream import TtsInjector
from dvgateway.transport.ws_pool import WsPool
from dvgateway.types import (
    CallEndedEvent,
    CallNewEvent,
    CallSession,
    LlmAdapter,
    Logger,
    Message,
    StreamDir,
    SttAdapter,
    TranscriptResult,
    TtsAdapter,
)


# ─── Fallback wrapper ──────────────────────────────────────────────────────


@dataclass
class _WithFallback:
    primary: SttAdapter | TtsAdapter | LlmAdapter
    fallback: SttAdapter | TtsAdapter | LlmAdapter | None = None


# ─── Pipeline State (per active call) ─────────────────────────────────────


@dataclass
class _CallPipelineState:
    session: CallSession
    conversation_history: list[Message] = field(default_factory=list)
    stt_adapter: SttAdapter | None = None
    audio_stream_unsubscribe: Callable[[], None] = lambda: None


# ─── Builder Stages ──────────────────────────────────────────────────────


class SttStage:
    """First stage: STT adapter configured."""

    def __init__(self, adapter: SttAdapter, builder: PipelineBuilder) -> None:
        self._adapter = adapter
        self._fallback: SttAdapter | None = None
        self._builder = builder

    def fallback(self, adapter: SttAdapter) -> SttStage:
        self._fallback = adapter
        return self

    def _to_wrapped(self) -> _WithFallback:
        return _WithFallback(primary=self._adapter, fallback=self._fallback)

    def llm(self, adapter: LlmAdapter) -> LlmStage:
        return LlmStage(adapter, self._to_wrapped(), self._builder)

    def for_conference(self) -> ConferencePipelineBuilder:
        return ConferencePipelineBuilder(self._builder, self._to_wrapped())


class LlmStage:
    """Second stage: LLM adapter configured."""

    def __init__(self, adapter: LlmAdapter, stt: _WithFallback, builder: PipelineBuilder) -> None:
        self._adapter = adapter
        self._fallback: LlmAdapter | None = None
        self._stt = stt
        self._builder = builder

    def fallback(self, adapter: LlmAdapter) -> LlmStage:
        self._fallback = adapter
        return self

    def _to_wrapped(self) -> _WithFallback:
        return _WithFallback(primary=self._adapter, fallback=self._fallback)

    def tts(self, adapter: TtsAdapter) -> TtsStage:
        return TtsStage(adapter, self._stt, self._to_wrapped(), self._builder)


class TtsStage:
    """Third stage: TTS adapter configured. Can set callbacks and start."""

    def __init__(
        self,
        adapter: TtsAdapter,
        stt: _WithFallback,
        llm: _WithFallback,
        builder: PipelineBuilder,
    ) -> None:
        self._adapter = adapter
        self._fallback: TtsAdapter | None = None
        self._stt = stt
        self._llm = llm
        self._builder = builder

    def fallback(self, adapter: TtsAdapter) -> TtsStage:
        self._fallback = adapter
        return self

    def audio_filter(self, dir: StreamDir) -> TtsStage:
        self._builder._dir = dir
        return self

    def on_new_call(self, handler: Callable[[CallSession], None | Awaitable[None]]) -> TtsStage:
        self._builder._on_new_call = handler
        return self

    def on_call_ended(self, handler: Callable[[str, float], None | Awaitable[None]]) -> TtsStage:
        self._builder._on_call_ended = handler
        return self

    def on_transcript(
        self,
        handler: Callable[[TranscriptResult, CallSession], None | Awaitable[None]],
    ) -> TtsStage:
        self._builder._on_transcript = handler
        return self

    def on_error(self, handler: Callable[[Exception, str | None], None | Awaitable[None]]) -> TtsStage:
        self._builder._on_error = handler
        return self

    async def start(self) -> None:
        tts_wrapped = _WithFallback(primary=self._adapter, fallback=self._fallback)
        await self._builder._start_pipeline(self._stt, self._llm, tts_wrapped)

    async def stop(self) -> None:
        self._builder._stop_pipeline()


# ─── Conference Pipeline (STT-only, no LLM/TTS) ───────────────────────────


class ConferencePipelineBuilder:
    """Conference-only pipeline: STT without LLM/TTS."""

    def __init__(self, builder: PipelineBuilder, stt: _WithFallback) -> None:
        self._builder = builder
        self._stt = stt
        self._on_transcript: Callable[[TranscriptResult, CallSession], None | Awaitable[None]] | None = None
        self._on_error: Callable[[Exception, str | None], None | Awaitable[None]] | None = None

    def on_transcript(
        self,
        handler: Callable[[TranscriptResult, CallSession], None | Awaitable[None]],
    ) -> ConferencePipelineBuilder:
        self._on_transcript = handler
        return self

    def on_error(
        self,
        handler: Callable[[Exception, str | None], None | Awaitable[None]],
    ) -> ConferencePipelineBuilder:
        self._on_error = handler
        return self

    async def start(self) -> None:
        await self._builder._start_conference_pipeline(self._stt, self._on_transcript, self._on_error)


# ─── Main Builder ──────────────────────────────────────────────────────────


class PipelineBuilder:
    """Main pipeline builder — wires STT → LLM → TTS with call lifecycle management."""

    def __init__(
        self,
        ws_pool: WsPool,
        call_events: CallEventStream,
        tts_injector: TtsInjector,
        logger: Logger,
    ) -> None:
        self._ws_pool = ws_pool
        self._call_events = call_events
        self._tts_injector = tts_injector
        self._logger = logger

        self._dir: StreamDir = "both"
        self._on_new_call: Callable[[CallSession], None | Awaitable[None]] | None = None
        self._on_call_ended: Callable[[str, float], None | Awaitable[None]] | None = None
        self._on_transcript: Callable[[TranscriptResult, CallSession], None | Awaitable[None]] | None = None
        self._on_error: Callable[[Exception, str | None], None | Awaitable[None]] | None = None
        self._active_calls: dict[str, _CallPipelineState] = {}
        self._running = False
        self._stop_event: asyncio.Event | None = None

    def stt(self, adapter: SttAdapter) -> SttStage:
        """Set the STT adapter (first step in fluent chain)."""
        return SttStage(adapter, self)

    async def _start_pipeline(
        self,
        stt: _WithFallback,
        llm: _WithFallback,
        tts: _WithFallback,
    ) -> None:
        if self._running:
            raise RuntimeError("Pipeline is already running. Call stop() first.")
        self._running = True
        self._stop_event = asyncio.Event()

        system_prompt = getattr(llm.primary, "system_prompt", None) or (
            "You are a helpful voice assistant. Keep answers concise and conversational."
        )

        async def handle_event(event: object) -> None:
            if isinstance(event, CallNewEvent):
                session = event.session
                self._logger.info({"linked_id": session.linked_id}, "New call — starting AI pipeline")
                if self._on_new_call:
                    result = self._on_new_call(session)
                    if asyncio.iscoroutine(result):
                        await result
                asyncio.ensure_future(self._handle_call(session, stt, llm, tts, system_prompt))

            elif isinstance(event, CallEndedEvent):
                state = self._active_calls.get(event.linked_id)
                if state:
                    state.audio_stream_unsubscribe()
                    if state.stt_adapter:
                        await state.stt_adapter.stop()
                    del self._active_calls[event.linked_id]
                if self._on_call_ended:
                    result = self._on_call_ended(event.linked_id, event.duration_sec)
                    if asyncio.iscoroutine(result):
                        await result
                self._logger.info(
                    {"linked_id": event.linked_id, "duration": event.duration_sec},
                    "Call ended",
                )

        unsub = self._call_events.on(handle_event)

        # Keep pipeline alive until stop() is called
        await self._stop_event.wait()
        unsub()

    async def _handle_call(
        self,
        session: CallSession,
        stt: _WithFallback,
        llm: _WithFallback,
        tts: _WithFallback,
        system_prompt: str,
    ) -> None:
        conversation_history: list[Message] = [
            Message(role="system", content=system_prompt),
        ]

        audio_stream = create_audio_stream(self._ws_pool, session.linked_id, self._logger, self._dir)

        stt_adapter = stt.primary
        assert isinstance(stt_adapter, SttAdapter)

        state = _CallPipelineState(
            session=session,
            conversation_history=conversation_history,
            stt_adapter=stt_adapter,
            audio_stream_unsubscribe=audio_stream.unsubscribe,
        )
        self._active_calls[session.linked_id] = state

        def on_transcript_cb(result: TranscriptResult) -> None:
            if not result.is_final:
                return
            asyncio.ensure_future(
                self._process_transcript(result, session, conversation_history, llm, tts)
            )

        stt_adapter.on_transcript(on_transcript_cb)

        try:
            await stt_adapter.start_stream(session.linked_id, audio_stream)
        except Exception as e:
            self._logger.error({"err": str(e), "linked_id": session.linked_id}, "STT stream error")
            if self._on_error:
                result = self._on_error(e, session.linked_id)
                if asyncio.iscoroutine(result):
                    await result

    async def _process_transcript(
        self,
        result: TranscriptResult,
        session: CallSession,
        conversation_history: list[Message],
        llm: _WithFallback,
        tts: _WithFallback,
    ) -> None:
        try:
            if self._on_transcript:
                cb_result = self._on_transcript(result, session)
                if asyncio.iscoroutine(cb_result):
                    await cb_result

            self._logger.debug({"linked_id": session.linked_id}, "Transcript finalized, calling LLM")
            conversation_history.append(Message(role="user", content=result.text))

            llm_adapter = llm.primary
            assert isinstance(llm_adapter, LlmAdapter)

            response_text = ""
            async for token in llm_adapter.chat(list(conversation_history)):
                response_text += token
                if response_text == token and token:
                    self._logger.debug({"linked_id": session.linked_id}, "LLM first token received, starting TTS")

            if not response_text:
                return

            conversation_history.append(Message(role="assistant", content=response_text))

            tts_adapter = tts.primary
            assert isinstance(tts_adapter, TtsAdapter)

            audio_chunks = tts_adapter.synthesize(response_text)
            await self._tts_injector.inject_streaming(session.linked_id, audio_chunks)

        except Exception as e:
            self._logger.error({"err": str(e), "linked_id": session.linked_id}, "Pipeline error")
            if self._on_error:
                cb_result = self._on_error(e, session.linked_id)
                if asyncio.iscoroutine(cb_result):
                    await cb_result

    async def _start_conference_pipeline(
        self,
        stt: _WithFallback,
        on_transcript: Callable[[TranscriptResult, CallSession], None | Awaitable[None]] | None,
        on_error: Callable[[Exception, str | None], None | Awaitable[None]] | None,
    ) -> None:
        if self._running:
            raise RuntimeError("Pipeline is already running.")
        self._running = True
        self._stop_event = asyncio.Event()

        async def handle_event(event: object) -> None:
            from dvgateway.types import ConfJoinEvent, ConfLeaveEvent

            if isinstance(event, (CallNewEvent, ConfJoinEvent)):
                linked_id = event.session.linked_id if isinstance(event, CallNewEvent) else event.linked_id
                fake_session = (
                    event.session
                    if isinstance(event, CallNewEvent)
                    else CallSession(linked_id=linked_id, conf_id=event.conf_id)
                )

                audio_stream = create_audio_stream(self._ws_pool, linked_id, self._logger, self._dir)
                stt_adapter = stt.primary
                assert isinstance(stt_adapter, SttAdapter)

                def transcript_cb(result: TranscriptResult) -> None:
                    if on_transcript:
                        r = on_transcript(result, fake_session)
                        if asyncio.iscoroutine(r):
                            asyncio.ensure_future(r)

                stt_adapter.on_transcript(transcript_cb)

                try:
                    await stt_adapter.start_stream(linked_id, audio_stream)
                except Exception as e:
                    if on_error:
                        r = on_error(e, linked_id)
                        if asyncio.iscoroutine(r):
                            await r

            elif isinstance(event, (CallEndedEvent, ConfLeaveEvent)):
                linked_id = event.linked_id
                state = self._active_calls.get(linked_id)
                if state:
                    state.audio_stream_unsubscribe()
                    if state.stt_adapter:
                        await state.stt_adapter.stop()
                    del self._active_calls[linked_id]

        unsub = self._call_events.on(handle_event)
        await self._stop_event.wait()
        unsub()

    def _stop_pipeline(self) -> None:
        self._running = False
        for state in self._active_calls.values():
            state.audio_stream_unsubscribe()
            if state.stt_adapter:
                asyncio.ensure_future(state.stt_adapter.stop())
        self._active_calls.clear()
        if self._stop_event:
            self._stop_event.set()
