"""Streaming TTS session for text-in/audio-out streaming."""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union

from kugelaudio.exceptions import (
    AuthenticationError,
    InsufficientCreditsError,
    KugelAudioError,
)
from kugelaudio.models import AudioChunk, StreamConfig

logger = logging.getLogger(__name__)

# Default timeout for waiting on audio generation responses
_DEFAULT_RECV_TIMEOUT_S = 30.0
# Short poll interval for checking if more messages are available
_POLL_TIMEOUT_S = 0.05


class StreamingSession:
    """WebSocket session for streaming text input and audio output.

    This allows streaming text (e.g., from an LLM) and receiving audio
    as it's generated. Text is buffered and processed when sentence
    boundaries are detected or when explicitly flushed.

    Example:
        async with client.tts.streaming_session(voice_id=123) as session:
            async for token in llm_stream:
                async for chunk in session.send(token):
                    play_audio(chunk)

            # Flush remaining text
            async for chunk in session.flush():
                play_audio(chunk)
    """

    def __init__(
        self,
        api_key: str,
        tts_url: str,
        config: Optional[StreamConfig] = None,
    ):
        self._api_key = api_key
        self._tts_url = tts_url
        self._config = config or StreamConfig()
        self._ws = None
        self._session_id: Optional[str] = None
        self._is_started = False
        self._config_sent = False
        self._pending_messages: List[Dict[str, Any]] = []

    async def __aenter__(self) -> StreamingSession:
        await self.connect()
        return self

    async def __aexit__(self, *args) -> None:
        await self.close()

    async def connect(self) -> None:
        """Connect to the streaming WebSocket endpoint."""
        try:
            import websockets
        except ImportError:
            raise ImportError(
                "websockets required. Install with: pip install websockets"
            )

        ws_url = self._tts_url.replace("https://", "wss://").replace("http://", "ws://")
        ws_url = f"{ws_url}/ws/tts/stream?api_key={self._api_key}"

        try:
            self._ws = await websockets.connect(ws_url, compression=None)
            self._is_started = True
            logger.debug("Streaming session connected")
        except websockets.exceptions.InvalidStatusCode as e:
            if e.status_code == 401:
                raise AuthenticationError("Invalid API key")
            elif e.status_code == 403:
                raise InsufficientCreditsError("Insufficient credits")
            raise KugelAudioError(f"Connection failed: {e}")

    async def _ensure_config_sent(self) -> None:
        """Send session config on first text message.

        The server's /ws/tts/stream endpoint does NOT send a session_started
        response. Instead, it reads config from the first messages in the
        main loop. We send config fields alongside the first text message.
        """
        if self._config_sent:
            return
        self._config_sent = True

    async def send(
        self,
        text: str,
        flush: bool = False,
    ) -> AsyncIterator[AudioChunk]:
        """Send text and yield any generated audio chunks.

        Args:
            text: Text to add to buffer
            flush: Force flush the buffer

        Yields:
            AudioChunk as audio is generated
        """
        if not self._is_started:
            await self.connect()

        # Build message
        msg: Dict[str, Any] = {"text": text, "flush": flush}

        # Attach config on first message (server reads config from message fields)
        if not self._config_sent:
            config = self._config.to_dict()
            msg.update(config)
            self._config_sent = True

        await self._ws.send(json.dumps(msg))

        # Collect any generated audio
        async for chunk in self._receive_until_idle():
            yield chunk

    async def flush(self) -> AsyncIterator[AudioChunk]:
        """Flush the text buffer and yield remaining audio.

        Yields:
            AudioChunk as remaining audio is generated
        """
        if not self._is_started or not self._ws:
            return

        await self._ws.send(json.dumps({"flush": True}))

        async for chunk in self._receive_until_idle():
            yield chunk

    async def _receive_until_idle(self) -> AsyncIterator[AudioChunk]:
        """Receive messages until we get a chunk_complete or session_closed.

        The server sends:
          - {"generation_started": true, ...} when synthesis begins
          - {"audio": "<base64>", ...}       for each audio chunk
          - {"chunk_complete": true, ...}    when a text chunk is done
          - {"session_closed": true, ...}    on session end
          - {"error": "..."}                 on error

        We use a two-tier timeout strategy:
          1. _DEFAULT_RECV_TIMEOUT_S (30s) for the first message after we
             send text (the server may be processing / queued).
          2. _POLL_TIMEOUT_S (50ms) after receiving at least one audio chunk
             to check if more messages are immediately available before
             returning control to the caller.
        """
        received_any_audio = False

        while True:
            timeout = _POLL_TIMEOUT_S if received_any_audio else _DEFAULT_RECV_TIMEOUT_S
            try:
                msg = await asyncio.wait_for(self._ws.recv(), timeout=timeout)
                data = json.loads(msg)

                if data.get("error"):
                    raise KugelAudioError(data["error"])

                if data.get("generation_started"):
                    # Server acknowledged it started generating - reset timeout
                    received_any_audio = False
                    continue

                if data.get("audio"):
                    received_any_audio = True
                    yield AudioChunk.from_dict(data)

                if data.get("chunk_complete") or data.get("session_closed"):
                    # Generation complete for this text chunk
                    break

            except asyncio.TimeoutError:
                # After long timeout with no response: give up
                # After short poll timeout with audio received: done for now
                break
            except Exception as e:
                if "ConnectionClosed" in str(type(e)):
                    break
                raise

    async def close(self) -> Dict[str, Any]:
        """Close the session and return stats.

        Returns:
            Session statistics
        """
        stats = {}

        if self._ws:
            try:
                # Send close command
                await self._ws.send(json.dumps({"close": True}))

                # Wait for session stats
                msg = await asyncio.wait_for(self._ws.recv(), timeout=5.0)
                data = json.loads(msg)

                if data.get("session_closed"):
                    stats = data

            except Exception:
                pass
            finally:
                await self._ws.close()
                self._ws = None
                self._is_started = False

        return stats


