from __future__ import annotations

import asyncio

from livekit.agents import tts, utils
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS, APIConnectOptions

from app.providers.tts import CHANNELS, SAMPLE_RATE
from app.providers.tts import FakeTTS as FakeProviderTTS


class FakeTTS(tts.TTS):
    def __init__(self) -> None:
        super().__init__(
            capabilities=tts.TTSCapabilities(streaming=False, aligned_transcript=False),
            sample_rate=SAMPLE_RATE,
            num_channels=CHANNELS,
        )
        self._provider_client = FakeProviderTTS()

    @property
    def model(self) -> str:
        return "fake-tone"

    @property
    def provider(self) -> str:
        return "fake"

    def synthesize(
        self,
        text: str,
        *,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> ChunkedStream:
        return ChunkedStream(tts=self, input_text=text, conn_options=conn_options)

    async def aclose(self) -> None:
        return None


class ChunkedStream(tts.ChunkedStream):
    def __init__(
        self,
        *,
        tts: FakeTTS,
        input_text: str,
        conn_options: APIConnectOptions,
    ) -> None:
        super().__init__(tts=tts, input_text=input_text, conn_options=conn_options)
        self._tts: FakeTTS = tts

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        output_emitter.initialize(
            request_id=utils.shortuuid(),
            sample_rate=self._tts.sample_rate,
            num_channels=self._tts.num_channels,
            mime_type="audio/pcm",
        )

        audio = await asyncio.to_thread(
            self._tts._provider_client.synthesize_pcm16k,
            self._input_text,
        )
        output_emitter.push(audio)
        output_emitter.flush()
