import numpy as np
import spaces
import torch
import gradio as gr
import tempfile
import scipy.io.wavfile as wavfile
from transformers import VitsModel, AutoTokenizer
# ========== FIX: Monkey patch to handle bool schemas ==========
import gradio_client.utils as client_utils
original_get_type = client_utils.get_type
def patched_get_type(schema):
    if isinstance(schema, bool):
        return "bool"
    return original_get_type(schema)
client_utils.get_type = patched_get_type
# Also patch the _json_schema_to_python_type function to handle bool
original_json_schema = client_utils._json_schema_to_python_type
def patched_json_schema(schema, defs=None):
    if isinstance(schema, bool):
        return "bool"
    return original_json_schema(schema, defs)
client_utils._json_schema_to_python_type = patched_json_schema
# =============================================================

# NOTE: Two Dinka variants exist on MMS:
#   facebook/mms-tts-dik  -> Dinka, Southwestern (Rek) - standard written Dinka
#   facebook/mms-tts-dip  -> Dinka, Northeastern
# Swap MODEL_ID below if you need the other dialect.
MODEL_ID = "facebook/mms-tts-dik"

EXAMPLES = [
    "Ɣɛn anɔŋ mäth cɔl Yar.",
    "Alaak aciɛ̈ɛ̈r në Nairobi, Kenya.",
    "Ɣɛn lɔ Juba, miäkduur.",
    "Duɔ̈kkë wenhïïm ye yiëk këriëëc ëbën, kuat kë yakë wiɛ̈ckë, röökkë ku röökkë ku lëkkë Nhialic lɛc. Dɔ̈ɔ̈r Nhialic, dɔ̈ɔ̈r cie raan ŋic raan ëbën, abï we tiit ku nhïïmkun tit rin cï wek mat kek Jethu raan cï lɔc ku dɔc.",
    "Nairobi akor në pɛɛi de dhetem ku pɛɛi de dhorou.",
]

print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = VitsModel.from_pretrained(MODEL_ID)
model.eval()
print("Model loaded successfully.")

@spaces.GPU
def synthesize(text, seed):
    try:
        if text is None or not text.strip():
            raise gr.Error("Please enter Dinka text first.")

        device = "cuda" if torch.cuda.is_available() else "cpu"
        print("Running on:", device)
        model.to(device)
        torch.manual_seed(int(seed))
        inputs = tokenizer(text, return_tensors="pt").to(device)

        with torch.no_grad():
            waveform = model(**inputs).waveform

        waveform = waveform.squeeze().cpu().numpy().astype(np.float32)
        sample_rate = model.config.sampling_rate

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as file:
            wav_path = file.name

        wavfile.write(wav_path, sample_rate, waveform)
        return (sample_rate, waveform), wav_path

    except Exception as e:
        print("ERROR:", e)
        raise gr.Error(f"Speech generation failed: {str(e)}")

CUSTOM_CSS = """
.header {
    text-align:center;
    padding:20px;
    border-radius:12px;
    background:linear-gradient(90deg, #1565C0, #2E7D32);
    color:white;
    margin-bottom:20px;
}
.header h1, .header p { color:white; }
"""

with gr.Blocks(css=CUSTOM_CSS, title="Dinka Text-to-Speech") as demo:
    gr.HTML("""
        <div class="header">
        <h1>🇸🇸 Dinka Text-to-Speech</h1>
        <p>Powered by Meta MMS (facebook/mms-tts-dik)</p>
        </div>
    """)

    text = gr.Textbox(label="Dinka Text", placeholder="Type Dinka text here...", lines=5)

    with gr.Accordion("Advanced Options", open=False):
        seed = gr.Number(label="Seed", value=42, precision=0)

    generate = gr.Button("Generate Speech", variant="primary")
    audio = gr.Audio(label="Generated Speech", type="numpy")
    download = gr.File(label="Download WAV")

    generate.click(fn=synthesize, inputs=[text, seed], outputs=[audio, download])

    gr.Markdown("### Dinka Examples")
    gr.Examples(examples=[[x] for x in EXAMPLES], inputs=[text])
    gr.Markdown("Model: https://huggingface.co/facebook/mms-tts-dik")

if __name__ == "__main__":
    # Try different launch options
    try:
        demo.launch(show_api=False)
    except TypeError:
        try:
            demo.launch(api_name=False)
        except TypeError:
            try:
                demo.launch(api=False)
            except TypeError:
                demo.launch()