import os
os.environ["COQUI_TOS_AGREED"] = "1"

import spaces
import gradio as gr
import torch

# --- Compatibility shim ---
# Newer `transformers` (5.x, needed here for gradio's huggingface-hub
# requirement) removed `isin_mps_friendly` from transformers.pytorch_utils.
# coqui-tts still imports it. This restores the function before coqui-tts
# is imported, so both libraries work together in this environment.
import transformers.pytorch_utils as _ptu
if not hasattr(_ptu, "isin_mps_friendly"):
    def _isin_mps_friendly(elements, test_elements):
        return torch.isin(elements, test_elements)
    _ptu.isin_mps_friendly = _isin_mps_friendly

from TTS.api import TTS

print("Loading XTTS-v2 model...")
device = "cuda" if torch.cuda.is_available() else "cpu"
tts_model = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
print("Model ready.")

LANGUAGES = {
    "English": "en",
    "Hindi": "hi",
    "Spanish": "es",
    "French": "fr",
    "German": "de",
    "Arabic": "ar",
    "Chinese": "zh-cn",
    "Japanese": "ja",
    "Korean": "ko",
}


@spaces.GPU(duration=60)
def clone_voice(reference_audios, text, language_label):
    if not reference_audios:
        raise gr.Error("Please upload at least one voice sample (longer, clean samples work best).")
    if not text or not text.strip():
        raise gr.Error("Please enter the text you want spoken in the cloned voice.")

    try:
        lang_code = LANGUAGES[language_label]
        output_path = "/tmp/cloned_output.wav"

        # Accept either a single file or multiple — more reference clips
        # generally improves how closely the output matches the voice.
        speaker_wavs = reference_audios if isinstance(reference_audios, list) else [reference_audios]

        tts_model.tts_to_file(
            text=text,
            speaker_wav=speaker_wavs,
            language=lang_code,
            file_path=output_path,
        )

        return output_path

    except Exception as e:
        import traceback
        traceback.print_exc()
        raise gr.Error(f"Voice cloning failed: {e}")


css = """
#header {
    text-align: center;
    padding: 24px 0 8px;
}
#header h1 {
    font-size: 32px;
    font-weight: 700;
    background: linear-gradient(135deg, #6366f1, #06b6d4);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 4px;
}
#header p {
    color: #888;
    font-size: 14px;
}
#license-note {
    text-align: center;
    font-size: 12px;
    color: #999;
    padding: 0 20px 16px;
}
#run-btn {
    background: linear-gradient(135deg, #6366f1, #06b6d4) !important;
    color: white !important;
    font-weight: 600 !important;
    border: none !important;
}
"""

with gr.Blocks(title="Peace Network Voice Clone", css=css) as demo:
    gr.HTML(
        """
        <div id="header">
            <h1>🗣️ Peace Network Voice Clone</h1>
            <p>Apni awaaz ka sample daaliye — kuch bhi text usi awaaz mein bol dega.</p>
        </div>
        <div id="license-note">
            Non-commercial use only (Coqui Public Model License). Personal projects ke liye theek hai, client/commercial kaam ke liye alag license chahiye.
        </div>
        """
    )

    reference = gr.File(
        label="Your voice samples (upload 2-3 clean clips, 20-30 sec each, works much better than one short clip)",
        file_count="multiple",
        file_types=["audio"],
        type="filepath",
    )

    text_input = gr.Textbox(
        label="Text to speak in your cloned voice",
        placeholder="Type what you want your cloned voice to say...",
        lines=5,
    )

    language = gr.Dropdown(
        choices=list(LANGUAGES.keys()), value="Hindi", label="Language"
    )

    btn = gr.Button("🗣️ Clone & Generate", variant="primary", elem_id="run-btn")
    output_audio = gr.Audio(label="Cloned voice output", type="filepath")

    btn.click(clone_voice, inputs=[reference, text_input, language], outputs=output_audio)

demo.queue().launch()