#!/usr/bin/env python3
import argparse
import asyncio
import contextlib
import gc
import os
import queue
import re
import subprocess
import sys
import threading
import time
import yaml
import json

from fastapi.responses import StreamingResponse
from loguru import logger
from openedai import OpenAIStub, BadRequestError, ServiceUnavailableError
from pydantic import BaseModel
from typing import Optional
import uvicorn

# Try to import XTTS dependencies (might not be available in minimal installations)
try:
    import torch
    from TTS.tts.configs.xtts_config import XttsConfig
    from TTS.tts.models.xtts import Xtts
    from TTS.utils.manage import ModelManager
    from TTS.tts.layers.xtts.tokenizer import split_sentence
    from langdetect import detect
    XTTS_AVAILABLE = True
except ImportError:
    XTTS_AVAILABLE = False
    torch = None
    XttsConfig = None
    Xtts = None
    ModelManager = None
    split_sentence = None
    detect = None

# Try to import Qwen3-TTS dependencies
try:
    import torch
    from qwen_tts import Qwen3TTSModel
    QWEN_TTS_AVAILABLE = True
except ImportError:
    QWEN_TTS_AVAILABLE = False
    if torch is None:
        try:
            import torch
        except ImportError:
            torch = None
    Qwen3TTSModel = None

@contextlib.asynccontextmanager
async def lifespan(app):
    # Startup: Initialize voice caches in each worker process
    global voice_to_model_cache, voices_cache

    default_exists('config/pre_process_map.yaml')
    default_exists('config/voice_to_speaker.yaml')

    # Build voice-to-model cache for fast lookups
    with open('config/voice_to_speaker.yaml', 'r', encoding='utf8') as file:
        voice_map = yaml.safe_load(file)
        for model_id, voices in voice_map.items():
            if isinstance(voices, dict):
                for voice_name in voices.keys():
                    # First match wins (for duplicate voice names across models)
                    if voice_name not in voice_to_model_cache:
                        voice_to_model_cache[voice_name] = model_id
    print(f"Voice-to-model cache initialized with {len(voice_to_model_cache)} voices")

    # Build voices cache for /v1/voices endpoint
    models_data = []
    for model_id, voices in voice_map.items():
        if isinstance(voices, dict):
            voice_list = list(voices.keys())

            model_info = {
                "id": model_id,
                "object": "model",
                "created": 1700000000,
                "owned_by": "uncloseai",
                "voices": voice_list,
                "voice_count": len(voice_list)
            }

            # Add engine-specific metadata
            if model_id == 'tts-1':
                model_info["engine"] = "piper"
                model_info["description"] = "Fast neural TTS with 100+ voices"
                model_info["sample_rate"] = 22050
            elif model_id == 'tts-1-hd':
                model_info["engine"] = "xtts"
                model_info["description"] = "High-quality voice cloning TTS"
                model_info["sample_rate"] = 24000
            elif model_id == 'tts-1-silero':
                model_info["engine"] = "silero"
                model_info["description"] = "Fast multilingual TTS (en, ru, de, es, fr)"
                model_info["sample_rate"] = 48000
            elif model_id == 'tts-1-kokoro':
                model_info["engine"] = "kokoro"
                model_info["description"] = "Lightweight decoder-only TTS (82M params)"
                model_info["sample_rate"] = 24000
            elif model_id == 'tts-1-qwen':
                model_info["engine"] = "qwen3-tts"
                model_info["description"] = "State-of-the-art TTS with voice cloning (1.7B params, 10 languages)"
                model_info["sample_rate"] = 24000

            models_data.append(model_info)

    voices_cache = {
        "object": "list",
        "data": models_data
    }
    print(f"/v1/voices cache initialized with {len(models_data)} models")

    yield

    # Shutdown: Cleanup
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    except:
        pass

# We return 'mps' but currently XTTS will not work with mps devices as the cuda support is incomplete
def auto_torch_device():
    try:
        import torch
        return 'cuda' if torch.cuda.is_available() else 'mps' if ( torch.backends.mps.is_available() and torch.backends.mps.is_built() ) else 'cpu'

    except:
        return 'none'

app = OpenAIStub(lifespan=lifespan)
xtts = None
silero_model = None
silero_speakers = {}
kokoro_pipeline = None
kokoro_lang = None
qwen_model = None
qwen_voice_prompts = {}  # Cache for voice clone prompts

# Default args for worker processes (will be overridden in __main__)
class DefaultArgs:
    xtts_device = None  # Will be set after torch import
    use_deepspeed = False
    unload_timer = None
    log_level = 'INFO'
    host = '0.0.0.0'
    port = 8000
    preload = None
    no_cache_speaker = False

args = DefaultArgs()

# Set default device after torch is available (for worker processes)
# Main process will override this in __main__ with argparse
if args.xtts_device is None:
    try:
        detected_device = auto_torch_device()
        args.xtts_device = detected_device
        logger.debug(f"Worker process initialized with device: {detected_device}")
    except Exception as e:
        logger.warning(f"Failed to detect torch device: {e}, falling back to CPU")
        args.xtts_device = 'cpu'

# Voice-to-model lookup cache (loaded at startup)
voice_to_model_cache = {}

# Cached voice data for /v1/voices endpoint (loaded at startup)
voices_cache = None

# Semaphores to limit concurrent model loading (prevent thread pool exhaustion)
silero_load_semaphore = asyncio.Semaphore(1)  # Only one Silero model load at a time
kokoro_load_semaphore = asyncio.Semaphore(1)  # Only one Kokoro model load at a time
qwen_load_semaphore = asyncio.Semaphore(1)    # Only one Qwen model load at a time

