import os
import io
import spaces
import torch
import torchaudio
import numpy as np
import gradio as gr
from pocket_tts import TTSModel

# 1. Load the Pocket-TTS model at startup (CPU-only, ~100M params)
print("Loading Kyutai Pocket-TTS model...")
tts_model = TTSModel.load_model()
print("Pocket-TTS model loaded successfully!")

# Official Kyutai Pocket-TTS English voice list
# See: https://huggingface.co/kyutai/tts-voices
VOICES = [
    "alba",
    "anna",
    "azelma",
    "bill_boerst",
    "caro_davy",
    "charles",
    "cosette",
    "eponine",
    "eve",
    "fantine",
    "george",
    "jane",
    "jean",
    "javert",
    "marius",
    "mary",
    "michael",
    "paul",
    "peter_yearsley",
    "stuart_bell",
    "vera",
]

# Pre-cache voice states at startup for faster inference
print("Pre-caching voice states...")
voice_states = {}
for voice_name in VOICES:
    try:
        voice_states[voice_name] = tts_model.get_state_for_audio_prompt(voice_name)
        print(f"  Cached voice: {voice_name}")
    except Exception as e:
        print(f"  Warning: Could not cache voice '{voice_name}': {e}")
print("Voice states cached!")


def change_speed_pitch_preserved(audio_np: np.ndarray, sample_rate: int, speed: float) -> np.ndarray:
    """
    Adjusts speech speed while preserving pitch/formants without echo.
    Tries backends in quality order: pyrubberband (RubberBand) > sox tempo -s (WSOLA) > audiotsm WSOLA.
    Deliberately AVOIDS librosa.effects.time_stretch (phase vocoder) -> that is the echo source.
    speed >1 = faster/shorter, speed <1 = slower/longer
    """
    if abs(speed - 1.0) < 0.02:
        return audio_np
    # Clamp to avoid extreme WSOLA artifacts
    speed = float(np.clip(speed, 0.5, 2.0))
    x = audio_np.astype(np.float32)
    # 1) Best quality: RubberBand (if binary available)
    try:
        import pyrubberband as rb
        # rb.time_stretch expects rate = speed ( >1 faster)
        y = rb.time_stretch(x, sample_rate, rate=speed)
        return y.astype(np.float32)
    except Exception as e:
        print(f"pyrubberband unavailable: {e}")
    # 2) SoX WSOLA via torchaudio (fast, good, but deprecated API)
    try:
        tensor = torch.from_numpy(x).float()
        if tensor.dim() == 1:
            tensor = tensor.unsqueeze(0)
        # Newer torchaudio >=2.4 moved sox_effects; try both
        try:
            import torchaudio.sox_effects as sox_effects
            effects = [["tempo", "-s", str(speed)]]
            stretched, _ = sox_effects.apply_effects_tensor(tensor, sample_rate, effects)
            return stretched.squeeze(0).numpy().astype(np.float32)
        except Exception:
            # Fallback: torchaudio.functional - not tempo, skip
            raise
    except Exception as e:
        print(f"SoX tempo fallback: {e}")
    # 3) Pure-python WSOLA via audiotsm (no system deps, no phase-vocoder echo)
    try:
        from audiotsm import wsola
        from audiotsm.io.array import ArrayReader, ArrayWriter
        # audiotsm expects (channels, samples) float32
        channels = 1
        reader = ArrayReader(x[np.newaxis, :])
        writer = ArrayWriter(channels)
        # speed >1 = faster, so wsola speed param is same
        tsm = wsola(channels=channels, speed=speed)
        tsm.run(reader, writer)
        y = writer.data[0]
        return y.astype(np.float32)
    except Exception as e:
        print(f"audiotsm WSOLA fallback failed: {e}")
    # 4) Last resort: no DSP, return original and let client do playbackRate
    print("All time-stretch backends failed; returning original audio (use client playbackRate)")
    return audio_np

@spaces.GPU
def synthesize(text: str, voice: str, speed: float = 1.0):
    """
    Generates audio from Pocket-TTS and applies clean server-side speed adjustment.
    """
    if not text or not text.strip():
        raise gr.Error("Text prompt cannot be empty.")

    clean_text = text.strip()
    clean_voice = voice.lower().strip()
    if clean_voice not in VOICES:
        clean_voice = "alba"

    speed_factor = max(0.5, min(2.0, float(speed) if speed else 1.0))

    # 1. Get the cached voice state, or load it on demand
    if clean_voice in voice_states:
        voice_state = voice_states[clean_voice]
    else:
        voice_state = tts_model.get_state_for_audio_prompt(clean_voice)

    # 2. Generate audio using the official API
    audio_tensor = tts_model.generate_audio(voice_state, clean_text)

    # Convert to numpy
    audio_np = audio_tensor.numpy().astype(np.float32)
    sample_rate = tts_model.sample_rate

    # 3. Adjust speed with pitch preservation
    if abs(speed_factor - 1.0) >= 0.02:
        audio_np = change_speed_pitch_preserved(audio_np, sample_rate, speed_factor)

    # Normalize audio to prevent clipping
    max_val = np.max(np.abs(audio_np))
    if max_val > 0:
        audio_np = (audio_np / max_val) * 0.95

    # Return in Gradio (sample_rate, numpy_int16_array) format
    int16_audio = (audio_np * 32767).astype(np.int16)
    return (sample_rate, int16_audio)


# --- Gradio UI & API Interface ---
with gr.Blocks(title="Kyutai Pocket-TTS Server") as demo:
    gr.Markdown("# 🎙️ Kyutai Pocket-TTS Server with Speed Control")

    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="Text to Synthesize",
                placeholder="Enter text to speak...",
                lines=4,
                value="The quick brown fox jumps over the lazy dog."
            )
            voice_input = gr.Dropdown(
                label="Voice",
                choices=VOICES,
                value="alba"
            )
            speed_slider = gr.Slider(
                label="Speed Multiplier",
                minimum=0.5,
                maximum=2.0,
                step=0.05,
                value=1.0
            )
            generate_btn = gr.Button("Generate Speech", variant="primary")

        with gr.Column():
            audio_output = gr.Audio(label="Synthesized Audio", type="numpy")

    generate_btn.click(
        fn=synthesize,
        inputs=[text_input, voice_input, speed_slider],
        outputs=audio_output,
        api_name="predict"
    )

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)