import torch
import numpy as np
import asyncio
import uuid
import sys
import traceback
from transformers import VitsModel, AutoTokenizer
from scipy.signal import resample
from livekit.agents import tts
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS

class MmsTTS(tts.TTS):
    def __init__(self, model_id: str):
        super().__init__(
            capabilities=tts.TTSCapabilities(streaming=False),
            sample_rate=48000,
            num_channels=1,
        )
        print(f"Loading MMS-TTS model '{model_id}' (first time may take a minute)...", flush=True)
        self._model = VitsModel.from_pretrained(model_id)
        self._tokenizer = AutoTokenizer.from_pretrained(model_id)
        self._orig_rate = self._model.config.sampling_rate
        print(f"Voice ready -> {model_id} (native rate: {self._orig_rate}Hz)", flush=True)

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


class MmsStream(tts.ChunkedStream):
    def __init__(self, *, tts: "MmsTTS", 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"[MmsTTS] _run called for: '{self._input_text[:80]}'", flush=True)
            
            # Generate model input
            inputs = self._tts._tokenizer(self._input_text, return_tensors="pt")

            def _generate():
                with torch.no_grad():
                    return self._tts._model(**inputs).waveform[0].numpy()

            print(f"[MmsTTS] Starting synthesis...", flush=True)
            waveform = await asyncio.to_thread(_generate)
            print(f"[MmsTTS] Waveform generated: {waveform.shape}, range [{waveform.min():.3f}, {waveform.max():.3f}]", flush=True)

            # Resample original rate → 48 kHz
            num_samples = int(len(waveform) * 48000 / self._tts._orig_rate)
            waveform_48k = resample(waveform, num_samples).astype(np.float32)

            # Convert to 16-bit PCM bytes
            pcm_int16 = (waveform_48k * 32767).clip(-32768, 32767).astype(np.int16)
            pcm_bytes = pcm_int16.tobytes()

            print(f"[MmsTTS] PCM ready: {len(pcm_int16)} samples @ 48000Hz, {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]
                output_emitter.push(chunk)
                await asyncio.sleep(0.005)  # small yield to event loop
                
            output_emitter.flush()
            print(f"[MmsTTS] Audio pushed in chunks and flushed successfully!", flush=True)
            
        except Exception as e:
            print(f"[MmsTTS] ERROR in _run: {e}", flush=True)
            traceback.print_exc()
            sys.stdout.flush()
            raise  # Re-raise so LiveKit knows TTS failed