def unload_model():
    import torch, gc
    global xtts
    if xtts:
        logger.info("Unloading model")
        xtts.xtts.to('cpu') # this was required to free up GPU memory... 
        del xtts
        xtts = None
        gc.collect()
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()

class xtts_wrapper():
    check_interval: int = 1 # too aggressive?

    def __init__(self, model_name, device, model_path=None, unload_timer=None):
        self.model_name = model_name
        self.unload_timer = unload_timer
        self.last_used = time.time()
        self.timer = None
        self.lock = threading.Lock()

        logger.info(f"Loading model {self.model_name} to {device}")

        if model_path is None:
            model_path = ModelManager().download_model(model_name)[0]

        config_path = os.path.join(model_path, 'config.json')
        config = XttsConfig()
        config.load_json(config_path)
        self.xtts = Xtts.init_from_config(config)
        self.xtts.load_checkpoint(config, checkpoint_dir=model_path, use_deepspeed=args.use_deepspeed)  # XXX there are no prebuilt deepspeed wheels??
        self.xtts = self.xtts.to(device=device)
        self.xtts.eval()

        if self.unload_timer:
            logger.info(f"Setting unload timer to {self.unload_timer} seconds")
            self.last_used = time.time()
            self.check_idle()

    def check_idle(self):
        with self.lock:
            if time.time() - self.last_used >= self.unload_timer:
                print("Unloading TTS model due to inactivity")
                unload_model()
            else:
                # Reschedule the check
                self.timer = threading.Timer(self.check_interval, self.check_idle)
                self.timer.daemon = True
                self.timer.start()

    def tts(self, text, language, audio_path, **hf_generate_kwargs):
        with torch.no_grad():
            self.last_used = time.time()
            tokens = 0
            try:
                with self.lock:
                    logger.debug(f"generating [{language}]: {[text]}")

                    gpt_cond_latent, speaker_embedding = self.xtts.get_conditioning_latents(audio_path=audio_path) # not worth caching calls, it's < 0.001s after model is loaded
                    pcm_stream = self.xtts.inference_stream(text, language, gpt_cond_latent, speaker_embedding, **hf_generate_kwargs)
                    self.last_used = time.time()

                while True:
                    with self.lock:
                        yield next(pcm_stream).cpu().numpy().tobytes()
                        self.last_used = time.time()
                    tokens += 1

            except StopIteration:
                pass

            finally:
                logger.debug(f"Generated {tokens} tokens in {time.time() - self.last_used:.2f}s @ {tokens / (time.time() - self.last_used):.2f} T/s")
                self.last_used = time.time()

class silero_wrapper():
    """Wrapper for Silero TTS models

    Silero torch.hub.load returns: (model, example_text)
    The model has a method apply_tts(text, speaker, sample_rate)
    """
    def __init__(self, language='en', speaker='v3_en', device='cpu'):
        self.language = language
        self.speaker = speaker
        self.device = device

        logger.info(f"Loading Silero model for {language} with speaker {speaker} on {device}")

        import torch
        try:
            # torch.hub.load returns (model, example_text)
            self.model, example_text = torch.hub.load(
                repo_or_dir='snakers4/silero-models',
                model='silero_tts',
                language=language,
                speaker=speaker,
                verbose=False
            )

            logger.info(f"Model loaded, type: {type(self.model)}")
            self.model.to(device)  # Move to device (in-place for Silero)
            self.sample_rate = 48000  # Silero uses 48kHz
            logger.info(f"Successfully loaded Silero {language}/{speaker}, example: {example_text}")
        except Exception as e:
            logger.error(f"Failed to load Silero model: {e}")
            raise

    def tts(self, text, speaker_id='lj_16khz'):
        """Generate speech from text"""
        import torch

        logger.info(f"Silero tts() called: model={self.model}, speaker_id={speaker_id}")

        if self.model is None:
            raise RuntimeError("Silero model is None - model failed to load")

        if not hasattr(self.model, 'apply_tts'):
            logger.error(f"Model has no apply_tts method. Model type: {type(self.model)}, dir: {dir(self.model)}")
            raise AttributeError(f"Silero model {type(self.model)} has no apply_tts method")

        with torch.no_grad():
            # Use model's apply_tts method
            audio = self.model.apply_tts(
                text=text,
                speaker=speaker_id,
                sample_rate=self.sample_rate
            )
            # audio is a tensor, convert to numpy float32
            return audio.cpu().numpy().tobytes()

class kokoro_wrapper():
    """Wrapper for Kokoro TTS model

    Kokoro is a lightweight decoder-only TTS model (82M params)
    Output: 24kHz audio
    """
    def __init__(self, lang_code='a', device='cpu'):
        self.lang_code = lang_code
        self.device = device
        self.sample_rate = 24000  # Kokoro outputs 24kHz

        logger.info(f"Loading Kokoro TTS pipeline for language '{lang_code}' on device '{device}'")

        try:
            from kokoro import KPipeline
            import numpy as np

            # KPipeline will use default repo_id if not specified
            # Pass device to KPipeline (supports 'cpu' or 'cuda')
            self.pipeline = KPipeline(lang_code=lang_code, device=device)
            logger.info(f"Successfully loaded Kokoro pipeline for lang={lang_code} on {device}")
        except Exception as e:
            logger.error(f"Failed to load Kokoro model: {e}")
            raise

    def tts(self, text, voice='af_heart', speed=1.0):
        """Generate speech from text using Kokoro"""
        import numpy as np

        logger.info(f"Kokoro tts() called: text length={len(text)}, voice={voice}, speed={speed}")

        try:
            # Generate audio using Kokoro pipeline
            generator = self.pipeline(text, voice=voice, speed=speed)

            # Collect all audio chunks
            audio_chunks = []
            chunk_count = 0
            for _, _, audio in generator:
                if audio is not None and len(audio) > 0:
                    audio_chunks.append(audio)
                    chunk_count += 1
                    if chunk_count % 10 == 0:
                        logger.debug(f"Kokoro generated {chunk_count} chunks so far...")

            # Concatenate all chunks
            if len(audio_chunks) > 0:
                full_audio = np.concatenate(audio_chunks)
                logger.info(f"Kokoro generation complete: {chunk_count} chunks, {len(full_audio)} samples")
                # Convert float32 numpy array to bytes
                return full_audio.astype(np.float32).tobytes()
            else:
                logger.warning("Kokoro generated no audio")
                return b''

        except Exception as e:
            logger.error(f"Kokoro TTS generation failed: {e}")
            raise

