import asyncio
import uuid
import sys
import traceback
import io
import edge_tts
from pydub import AudioSegment
from livekit.agents import tts
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS

class EdgeTTS(tts.TTS):
    def __init__(self, voice_id: str):
        super().__init__(
            capabilities=tts.TTSCapabilities(streaming=False),
            sample_rate=48000,
            num_channels=1,
        )
        self.voice_id = voice_id
        print(f"[EdgeTTS] Initialized with voice '{voice_id}'", flush=True)

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


class EdgeStream(tts.ChunkedStream):
    def __init__(self, *, tts: "EdgeTTS", input_text: str, conn_options=DEFAULT_API_CONNECT_OPTIONS):
        super().__init__(tts=tts, input_text=input_text, conn_options=conn_options)
        self._tts = tts
        self._input_text = input_text

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        try:
            print(f"[EdgeTTS] Generating for: '{self._input_text[:80]}'", flush=True)
            
            # Generate Edge TTS audio directly to memory (mp3)
            communicate = edge_tts.Communicate(self._input_text, self._tts.voice_id)
            audio_data = b""
            async for chunk in communicate.stream():
                if chunk["type"] == "audio":
                    audio_data += chunk["data"]
            
            # Decode using pydub
            audio_seg = AudioSegment.from_file(io.BytesIO(audio_data), format="mp3")
            # Convert to 48kHz, mono, 16-bit PCM
            audio_seg = audio_seg.set_frame_rate(48000).set_channels(1).set_sample_width(2)
            
            pcm_bytes = audio_seg.raw_data
            
            print(f"[EdgeTTS] PCM ready: {len(pcm_bytes)} bytes", flush=True)

            output_emitter.initialize(
                request_id=str(uuid.uuid4()),
                sample_rate=48000,
                num_channels=1,
                mime_type="audio/pcm",
            )
            
            # Chunk the output (48kHz * 20ms = 960 samples per chunk * 2 bytes = 1920 bytes)
            chunk_size = 1920
            for i in range(0, len(pcm_bytes), chunk_size):
                chunk = pcm_bytes[i:i+chunk_size]
                if len(chunk) > 0:
                    output_emitter.push_audio(chunk)
            
        except Exception as e:
            print(f"[EdgeTTS] Error in synthesis: {e}", file=sys.stderr)
            traceback.print_exc()