class StreamingSessionSync:
    """Synchronous wrapper for StreamingSession."""

    def __init__(self, session: StreamingSession):
        self._session = session
        self._loop = asyncio.new_event_loop()

    def __enter__(self) -> StreamingSessionSync:
        self._loop.run_until_complete(self._session.connect())
        return self

    def __exit__(self, *args) -> None:
        self._loop.run_until_complete(self._session.close())
        self._loop.close()

    def send(self, text: str, flush: bool = False):
        """Send text and return generated audio chunks."""

        async def collect():
            chunks = []
            async for chunk in self._session.send(text, flush=flush):
                chunks.append(chunk)
            return chunks

        return self._loop.run_until_complete(collect())

    def flush(self):
        """Flush buffer and return remaining audio chunks."""

        async def collect():
            chunks = []
            async for chunk in self._session.flush():
                chunks.append(chunk)
            return chunks

        return self._loop.run_until_complete(collect())

    def close(self) -> Dict[str, Any]:
        """Close the session."""
        return self._loop.run_until_complete(self._session.close())


class MultiContextSession:
    """WebSocket session for multi-context TTS streaming.

    Allows managing up to 5 independent audio generation contexts over
    a single WebSocket connection. Each context has its own text buffer,
    voice settings, and generation queue.

    Use cases:
    - Multi-speaker conversations with different voices
    - Pre-buffering audio while another stream plays
    - Interleaved audio generation for dynamic conversations

    Example:
        async with client.tts.multi_context_session() as session:
            # Create contexts with different voices
            await session.create_context("narrator", voice_id=123)
            await session.create_context("character", voice_id=456)

            # Send text to different speakers
            async for chunk in session.send("narrator", "The story begins."):
                play_audio("narrator", chunk)

            async for chunk in session.send("character", "Hello!"):
                play_audio("character", chunk)

            # Close specific context
            await session.close_context("narrator")
    """

    def __init__(
        self,
        api_key: str,
        tts_url: str,
        default_voice_id: Optional[int] = None,
        sample_rate: int = 24000,
        cfg_scale: float = 2.0,
        max_new_tokens: int = 2048,
        normalize: bool = True,
        inactivity_timeout: float = 20.0,
    ):
        self._api_key = api_key
        self._tts_url = tts_url
        self._default_voice_id = default_voice_id
        self._sample_rate = sample_rate
        self._cfg_scale = cfg_scale
        self._max_new_tokens = max_new_tokens
        self._normalize = normalize
        self._inactivity_timeout = inactivity_timeout
        self._ws = None
        self._session_id: Optional[str] = None
        self._is_started = False
        self._contexts: set[str] = set()
        self._pending_messages: List[Dict[str, Any]] = []

    async def __aenter__(self) -> "MultiContextSession":
        await self.connect()
        return self

    async def __aexit__(self, *args) -> None:
        await self.close()

    async def connect(self) -> None:
        """Connect to the multi-context WebSocket endpoint."""
        try:
            import websockets
        except ImportError:
            raise ImportError(
                "websockets required. Install with: pip install websockets"
            )

        ws_url = self._tts_url.replace("https://", "wss://").replace("http://", "ws://")
        ws_url = f"{ws_url}/ws/tts/multi?api_key={self._api_key}"

        try:
            self._ws = await websockets.connect(ws_url)
        except websockets.exceptions.InvalidStatusCode as e:
            if e.status_code == 401:
                raise AuthenticationError("Invalid API key")
            elif e.status_code == 403:
                raise InsufficientCreditsError("Insufficient credits")
            raise KugelAudioError(f"Connection failed: {e}")

    async def _start_session(
        self, context_id: str, voice_id: Optional[int] = None
    ) -> None:
        """Start the session with first context creation.

        The server's /ws/tts/multi endpoint does NOT send a ``session_started``
        message.  Instead, the first message that references a ``context_id``
        triggers an implicit ``create_context`` on the server side, which
        responds with ``{"context_created": true, "context_id": "..."}``.

        Per-context voice / generation settings must be nested inside a
        ``voice_settings`` dict (see ``create_context`` helper in the server).
        """
        if self._is_started:
            return

        if not self._ws:
            await self.connect()

        # Build the first message — this implicitly creates the context
        effective_voice_id = voice_id or self._default_voice_id
        init_msg: Dict[str, Any] = {
            "text": " ",
            "context_id": context_id,
            "sample_rate": self._sample_rate,
            "normalize": self._normalize,
        }
        if effective_voice_id is not None:
            init_msg["voice_settings"] = {
                "voice_id": effective_voice_id,
                "cfg_scale": self._cfg_scale,
                "max_new_tokens": self._max_new_tokens,
            }

        await self._ws.send(json.dumps(init_msg))

        # Wait for context_created confirmation (no session_started event)
        while True:
            msg = await asyncio.wait_for(
                self._ws.recv(), timeout=_DEFAULT_RECV_TIMEOUT_S
            )
            data = json.loads(msg)

            if data.get("error"):
                raise KugelAudioError(data["error"])

            if data.get("context_created"):
                self._contexts.add(data["context_id"])
                self._is_started = True
                logger.debug(
                    "Multi-context session started, first context: %s",
                    data["context_id"],
                )
                break

    async def create_context(
        self,
        context_id: str,
        voice_id: Optional[int] = None,
    ) -> None:
        """Create a new context with optional voice override.

        Args:
            context_id: Unique identifier for this context
            voice_id: Optional voice ID (uses default if not specified)
        """
        if not self._is_started:
            await self._start_session(context_id, voice_id)
            return

        if context_id in self._contexts:
            return  # Already exists

        # Send context initialization — voice_id must be in voice_settings
        effective_voice_id = voice_id or self._default_voice_id
        msg: Dict[str, Any] = {"text": " ", "context_id": context_id}
        if effective_voice_id is not None:
            msg["voice_settings"] = {
                "voice_id": effective_voice_id,
                "cfg_scale": self._cfg_scale,
                "max_new_tokens": self._max_new_tokens,
            }

        await self._ws.send(json.dumps(msg))

        # Wait for confirmation
        while True:
            response = await self._ws.recv()
            data = json.loads(response)

            if data.get("error"):
                raise KugelAudioError(data["error"])

            if data.get("context_created") and data.get("context_id") == context_id:
                self._contexts.add(context_id)
                break

    async def send(
        self,
        context_id: str,
        text: str,
        flush: bool = False,
    ) -> AsyncIterator[AudioChunk]:
        """Send text to a specific context and yield audio chunks.

        Args:
            context_id: Context to send text to
            text: Text to synthesize
            flush: Force flush the buffer

        Yields:
            AudioChunk as audio is generated
        """
        if not self._is_started:
            await self._start_session(context_id)

        if context_id not in self._contexts:
            await self.create_context(context_id)

        await self._ws.send(
            json.dumps(
                {
                    "text": text,
                    "context_id": context_id,
                    "flush": flush,
                }
            )
        )

        # Collect audio for this context.
        # Note: flush only triggers chunk_complete on the server, NOT is_final.
        # is_final is only sent when a context is closed via close_context.
        async for chunk in self._receive_audio(context_id):
            yield chunk

    async def flush(self, context_id: str) -> AsyncIterator[AudioChunk]:
        """Flush a specific context's buffer.

        Args:
            context_id: Context to flush

        Yields:
            AudioChunk as remaining audio is generated
        """
        if not self._is_started or context_id not in self._contexts:
            return

        await self._ws.send(
            json.dumps(
                {
                    "flush": True,
                    "context_id": context_id,
                }
            )
        )

        # flush only triggers chunk_complete, not is_final
        async for chunk in self._receive_audio(context_id):
            yield chunk

    async def close_context(self, context_id: str) -> AsyncIterator[AudioChunk]:
        """Close a specific context and get remaining audio.

        Args:
            context_id: Context to close

        Yields:
            AudioChunk for any remaining buffered text
        """
        if not self._is_started or context_id not in self._contexts:
            return

        await self._ws.send(
            json.dumps(
                {
                    "close_context": True,
                    "context_id": context_id,
                }
            )
        )

        async for chunk in self._receive_audio(context_id, wait_for_close=True):
            yield chunk

        self._contexts.discard(context_id)

    async def keep_alive(self, context_id: str) -> None:
        """Reset inactivity timeout for a context.

        Args:
            context_id: Context to keep alive
        """
        if self._is_started and context_id in self._contexts:
            await self._ws.send(
                json.dumps(
                    {
                        "text": "",
                        "context_id": context_id,
                    }
                )
            )

    async def _receive_audio(
        self,
        context_id: str,
        wait_for_final: bool = False,
        wait_for_close: bool = False,
    ) -> AsyncIterator[AudioChunk]:
        """Receive audio messages for a specific context.

        The server sends (per context):
          - {"generation_started": true, "context_id": "...", ...}
          - {"audio": "<b64>", "context_id": "...", ...}
          - {"chunk_complete": true, "context_id": "...", ...}
          - {"is_final": true, "context_id": "..."}   (after flush/close)
          - {"context_closed": true, "context_id": "..."}
        """
        received_any_audio = False

        while True:
            # Use longer timeout when waiting for a definitive server signal
            if wait_for_final or wait_for_close:
                timeout = _DEFAULT_RECV_TIMEOUT_S
            else:
                timeout = (
                    _POLL_TIMEOUT_S if received_any_audio else _DEFAULT_RECV_TIMEOUT_S
                )

            try:
                msg = await asyncio.wait_for(self._ws.recv(), timeout=timeout)
                data = json.loads(msg)

                if data.get("error"):
                    raise KugelAudioError(data["error"])

                msg_ctx = data.get("context_id")

                # Skip messages for other contexts — stash if needed later
                if msg_ctx != context_id:
                    self._pending_messages.append(data)
                    continue

                if data.get("generation_started"):
                    received_any_audio = False
                    continue

                # Only yield audio for matching context
                if data.get("audio"):
                    received_any_audio = True
                    yield AudioChunk.from_dict(data)

                # chunk_complete = one text-chunk done
                if data.get("chunk_complete"):
                    if not wait_for_final and not wait_for_close:
                        break
                    continue

                # is_final = all queued chunks done for this context
                if data.get("is_final"):
                    if wait_for_final:
                        break

                # context_closed = context removed
                if data.get("context_closed"):
                    if wait_for_close:
                        break

            except asyncio.TimeoutError:
                # Short poll timed out — nothing more right now
                break
            except Exception as e:
                if "ConnectionClosed" in str(type(e)):
                    break
                raise

    async def close(self) -> Dict[str, Any]:
        """Close the session and return stats.

        Returns:
            Session statistics including total audio generated
        """
        stats = {}

        if self._ws:
            try:
                await self._ws.send(json.dumps({"close_socket": True}))

                # Wait for session stats
                msg = await asyncio.wait_for(self._ws.recv(), timeout=5.0)
                data = json.loads(msg)

                if data.get("session_closed"):
                    stats = data

            except Exception:
                pass
            finally:
                await self._ws.close()
                self._ws = None
                self._is_started = False
                self._contexts.clear()

        return stats

    @property
    def active_contexts(self) -> set[str]:
        """Get the set of active context IDs."""
        return self._contexts.copy()

    @property
    def session_id(self) -> Optional[str]:
        """Get the session ID."""
        return self._session_id