class qwen3_wrapper():
    """Wrapper for Qwen3-TTS model

    Qwen3-TTS is a state-of-the-art TTS with voice cloning:
    - 1.7B parameters, 12Hz tokenizer
    - 10 languages: zh, en, ja, ko, de, fr, ru, pt, es, it
    - 97ms first-packet latency
    - 3-second rapid voice cloning
    Output: Variable sample rate (typically 24kHz)
    """
    def __init__(self, model_name='Qwen/Qwen3-TTS-12Hz-1.7B-Base', device='cuda', dtype=None):
        self.model_name = model_name
        self.device = device
        self.sample_rate = None  # Set after first generation
        self.voice_prompts = {}  # Cache for reusable voice clone prompts

        logger.info(f"Loading Qwen3-TTS model '{model_name}' on device '{device}'")

        try:
            import torch
            from qwen_tts import Qwen3TTSModel

            # Determine dtype
            if dtype is None:
                if device == 'cuda' and torch.cuda.is_available():
                    dtype = torch.bfloat16
                else:
                    dtype = torch.float32

            # Try to use flash attention if available
            try:
                self.model = Qwen3TTSModel.from_pretrained(
                    model_name,
                    device_map=device,
                    dtype=dtype,
                    attn_implementation="flash_attention_2",
                )
                logger.info(f"Loaded Qwen3-TTS with FlashAttention 2")
            except Exception as fa_error:
                logger.warning(f"FlashAttention 2 not available ({fa_error}), using default attention")
                self.model = Qwen3TTSModel.from_pretrained(
                    model_name,
                    device_map=device,
                    dtype=dtype,
                )

            logger.info(f"Successfully loaded Qwen3-TTS model on {device}")
        except Exception as e:
            logger.error(f"Failed to load Qwen3-TTS model: {e}")
            raise

    def create_voice_prompt(self, ref_audio, ref_text, x_vector_only_mode=False):
        """Create a reusable voice clone prompt from reference audio/text.

        Args:
            ref_audio: Path to reference audio file, URL, or (numpy_array, sample_rate) tuple
            ref_text: Text spoken in the reference audio
            x_vector_only_mode: If True, use only speaker embedding (faster but lower quality)

        Returns:
            Voice clone prompt items for reuse
        """
        logger.info(f"Creating voice clone prompt from ref_audio={ref_audio}, ref_text={ref_text[:50]}...")
        return self.model.create_voice_clone_prompt(
            ref_audio=ref_audio,
            ref_text=ref_text,
            x_vector_only_mode=x_vector_only_mode,
        )

    def tts(self, text, language='English', ref_audio=None, ref_text=None, voice_prompt=None):
        """Generate speech from text using voice cloning.

        Args:
            text: Text to synthesize (string or list of strings)
            language: Language of the text
            ref_audio: Path/URL to reference audio (if voice_prompt not provided)
            ref_text: Text in reference audio (if voice_prompt not provided)
            voice_prompt: Pre-computed voice clone prompt (for efficiency)

        Returns:
            Audio data as bytes (float32 PCM)
        """
        import numpy as np

        logger.info(f"Qwen3-TTS generating: text length={len(text)}, language={language}")

        try:
            if voice_prompt is not None:
                # Use pre-computed voice prompt
                wavs, sr = self.model.generate_voice_clone(
                    text=text,
                    language=language,
                    voice_clone_prompt=voice_prompt,
                )
            elif ref_audio is not None and ref_text is not None:
                # Generate with inline reference
                wavs, sr = self.model.generate_voice_clone(
                    text=text,
                    language=language,
                    ref_audio=ref_audio,
                    ref_text=ref_text,
                )
            else:
                raise ValueError("Either voice_prompt or (ref_audio + ref_text) must be provided")

            self.sample_rate = sr
            logger.info(f"Qwen3-TTS generated {len(wavs)} audio segment(s) at {sr}Hz")

            # Concatenate all wav segments and convert to bytes
            if len(wavs) > 1:
                full_audio = np.concatenate(wavs)
            else:
                full_audio = wavs[0]

            return full_audio.astype(np.float32).tobytes()

        except Exception as e:
            logger.error(f"Qwen3-TTS generation failed: {e}")
            raise

def default_exists(filename: str):
    if not os.path.exists(filename):
        fpath, ext = os.path.splitext(filename)
        basename = os.path.basename(fpath)
        default = f"{basename}.default{ext}"
        
        logger.info(f"{filename} does not exist, setting defaults from {default}")

        with open(default, 'r', encoding='utf8') as from_file:
            with open(filename, 'w', encoding='utf8') as to_file:
                to_file.write(from_file.read())

