"""
Text-to-speech for TalkToDoc.
Converts translated text responses into audible speech in the patient's
selected language.

English, Yoruba, Igbo, and Hausa use YarnGPT, a local model trained
specifically on Nigerian-accented speech.

Nigerian Pidgin uses MMS-TTS (facebook/mms-tts-pcm) instead, since YarnGPT
has no Pidgin support. MMS-TTS is the only free option found with a
dedicated Pidgin checkpoint.

Note on the WavTokenizer checkpoint: YarnGPT downloads this itself on
first import, using a bare requests.get() with no error checking, which
can silently save a corrupted file if the download hiccups. The
Dockerfile pre-downloads the same two files at build time using
huggingface_hub's properly tested download function instead, so by the
time this file is imported, YarnGPT finds them already in place and
skips its own fragile download step. Locally (not in Docker), the first
import still triggers YarnGPT's own download as normal, this only
matters for the deployed container.

Performance note: YarnGPT's own generate_speech() function reloads its
full model from disk on every call, which is too slow for a live app.
This file loads the model once and reuses it, using the same generation
steps YarnGPT's own function uses internally, just without the reload.
Call preload_models() once when the app starts, so the first real request
isn't slow either.
"""

import torch
import torchaudio
import scipy.io.wavfile
from transformers import VitsModel, AutoTokenizer
from yarngpt.core import load_model_and_tokenizer, SPEAKER_MAPPING, AVAILABLE_SPEAKERS

# Default speaker used for each YarnGPT-supported language.
YARNGPT_SPEAKERS = {
    "english": "idera",
    "yoruba": "abayomi",
    "igbo": "chioma",
    "hausa": "amina",
}

_yarngpt_model = None
_yarngpt_tokenizer = None

_pidgin_model = None
_pidgin_tokenizer = None


def _get_yarngpt_model():
    global _yarngpt_model, _yarngpt_tokenizer
    if _yarngpt_model is None:
        _yarngpt_model, _yarngpt_tokenizer = load_model_and_tokenizer()
    return _yarngpt_model, _yarngpt_tokenizer


def _get_pidgin_model():
    global _pidgin_model, _pidgin_tokenizer
    if _pidgin_model is None:
        _pidgin_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-pcm")
        _pidgin_model = VitsModel.from_pretrained("facebook/mms-tts-pcm")
    return _pidgin_model, _pidgin_tokenizer


def preload_models():
    """Loads both TTS backends into memory ahead of time. Call this once
    when the Flask app starts, so the first real request isn't slow."""
    _get_yarngpt_model()
    _get_pidgin_model()


def _generate_yarngpt_speech(text, speaker, language, temperature=0.1, repetition_penalty=1.1, max_length=4000):
    """
    Same steps as yarngpt's own generate_speech(), but reuses the model
    already loaded by _get_yarngpt_model() instead of reloading it.
    """
    model_speaker = SPEAKER_MAPPING.get(speaker, speaker)
    if model_speaker not in AVAILABLE_SPEAKERS:
        raise ValueError(f"Unknown speaker: {speaker}")

    model, audio_tokenizer = _get_yarngpt_model()

    prompt = audio_tokenizer.create_prompt(text, language, model_speaker)
    input_ids = audio_tokenizer.tokenize_prompt(prompt)
    attention_mask = torch.ones_like(input_ids)

    output = model.generate(
        input_ids=input_ids,
        attention_mask=attention_mask,
        do_sample=True,
        temperature=temperature,
        repetition_penalty=repetition_penalty,
        max_length=max_length,
        pad_token_id=model.config.eos_token_id,
        eos_token_id=model.config.eos_token_id,
    )

    codes = audio_tokenizer.get_codes(output)
    audio = audio_tokenizer.get_audio(codes)
    return audio


def synthesize_speech(text, language, output_path):
    """
    text: the text to speak
    language: one of "english", "yoruba", "hausa", "igbo", "pidgin"
    output_path: where to save the resulting .wav file
    Returns output_path.
    """
    language = language.lower()

    if language in YARNGPT_SPEAKERS:
        speaker = YARNGPT_SPEAKERS[language]
        audio = _generate_yarngpt_speech(text, speaker=speaker, language=language)
        torchaudio.save(output_path, audio, sample_rate=24000)
        return output_path

    if language == "pidgin":
        model, tokenizer = _get_pidgin_model()
        inputs = tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            output = model(**inputs).waveform
        scipy.io.wavfile.write(
            output_path,
            rate=model.config.sampling_rate,
            data=output.numpy().squeeze(),
        )
        return output_path

    raise ValueError(f"Unsupported language: {language}")


if __name__ == "__main__":
    import sys

    if len(sys.argv) < 3:
        print('Usage: python tts.py "text" language [output_file]')
    else:
        input_text = sys.argv[1]
        selected_language = sys.argv[2]
        output_file = sys.argv[3] if len(sys.argv) > 3 else "output.wav"
        synthesize_speech(input_text, selected_language, output_file)
        print("Saved to", output_file)