"""High-level interface to Inworld's TTS streaming API.

Implements a similar interface to OpenAITTS for interchangeability.
"""

from __future__ import annotations

import base64
import json
import os
import wave
from contextlib import asynccontextmanager
from enum import Enum
from typing import AsyncGenerator, Optional

import httpx
from openai.helpers import LocalAudioPlayer
from pydantic import BaseModel, Field

from narrate_gpt.tts import AudioFormat


class AudioEncoding(str, Enum):
    """Supported audio encodings for Inworld TTS."""

    LINEAR16 = "LINEAR16"


class InworldTTSConfig(BaseModel):
    """Configuration for Inworld TTS."""

    api_key: Optional[str] = Field(
        default=None,
        description="Inworld API key (falls back to INWORLD_API_KEY env var)",
    )
    base_url: str = Field(
        default="https://api.inworld.ai/tts/v1/voice:stream",
        description="Inworld TTS streaming endpoint",
    )
    model_id: str = Field(
        default="inworld-tts-1.5-max", description="Inworld TTS model id"
    )
    voice_id: str = Field(default="Mark", description="Inworld voice id")
    audio_encoding: AudioEncoding = Field(
        default=AudioEncoding.LINEAR16,
        description="Audio encoding (LINEAR16 recommended)",
    )
    sample_rate_hertz: int = Field(
        default=24000,
        description="Sample rate for LINEAR16 audio (use 24000 for LocalAudioPlayer)",
    )
    channels: int = Field(default=1, description="Number of audio channels")


class _InworldStreamResponse:
    """Minimal streamed response interface compatible with LocalAudioPlayer."""

    def __init__(
        self,
        generator: AsyncGenerator[bytes, None],
    ) -> None:
        self._generator = generator

    async def iter_bytes(self, chunk_size: int = 1024) -> AsyncGenerator[bytes, None]:
        buffer = bytearray()
        async for chunk in self._generator:
            if not chunk:
                continue
            buffer.extend(chunk)
            while len(buffer) >= chunk_size:
                yield bytes(buffer[:chunk_size])
                del buffer[:chunk_size]
        if buffer:
            yield bytes(buffer)

    async def read(self) -> bytes:
        return b"".join([chunk async for chunk in self.iter_bytes()])


class InworldTTS:
    """Args:
    config: InworldTTSConfig object with TTS settings.
    """

    def __init__(self, config: Optional[InworldTTSConfig] = None) -> None:
        self.config = config or InworldTTSConfig()

    def _get_api_key(self) -> str:
        api_key = self.config.api_key or os.getenv("INWORLD_API_KEY")
        if not api_key:
            raise ValueError("Missing Inworld API key (set INWORLD_API_KEY)")
        return api_key

    @asynccontextmanager
    async def request(
        self,
        text: str,
        model: str | None = None,
        voice: str | None = None,
        format: AudioFormat | None = None,
        audio_encoding: AudioEncoding | None = None,
        sample_rate_hertz: int | None = None,
    ):
        if format is not None and format not in {AudioFormat.PCM, AudioFormat.WAV}:
            raise ValueError("Inworld TTS only supports PCM/WAV output")

        resolved_encoding = audio_encoding or self.config.audio_encoding
        resolved_sample_rate = sample_rate_hertz or self.config.sample_rate_hertz

        payload = {
            "text": text,
            "voice_id": voice or self.config.voice_id,
            "model_id": model or self.config.model_id,
            "audio_config": {
                "audio_encoding": resolved_encoding.value,
                "sample_rate_hertz": resolved_sample_rate,
            },
        }

        headers = {
            "Authorization": f"Basic {self._get_api_key()}",
            "Content-Type": "application/json",
            "Connection": "keep-alive",
        }

        async with httpx.AsyncClient(timeout=None) as client:
            async with client.stream(
                "POST",
                self.config.base_url,
                json=payload,
                headers=headers,
            ) as response:
                response.raise_for_status()

                async def pcm_generator() -> AsyncGenerator[bytes, None]:
                    async for line in response.aiter_lines():
                        if not line:
                            continue
                        try:
                            chunk = json.loads(line)
                        except json.JSONDecodeError:
                            continue
                        result = chunk.get("result") or {}
                        audio_content = result.get("audioContent")
                        if not audio_content:
                            continue
                        audio_chunk = base64.b64decode(audio_content)
                        if len(audio_chunk) > 44:
                            yield audio_chunk[44:]

                yield _InworldStreamResponse(pcm_generator())

    async def save(
        self,
        text: str,
        path: str,
        model: str | None = None,
        voice: str | None = None,
        format: AudioFormat | None = None,
        audio_encoding: AudioEncoding | None = None,
        sample_rate_hertz: int | None = None,
    ) -> None:
        """Convert text to speech and save it to a WAV file."""
        resolved_sample_rate = sample_rate_hertz or self.config.sample_rate_hertz

        async with self.request(
            text=text,
            model=model,
            voice=voice,
            format=format,
            audio_encoding=audio_encoding,
            sample_rate_hertz=resolved_sample_rate,
        ) as response:
            pcm_bytes = await response.read()

        with wave.open(path, "wb") as wf:
            wf.setnchannels(self.config.channels)
            wf.setsampwidth(2)
            wf.setframerate(resolved_sample_rate)
            wf.writeframes(pcm_bytes)

    async def play(
        self,
        text: str,
        model: str | None = None,
        voice: str | None = None,
        audio_encoding: AudioEncoding | None = None,
        sample_rate_hertz: int | None = None,
    ) -> None:
        """Convert text to speech and play it using LocalAudioPlayer.

        Note: LocalAudioPlayer uses a fixed 24kHz sample rate. For accurate
        playback, set sample_rate_hertz to 24000.
        """
        resolved_sample_rate = sample_rate_hertz or self.config.sample_rate_hertz
        if resolved_sample_rate != 24000:
            raise ValueError(
                "LocalAudioPlayer only supports 24000Hz audio; "
                "set sample_rate_hertz=24000"
            )

        async with self.request(
            text=text,
            model=model,
            voice=voice,
            format=AudioFormat.PCM,
            audio_encoding=audio_encoding,
            sample_rate_hertz=resolved_sample_rate,
        ) as response:
            await LocalAudioPlayer().play(response)


if __name__ == "__main__":
    import asyncio

    tts = InworldTTS()
    asyncio.run(tts.save("Hello from Inworld TTS!", "./inworld.wav"))