# Read pre process map on demand so it can be changed without restarting the server
def preprocess(raw_input):
    #logger.debug(f"preprocess: before: {[raw_input]}")
    default_exists('config/pre_process_map.yaml')
    with open('config/pre_process_map.yaml', 'r', encoding='utf8') as file:
        pre_process_map = yaml.safe_load(file)
        for a, b in pre_process_map:
            raw_input = re.sub(a, b, raw_input)

    raw_input = raw_input.strip()
    #logger.debug(f"preprocess: after: {[raw_input]}")
    return raw_input

def simple_sentence_split(text: str, max_length: int = 500) -> list[str]:
    """Split text into sentences for streaming TTS.

    Splits on every sentence boundary (.!?) for immediate streaming.
    Long sentences exceeding max_length are split at word boundaries.
    """
    # Split on sentence boundaries, keeping the punctuation
    parts = re.split(r'([.!?]+\s*)', text)

    result = []
    for i in range(0, len(parts), 2):
        sentence = parts[i].strip()
        punct = parts[i+1] if i+1 < len(parts) else ""

        if not sentence:
            continue

        full_sentence = (sentence + punct).strip()
        if full_sentence:
            result.append(full_sentence)

    # Split any sentences that exceed max_length at word boundaries
    final_result = []
    for sentence in result:
        if len(sentence) <= max_length:
            final_result.append(sentence)
        else:
            # Split at word boundaries
            words = sentence.split()
            chunk = ""
            for word in words:
                if len(chunk) + len(word) + 1 > max_length:
                    if chunk:
                        final_result.append(chunk.strip())
                    chunk = word
                else:
                    chunk += " " + word if chunk else word
            if chunk:
                final_result.append(chunk.strip())

    return final_result if final_result else [text]

# Auto-detect which model a voice belongs to (uses cached mapping)
def detect_model_from_voice(voice: str) -> str:
    """Find which model supports a given voice name.
    Returns the first model that has this voice, or None if not found.
    Uses voice_to_model_cache populated at startup for fast lookups.
    """
    global voice_to_model_cache
    return voice_to_model_cache.get(voice, None)

# Read voice map on demand so it can be changed without restarting the server
def map_voice_to_speaker(voice: str, model: str):
    default_exists('config/voice_to_speaker.yaml')
    with open('config/voice_to_speaker.yaml', 'r', encoding='utf8') as file:
        voice_map = yaml.safe_load(file)
        try:
            return voice_map[model][voice]

        except KeyError as e:
            raise BadRequestError(f"Error loading voice: {voice}, KeyError: {e}", param='voice')

class GenerateSpeechRequest(BaseModel):
    model: Optional[str] = None  # Auto-detected from voice if not provided
    input: str
    voice: str = "alloy"  # alloy, echo, fable, onyx, nova, and shimmer
    response_format: str = "mp3" # mp3, opus, aac, flac
    speed: float = 1.0 # 0.25 - 4.0

def build_ffmpeg_args(response_format, input_format, sample_rate):
    # Convert the output to the desired format using ffmpeg
    if input_format == 'WAV':
        ffmpeg_args = ["ffmpeg", "-loglevel", "error", "-f", "WAV", "-i", "-"]
    else:
        ffmpeg_args = ["ffmpeg", "-loglevel", "error", "-f", input_format, "-ar", sample_rate, "-ac", "1", "-i", "-"]
    
    if response_format == "mp3":
        ffmpeg_args.extend(["-f", "mp3", "-c:a", "libmp3lame", "-ab", "64k"])
    elif response_format == "opus":
        ffmpeg_args.extend(["-f", "ogg", "-c:a", "libopus"])
    elif response_format == "aac":
        ffmpeg_args.extend(["-f", "adts", "-c:a", "aac", "-ab", "64k"])
    elif response_format == "flac":
        ffmpeg_args.extend(["-f", "flac", "-c:a", "flac"])
    elif response_format == "wav":
        ffmpeg_args.extend(["-f", "wav", "-c:a", "pcm_s16le"])
    elif response_format == "webm":
        ffmpeg_args.extend(["-f", "webm", "-c:a", "libopus"])
    elif response_format == "pcm": # even though pcm is technically 'raw', we still use ffmpeg to adjust the speed
        ffmpeg_args.extend(["-f", "s16le", "-c:a", "pcm_s16le"])

    return ffmpeg_args

@app.get("/v1/models")
async def list_models():
    """List all available TTS models (OpenAI-compatible format)"""
    # Return minimal OpenAI-compatible model list (no extra fields)
    # Only tts-1-qwen is enabled by default
    return {
        "object": "list",
        "data": [
            {
                "id": "tts-1-qwen",
                "object": "model",
                "created": 1700000000,
                "owned_by": "uncloseai"
            }
            # Other models disabled by default:
            # {
            #     "id": "tts-1",
            #     "object": "model",
            #     "created": 1700000000,
            #     "owned_by": "uncloseai"
            # },
            # {
            #     "id": "tts-1-hd",
            #     "object": "model",
            #     "created": 1700000000,
            #     "owned_by": "uncloseai"
            # },
            # {
            #     "id": "tts-1-silero",
            #     "object": "model",
            #     "created": 1700000000,
            #     "owned_by": "uncloseai"
            # },
            # {
            #     "id": "tts-1-kokoro",
            #     "object": "model",
            #     "created": 1700000000,
            #     "owned_by": "uncloseai"
            # }
        ]
    }

