"""SSE stream parsing and audio collection utilities."""

from __future__ import annotations

import json
from collections.abc import AsyncIterator, Iterator
from contextlib import asynccontextmanager, contextmanager
from typing import TYPE_CHECKING, Any

from pydantic import ValidationError

from murmr._audio import create_wav_header
from murmr._errors import MurmrError
from murmr._types import AudioStreamChunk

if TYPE_CHECKING:
    import httpx


def _parse_sse_line(line: str) -> AudioStreamChunk | None:
    """Parse a single SSE line into an AudioStreamChunk, or None if not a data line."""
    trimmed = line.strip()
    if not trimmed or not trimmed.startswith("data: "):
        return None
    payload = trimmed[6:]
    if payload == "[DONE]":
        return None
    try:
        data: dict[str, Any] = json.loads(payload)
        return AudioStreamChunk.model_validate(data)
    except (json.JSONDecodeError, ValidationError):
        return None


@contextmanager
def sync_stream(response: httpx.Response) -> Iterator[Iterator[AudioStreamChunk]]:
    """Context manager that yields a sync iterator of AudioStreamChunk from an SSE response."""

    def _iter_chunks() -> Iterator[AudioStreamChunk]:
        for line in response.iter_lines():
            chunk = _parse_sse_line(line)
            if chunk is not None:
                yield chunk

    try:
        yield _iter_chunks()
    finally:
        response.close()


@asynccontextmanager
async def async_stream(
    response: httpx.Response,
) -> AsyncIterator[AsyncIterator[AudioStreamChunk]]:
    """Context manager that yields an async iterator of AudioStreamChunk from an SSE response."""

    async def _iter_chunks() -> AsyncIterator[AudioStreamChunk]:
        async for line in response.aiter_lines():
            chunk = _parse_sse_line(line)
            if chunk is not None:
                yield chunk

    try:
        yield _iter_chunks()
    finally:
        await response.aclose()


def collect_stream_as_wav_sync(response: httpx.Response) -> bytes:
    """Consume a sync SSE stream and return a complete WAV buffer."""
    pcm_parts: list[bytes] = []

    for line in response.iter_lines():
        chunk = _parse_sse_line(line)
        if chunk is not None:
            if chunk.error:
                response.close()
                raise MurmrError(chunk.error)
            audio_data = chunk.audio_bytes
            if audio_data:
                pcm_parts.append(audio_data)

    response.close()

    if not pcm_parts:
        return b""

    pcm = b"".join(pcm_parts)
    return create_wav_header(len(pcm)) + pcm


async def collect_stream_as_wav_async(response: httpx.Response) -> bytes:
    """Consume an async SSE stream and return a complete WAV buffer."""
    pcm_parts: list[bytes] = []

    async for line in response.aiter_lines():
        chunk = _parse_sse_line(line)
        if chunk is not None:
            if chunk.error:
                await response.aclose()
                raise MurmrError(chunk.error)
            audio_data = chunk.audio_bytes
            if audio_data:
                pcm_parts.append(audio_data)

    await response.aclose()

    if not pcm_parts:
        return b""

    pcm = b"".join(pcm_parts)
    return create_wav_header(len(pcm)) + pcm


def collect_stream_as_pcm_sync(response: httpx.Response) -> bytes:
    """Consume a sync SSE stream and return raw PCM bytes (no WAV header)."""
    pcm_parts: list[bytes] = []

    for line in response.iter_lines():
        chunk = _parse_sse_line(line)
        if chunk is not None:
            if chunk.error:
                response.close()
                raise MurmrError(chunk.error)
            audio_data = chunk.audio_bytes
            if audio_data:
                pcm_parts.append(audio_data)

    response.close()
    return b"".join(pcm_parts) if pcm_parts else b""


async def collect_stream_as_pcm_async(response: httpx.Response) -> bytes:
    """Consume an async SSE stream and return raw PCM bytes (no WAV header)."""
    pcm_parts: list[bytes] = []

    async for line in response.aiter_lines():
        chunk = _parse_sse_line(line)
        if chunk is not None:
            if chunk.error:
                await response.aclose()
                raise MurmrError(chunk.error)
            audio_data = chunk.audio_bytes
            if audio_data:
                pcm_parts.append(audio_data)

    await response.aclose()
    return b"".join(pcm_parts) if pcm_parts else b""
