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

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)
            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 to get raw PCM
            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(chunk)
                    await asyncio.sleep(0.005)
            
            output_emitter.flush()
            print(f"[EdgeTTS] Audio pushed and flushed successfully!", flush=True)

        except Exception as e:
            print(f"[EdgeTTS] Error in synthesis: {e}", file=sys.stderr)
            traceback.print_exc()
            # Fallback to local MMS-TTS if EdgeTTS fails
            try:
                print("[EdgeTTS] Falling back to local MMS-TTS...", flush=True)
                from tts_mms import MmsTTS
                fallback = MmsTTS("facebook/mms-tts-swh")
                fallback_stream = fallback.synthesize(self._input_text)
                async for audio_chunk in fallback_stream:
                    output_emitter.push(audio_chunk.frame.data.tobytes())
                output_emitter.flush()
            except Exception as fe:
                print(f"[EdgeTTS] Fallback error: {fe}", file=sys.stderr)
                raise