@app.get("/v1/voices")
async def list_voices():
    """List all available voices with model mapping and metadata (extended endpoint)"""
    global voices_cache

    # Return cached data if available
    if voices_cache is not None:
        return voices_cache

    # This should never happen since cache is populated at startup,
    # but provide fallback just in case
    logger.warning("/v1/voices called but cache not initialized - loading now")
    default_exists('config/voice_to_speaker.yaml')

    with open('config/voice_to_speaker.yaml', 'r', encoding='utf8') as file:
        voice_map = yaml.safe_load(file)

    models_data = []

    for model_id, voices in voice_map.items():
        if isinstance(voices, dict):
            voice_list = list(voices.keys())

            # Add model metadata with extended info
            model_info = {
                "id": model_id,
                "object": "model",
                "created": 1700000000,
                "owned_by": "uncloseai",
                "voices": voice_list,
                "voice_count": len(voice_list)
            }

            # Add engine-specific metadata
            if model_id == 'tts-1':
                model_info["engine"] = "piper"
                model_info["description"] = "Fast neural TTS with 100+ voices"
                model_info["sample_rate"] = 22050
            elif model_id == 'tts-1-hd':
                model_info["engine"] = "xtts"
                model_info["description"] = "High-quality voice cloning TTS"
                model_info["sample_rate"] = 24000
            elif model_id == 'tts-1-silero':
                model_info["engine"] = "silero"
                model_info["description"] = "Fast multilingual TTS (en, ru, de, es, fr)"
                model_info["sample_rate"] = 48000
            elif model_id == 'tts-1-kokoro':
                model_info["engine"] = "kokoro"
                model_info["description"] = "Lightweight decoder-only TTS (82M params)"
                model_info["sample_rate"] = 24000
            elif model_id == 'tts-1-qwen':
                model_info["engine"] = "qwen3-tts"
                model_info["description"] = "State-of-the-art TTS with voice cloning (1.7B params, 10 languages)"
                model_info["sample_rate"] = 24000

            models_data.append(model_info)

    voices_cache = {
        "object": "list",
        "data": models_data
    }

    return voices_cache

