import os
import io
import wave
import soundfile as sf
import numpy as np
import asyncio
import sys
import traceback
from typing import AsyncIterable
from livekit.agents import tts
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS
import uuid
from gradio_client import Client

# Fallback language mapping: API voice -> MMS model for fallback
FALLBACK_MMS = {
    "stano03/habari-xtts": "facebook/mms-tts-swh",
    "sta-03/habari-f5tts": "facebook/mms-tts-swh",
    "Stanley03/habari-yarngpt": "facebook/mms-tts-yor",
    "Stanley-07/habari-nursetoto": "facebook/mms-tts-swh",
    "Stanley-07/habari-kisii": "facebook/mms-tts-swh",
}


class ApiTTS(tts.TTS):
    def __init__(self, space_id: str, sample_rate: int = 24000, **kwargs):
        super().__init__(
            capabilities=tts.TTSCapabilities(streaming=False),
            sample_rate=sample_rate,
            num_channels=1,
        )
        self.space_id = space_id
        self.kwargs = kwargs
        self.client = None
        self._fallback_tts = None  # lazy MMS fallback
        self._api_failed = False   # track if API is down
        print(f"[ApiTTS] Initializing for: {space_id}", flush=True)

    def _get_fallback(self):
        """Lazy-load an MMS-TTS fallback engine."""
        if self._fallback_tts is None:
            from tts_mms import MmsTTS
            fallback_model = FALLBACK_MMS.get(self.space_id, "facebook/mms-tts-swh")
            print(f"[ApiTTS] Loading MMS fallback: {fallback_model}", flush=True)
            self._fallback_tts = MmsTTS(model_id=fallback_model)
        return self._fallback_tts

    def synthesize(self, text: str, *, conn_options=DEFAULT_API_CONNECT_OPTIONS) -> tts.ChunkedStream:
        # If we know the API is down, go straight to fallback
        if self._api_failed:
            print(f"[ApiTTS] API known down, using MMS fallback", flush=True)
            return self._get_fallback().synthesize(text, conn_options=conn_options)
        return _ApiChunkedStream(tts_engine=self, text=text, conn_options=conn_options)


class _ApiChunkedStream(tts.ChunkedStream):
    def __init__(self, *, tts_engine: ApiTTS, text: str, conn_options=DEFAULT_API_CONNECT_OPTIONS):
        super().__init__(tts=tts_engine, input_text=text, conn_options=conn_options)
        self.text = text
        self.tts_engine = tts_engine

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        try:
            # 1. Initialize client if not already done (with timeout)
            if self.tts_engine.client is None:
                print(f"[ApiTTS] Connecting to {self.tts_engine.space_id}...", flush=True)
                try:
                    self.tts_engine.client = await asyncio.wait_for(
                        asyncio.to_thread(Client, self.tts_engine.space_id),
                        timeout=120.0
                    )
                    print(f"[ApiTTS] Connected OK", flush=True)
                except (asyncio.TimeoutError, Exception) as conn_err:
                    print(f"[ApiTTS] Connection failed: {conn_err}", flush=True)
                    self.tts_engine._api_failed = True
                    # Use fallback
                    await self._run_fallback(output_emitter)
                    return

            # 2. Call the API with timeout and retry for "models still loading"
            print(f"[ApiTTS] Generating: '{self.text[:60]}'", flush=True)
            kwargs = dict(self.tts_engine.kwargs)
            kwargs["text"] = self.text

            result = None
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    result = await asyncio.wait_for(
                        asyncio.to_thread(
                            self.tts_engine.client.predict,
                            api_name="/generate",
                            **kwargs
                        ),
                        timeout=180.0
                    )
                    break  # Success
                except Exception as api_err:
                    err_msg = str(api_err)
                    if "still loading" in err_msg and attempt < max_retries - 1:
                        wait_time = 15 * (attempt + 1)
                        print(f"[ApiTTS] Models loading, retry {attempt+1}/{max_retries} in {wait_time}s...", flush=True)
                        await asyncio.sleep(wait_time)
                        continue
                    print(f"[ApiTTS] API call failed: {api_err}", flush=True)
                    self.tts_engine._api_failed = True
                    self.tts_engine.client = None
                    await self._run_fallback(output_emitter)
                    return

            audio_path = result
            print(f"[ApiTTS] Got audio: {audio_path}", flush=True)

            # 3. Read and convert audio
            data, sample_rate = sf.read(audio_path, dtype='float32')
            if len(data.shape) > 1:
                data = data.mean(axis=1)

            pcm_int16 = (data * 32767).clip(-32768, 32767).astype(np.int16)
            pcm_bytes = pcm_int16.tobytes()
            print(f"[ApiTTS] PCM: {len(pcm_int16)} samples @ {sample_rate}Hz", flush=True)

            # 4. Push audio
            output_emitter.initialize(
                request_id=str(uuid.uuid4()),
                sample_rate=sample_rate,
                num_channels=1,
                mime_type="audio/pcm"
            )
            output_emitter.push(pcm_bytes)
            output_emitter.flush()
            print(f"[ApiTTS] Audio pushed OK", flush=True)

        except Exception as e:
            print(f"[ApiTTS] Unexpected error: {e}", flush=True)
            traceback.print_exc()
            sys.stdout.flush()
            # Try fallback as last resort
            try:
                await self._run_fallback(output_emitter)
            except Exception as fb_err:
                print(f"[ApiTTS] Fallback also failed: {fb_err}", flush=True)
                raise

    async def _run_fallback(self, output_emitter: tts.AudioEmitter) -> None:
        """Use MMS-TTS fallback when API is unavailable."""
        print(f"[ApiTTS] Using MMS fallback for: '{self.text[:60]}'", flush=True)
        fallback = self.tts_engine._get_fallback()
        # Create MMS stream and directly call its _run with our emitter
        from tts_mms import MmsStream
        fallback_stream = MmsStream(tts=fallback, input_text=self.text)
        await fallback_stream._run(output_emitter)