@app.post("/v1/audio/speech", response_class=StreamingResponse)
async def generate_speech(request: GenerateSpeechRequest):
    global xtts, args
    if len(request.input) < 1:
        raise BadRequestError("Empty Input", param='input')

    input_text = preprocess(request.input)

    if len(input_text) < 1:
        raise BadRequestError("Input text empty after preprocess.", param='input')

    model = request.model
    voice = request.voice
    response_format = request.response_format.lower()
    speed = request.speed

    # Auto-detect model from voice if model not provided
    if model is None:
        detected_model = detect_model_from_voice(voice)
        if detected_model:
            logger.info(f"Auto-detected model '{detected_model}' for voice '{voice}'")
            model = detected_model
        else:
            raise BadRequestError(f"Voice '{voice}' not found in any model. Please specify a model.", param='voice')

    # Set the Content-Type header based on the requested format
    if response_format == "mp3":
        media_type = "audio/mpeg"
    elif response_format == "opus":
        media_type = "audio/ogg;codec=opus" # codecs?
    elif response_format == "aac":
        media_type = "audio/aac"
    elif response_format == "flac":
        media_type = "audio/x-flac"
    elif response_format == "wav":
        media_type = "audio/wav"
    elif response_format == "webm":
        media_type = "audio/webm;codecs=opus"
    elif response_format == "pcm":
        if model == 'tts-1': # piper
            media_type = "audio/pcm;rate=22050"
        elif model == 'tts-1-hd': # xtts
            media_type = "audio/pcm;rate=24000"
        elif model == 'tts-1-silero': # silero
            media_type = "audio/pcm;rate=48000"
        elif model == 'tts-1-kokoro': # kokoro
            media_type = "audio/pcm;rate=24000"
        elif model == 'tts-1-qwen': # qwen3-tts
            media_type = "audio/pcm;rate=24000"
    else:
        raise BadRequestError(f"Invalid response_format: '{response_format}'", param='response_format')

    ffmpeg_args = None

    # Use piper for tts-1, and if xtts_device == none use for all models.
    if model == 'tts-1' or args.xtts_device == 'none':
        voice_map = map_voice_to_speaker(voice, 'tts-1')
        try:
            piper_model = voice_map['model']

        except KeyError as e:
            raise ServiceUnavailableError(f"Configuration error: tts-1 voice '{voice}' is missing 'model:' setting. KeyError: {e}")

        speaker = voice_map.get('speaker', None)

        # Use absolute path without data-dir when model path is absolute
        if os.path.isabs(piper_model):
            tts_args = ["piper", "--model", str(piper_model), "--output-raw"]
        else:
            tts_args = ["piper", "--model", str(piper_model), "--data-dir", "voices", "--download-dir", "voices", "--output-raw"]
        if speaker:
            tts_args.extend(["--speaker", str(speaker)])
        if speed != 1.0:
            tts_args.extend(["--length-scale", f"{1.0/speed}"])

        # Debug logging
        logger.info(f"Piper command: {' '.join(tts_args)}")
        logger.info(f"Model file exists: {os.path.exists(piper_model)}")

        tts_proc = subprocess.Popen(tts_args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        tts_proc.stdin.write(bytearray(input_text.encode('utf-8')))
        tts_proc.stdin.close()

        # Log any stderr output from Piper
        if tts_proc.stderr:
            def log_stderr():
                stderr_output = tts_proc.stderr.read().decode('utf-8', errors='replace')
                if stderr_output.strip():
                    logger.error(f"Piper stderr: {stderr_output}")
            threading.Thread(target=log_stderr, daemon=True).start()

        try:
            with open(f"{piper_model}.json", 'r') as pvc_f:
                conf = json.load(pvc_f)
                sample_rate = str(conf['audio']['sample_rate'])

        except:
            sample_rate = '22050'
  
        ffmpeg_args = build_ffmpeg_args(response_format, input_format="s16le", sample_rate=sample_rate)

        # Pipe the output from piper/xtts to the input of ffmpeg
        ffmpeg_args.extend(["-"])
        ffmpeg_proc = subprocess.Popen(ffmpeg_args, stdin=tts_proc.stdout, stdout=subprocess.PIPE)

        return StreamingResponse(content=ffmpeg_proc.stdout, media_type=media_type)
    # Use xtts for tts-1-hd
    elif model == 'tts-1-hd':
        voice_map = map_voice_to_speaker(voice, 'tts-1-hd')
        try:
            tts_model = voice_map.pop('model')
            speaker = voice_map.pop('speaker')

        except KeyError as e:
            raise ServiceUnavailableError(f"Configuration error: tts-1-hd voice '{voice}' is missing setting. KeyError: {e}")

        if xtts and xtts.model_name != tts_model:
            unload_model()

        tts_model_path = voice_map.pop('model_path', None) # XXX changing this on the fly is ignored if you keep the same name

        if xtts is None:
            xtts = xtts_wrapper(tts_model, device=args.xtts_device, model_path=tts_model_path, unload_timer=args.unload_timer)

        ffmpeg_args = build_ffmpeg_args(response_format, input_format="f32le", sample_rate="24000")

        # tts speed doesn't seem to work well
        speed = voice_map.pop('speed', speed)
        if speed < 0.5:
            speed = speed / 0.5
            ffmpeg_args.extend(["-af", "atempo=0.5"]) 
        if speed > 1.0:
            ffmpeg_args.extend(["-af", f"atempo={speed}"]) 
            speed = 1.0

        # Pipe the output from piper/xtts to the input of ffmpeg
        ffmpeg_args.extend(["-"])

        language = voice_map.pop('language', 'auto')
        if language == 'auto':
            try:
                language = detect(input_text)
                if language not in [
                    'en', 'es', 'fr', 'de', 'it', 'pt', 'pl', 'tr',
                    'ru', 'nl', 'cs', 'ar', 'zh-cn', 'hu', 'ko', 'ja', 'hi'
                ]:
                    logger.debug(f"Detected language {language} not supported, defaulting to en")
                    language = 'en'
                else:
                    logger.debug(f"Detected language: {language}")
            except:
                language = 'en'
                logger.debug(f"Failed to detect language, defaulting to en")

        comment = voice_map.pop('comment', None) # ignored.

        hf_generate_kwargs = dict(
            speed=speed,
            **voice_map,
        )

        hf_generate_kwargs['enable_text_splitting'] = hf_generate_kwargs.get('enable_text_splitting', True) # change the default to true

        if hf_generate_kwargs['enable_text_splitting']:
            if language == 'zh-cn':
                split_lang = 'zh'
            else:
                split_lang = language
            all_text = split_sentence(input_text, split_lang, xtts.xtts.tokenizer.char_limits[split_lang])
        else:
            all_text = [input_text]

        ffmpeg_proc = subprocess.Popen(ffmpeg_args, stdin=subprocess.PIPE, stdout=subprocess.PIPE)

        in_q = queue.Queue() # speech pcm 
        ex_q = queue.Queue() # exceptions

        def get_speaker_samples(samples: str) -> list[str]:
            if os.path.isfile(samples):
                audio_path = [samples]
            elif os.path.isdir(samples):
                audio_path = [os.path.join(samples, sample) for sample in os.listdir(samples) if os.path.isfile(os.path.join(samples, sample))]

                if len(audio_path) < 1:
                    logger.error(f"No files found: {samples}")
                    raise ServiceUnavailableError(f"Invalid path: {samples}")
            else:
                logger.error(f"Invalid path: {samples}")
                raise ServiceUnavailableError(f"Invalid path: {samples}")
            
            return audio_path

        def exception_check(exq: queue.Queue):
            try:
                e = exq.get_nowait()
            except queue.Empty:
                return
            
            raise e

        def generator():
            # text -> in_q

            audio_path = get_speaker_samples(speaker)
            logger.debug(f"{voice} wav samples: {audio_path}")

            try:
                for text in all_text:
                    for chunk in xtts.tts(text=text, language=language, audio_path=audio_path, **hf_generate_kwargs):
                        exception_check(ex_q)
                        in_q.put(chunk)

            except BrokenPipeError as e: # client disconnect lands here
                logger.info("Client disconnected - 'Broken pipe'")

            except Exception as e:
                logger.error(f"Exception: {repr(e)}")
                raise e
        
            finally:
                in_q.put(None) # sentinel

        def out_writer(): 
            # in_q -> ffmpeg
            try:
                while True:
                    chunk = in_q.get()
                    if chunk is None: # sentinel
                        break
                    ffmpeg_proc.stdin.write(chunk) # BrokenPipeError from here on client disconnect

            except Exception as e: # BrokenPipeError
                ex_q.put(e)  # we need to get this exception into the generation loop
                ffmpeg_proc.kill()
                return
            
            finally:
                ffmpeg_proc.stdin.close()

        generator_worker = threading.Thread(target=generator, daemon=True)
        generator_worker.start()

        out_writer_worker = threading.Thread(target=out_writer, daemon=True)
        out_writer_worker.start()

        def cleanup():
            ffmpeg_proc.kill()
            # Only delete workers if they were created
            try:
                del generator_worker
            except NameError:
                pass
            try:
                del out_writer_worker
            except NameError:
                pass

        return StreamingResponse(content=ffmpeg_proc.stdout, media_type=media_type, background=cleanup)
    # Use Silero for tts-1-silero
    elif model == 'tts-1-silero':
        global silero_model, silero_speakers

        voice_map = map_voice_to_speaker(voice, 'tts-1-silero')
        language = voice_map.get('language', 'en')
        speaker_id = voice_map.get('speaker', 'en_0')
        silero_speaker_key = voice_map.get('silero_speaker', 'v4_en')

        # Create a unique key for this language+speaker combination
        model_key = f"{language}_{silero_speaker_key}"

        # Load Silero model if not already loaded or if language/speaker changed
        if silero_model is None or silero_speakers.get('current') != model_key:
            # Use semaphore to prevent multiple simultaneous model loads
            async with silero_load_semaphore:
                # Double-check after acquiring lock (another request may have loaded it)
                if silero_model is None or silero_speakers.get('current') != model_key:
                    logger.info(f"Loading/switching Silero model to {language}/{silero_speaker_key}")
                    # Run blocking model initialization in thread pool to avoid blocking event loop
                    silero_model = await asyncio.to_thread(silero_wrapper, language=language, speaker=silero_speaker_key, device='cpu')
                    silero_speakers['current'] = model_key

        # Generate audio (also blocking, so run in thread pool)
        audio_data = await asyncio.to_thread(silero_model.tts, input_text, speaker_id=speaker_id)

        # Silero outputs float32 PCM at 48000 Hz
        ffmpeg_args = build_ffmpeg_args(response_format, input_format="f32le", sample_rate="48000")

        # Apply speed adjustment if needed
        if speed != 1.0:
            ffmpeg_args.extend(["-af", f"atempo={speed}"])

        ffmpeg_args.extend(["-"])
        ffmpeg_proc = subprocess.Popen(ffmpeg_args, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        ffmpeg_proc.stdin.write(audio_data)
        ffmpeg_proc.stdin.close()

        return StreamingResponse(content=ffmpeg_proc.stdout, media_type=media_type)
    # Use Kokoro for tts-1-kokoro
    elif model == 'tts-1-kokoro':
        global kokoro_pipeline, kokoro_lang

        voice_map = map_voice_to_speaker(voice, 'tts-1-kokoro')
        lang_code = voice_map.get('lang_code', 'a')
        kokoro_voice = voice_map.get('kokoro_voice', 'af_heart')

        # Load Kokoro pipeline if not already loaded or if language changed
        if kokoro_pipeline is None or kokoro_lang != lang_code:
            # Use semaphore to prevent multiple simultaneous model loads
            async with kokoro_load_semaphore:
                # Double-check after acquiring lock (another request may have loaded it)
                if kokoro_pipeline is None or kokoro_lang != lang_code:
                    # Run blocking model initialization in thread pool to avoid blocking event loop
                    # Use GPU if available, fallback to CPU only if explicitly disabled
                    device = args.xtts_device if args.xtts_device != 'none' else 'cpu'
                    logger.info(f"Loading/switching Kokoro pipeline to language '{lang_code}' on device '{device}'")
                    kokoro_pipeline = await asyncio.to_thread(kokoro_wrapper, lang_code=lang_code, device=device)
                    kokoro_lang = lang_code

        # Split long text into sentences for streaming
        sentences = simple_sentence_split(input_text, max_length=500)
        logger.info(f"Split text into {len(sentences)} sentences for Kokoro streaming")

        # Kokoro outputs float32 PCM at 24000 Hz
        ffmpeg_args = build_ffmpeg_args(response_format, input_format="f32le", sample_rate="24000")
        ffmpeg_args.extend(["-"])
        ffmpeg_proc = subprocess.Popen(ffmpeg_args, stdin=subprocess.PIPE, stdout=subprocess.PIPE)

        # Use threading approach like XTTS to ensure proper sequential processing
        in_q = queue.Queue()  # audio chunks

        def generator():
            """Process sentences sequentially and feed to queue"""
            try:
                for idx, sentence in enumerate(sentences):
                    logger.debug(f"Processing sentence {idx+1}/{len(sentences)}: {len(sentence)} chars")
                    audio_bytes = kokoro_pipeline.tts(sentence, voice=kokoro_voice, speed=speed)
                    in_q.put(audio_bytes)
                    logger.debug(f"Kokoro: queued sentence {idx+1}/{len(sentences)}")
            except Exception as e:
                logger.error(f"Kokoro streaming error: {e}")
            finally:
                in_q.put(None)  # sentinel
                logger.info(f"Kokoro streaming complete: {len(sentences)} sentences processed")

        def out_writer():
            """Write audio from queue to ffmpeg stdin"""
            try:
                while True:
                    chunk = in_q.get()
                    if chunk is None:  # sentinel
                        break
                    ffmpeg_proc.stdin.write(chunk)
            except Exception as e:
                logger.error(f"Kokoro ffmpeg write error: {e}")
                ffmpeg_proc.kill()
            finally:
                ffmpeg_proc.stdin.close()

        generator_worker = threading.Thread(target=generator, daemon=True)
        generator_worker.start()

        out_writer_worker = threading.Thread(target=out_writer, daemon=True)
        out_writer_worker.start()

        return StreamingResponse(content=ffmpeg_proc.stdout, media_type=media_type)
    # Use Qwen3-TTS for tts-1-qwen
    elif model == 'tts-1-qwen':
        global qwen_model, qwen_voice_prompts

        if not QWEN_TTS_AVAILABLE:
            raise ServiceUnavailableError("Qwen3-TTS is not available. Install with: pip install qwen-tts")

        voice_map = map_voice_to_speaker(voice, 'tts-1-qwen')
        ref_audio = voice_map.get('ref_audio')
        ref_text = voice_map.get('ref_text')
        language = voice_map.get('language', 'English')

        # Load Qwen model if not already loaded
        if qwen_model is None:
            async with qwen_load_semaphore:
                if qwen_model is None:
                    device = args.xtts_device if args.xtts_device != 'none' else 'cpu'
                    logger.info(f"Loading Qwen3-TTS model on device '{device}'")
                    qwen_model = await asyncio.to_thread(
                        qwen3_wrapper,
                        model_name='Qwen/Qwen3-TTS-12Hz-1.7B-Base',
                        device=device
                    )

        # Create or retrieve cached voice prompt
        voice_prompt = None
        if ref_audio and ref_text:
            cache_key = f"{voice}_{ref_audio}"
            if cache_key not in qwen_voice_prompts:
                logger.info(f"Creating voice prompt for '{voice}'")
                qwen_voice_prompts[cache_key] = await asyncio.to_thread(
                    qwen_model.create_voice_prompt,
                    ref_audio=ref_audio,
                    ref_text=ref_text
                )
            voice_prompt = qwen_voice_prompts[cache_key]
        else:
            raise BadRequestError(f"Voice '{voice}' requires ref_audio and ref_text configuration", param='voice')

        # Split text into sentences for streaming (first audio arrives faster)
        sentences = simple_sentence_split(input_text, max_length=500)
        logger.info(f"Split text into {len(sentences)} sentences for Qwen streaming")

        # Qwen outputs float32 PCM at ~24kHz
        ffmpeg_args = build_ffmpeg_args(response_format, input_format="f32le", sample_rate="24000")

        # Apply speed adjustment if needed
        if speed != 1.0:
            ffmpeg_args.extend(["-af", f"atempo={speed}"])

        ffmpeg_args.extend(["-"])
        ffmpeg_proc = subprocess.Popen(ffmpeg_args, stdin=subprocess.PIPE, stdout=subprocess.PIPE)

        # Use queue + threading for sentence-by-sentence streaming
        in_q = queue.Queue()

        def generator():
            """Process sentences sequentially and feed to queue"""
            try:
                for idx, sentence in enumerate(sentences):
                    logger.debug(f"Qwen processing sentence {idx+1}/{len(sentences)}: {len(sentence)} chars")
                    audio_bytes = qwen_model.tts(
                        text=sentence,
                        language=language,
                        voice_prompt=voice_prompt
                    )
                    in_q.put(audio_bytes)
                    logger.debug(f"Qwen: queued sentence {idx+1}/{len(sentences)}")
            except Exception as e:
                logger.error(f"Qwen streaming error: {e}")
            finally:
                in_q.put(None)  # sentinel
                logger.info(f"Qwen streaming complete: {len(sentences)} sentences processed")

        def out_writer():
            """Write audio from queue to ffmpeg stdin"""
            try:
                while True:
                    chunk = in_q.get()
                    if chunk is None:  # sentinel
                        break
                    ffmpeg_proc.stdin.write(chunk)
            except Exception as e:
                logger.error(f"Qwen ffmpeg write error: {e}")
                ffmpeg_proc.kill()
            finally:
                ffmpeg_proc.stdin.close()

        generator_worker = threading.Thread(target=generator, daemon=True)
        generator_worker.start()

        out_writer_worker = threading.Thread(target=out_writer, daemon=True)
        out_writer_worker.start()

        return StreamingResponse(content=ffmpeg_proc.stdout, media_type=media_type)
    else:
        raise BadRequestError("No such model, must be tts-1-qwen (default), tts-1, tts-1-hd, tts-1-silero, or tts-1-kokoro.", param='model')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='uncloseai-speech API Server',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--xtts_device', action='store', default=auto_torch_device(), help="Set the device for the xtts model. The special value of 'none' will use piper for all models.")
    parser.add_argument('--preload', action='store', default=None, help="Preload a model (Ex. 'xtts' or 'xtts_v2.0.2'). By default it's loaded on first use.")
    parser.add_argument('--unload-timer', action='store', default=None, type=int, help="Idle unload timer for the XTTS model in seconds, Ex. 900 for 15 minutes")
    parser.add_argument('--use-deepspeed', action='store_true', default=False, help="Use deepspeed with xtts (this option is unsupported)")
    parser.add_argument('--no-cache-speaker', action='store_true', default=False, help="Don't use the speaker wav embeddings cache")
    parser.add_argument('-W', '--workers', action='store', default=4, type=int, help="Number of uvicorn worker processes for concurrent request handling")
    parser.add_argument('-P', '--port', action='store', default=8000, type=int, help="Server tcp port")
    parser.add_argument('-H', '--host', action='store', default='0.0.0.0', help="Host to listen on, Ex. 0.0.0.0")
    parser.add_argument('-L', '--log-level', default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], help="Set the log level")

    args = parser.parse_args()

    logger.remove()
    logger.add(sink=sys.stderr, level=args.log_level)

    if args.preload and not XTTS_AVAILABLE:
        logger.error("Cannot preload XTTS model - XTTS dependencies not available")
    elif args.preload:
        xtts = xtts_wrapper(args.preload, device=args.xtts_device, unload_timer=args.unload_timer)

    # Register only Qwen by default (other models disabled)
    app.register_model('tts-1-qwen')
    # To enable other models, uncomment below:
    # app.register_model('tts-1')
    # app.register_model('tts-1-hd')
    # app.register_model('tts-1-silero')
    # app.register_model('tts-1-kokoro')

    # Use multiple workers for true concurrency (each worker = separate process with own GIL)
    # This prevents thread pool exhaustion and allows concurrent model loading
    # Must use import string format for workers to function
    uvicorn.run("speech:app", host=args.host, port=args.port, workers=args.workers, timeout_keep_alive=300)