import gradio as gr
import torch
import numpy as np
import scipy.io.wavfile
import tempfile
import os
from transformers import VitsModel, AutoTokenizer
from gtts import gTTS

# ============================================================
# 🎙️ Hindi Text-to-Speech Studio
# Two Models: Meta MMS-TTS + Google TTS
# ============================================================

# --- Global Model Loading (loaded once at startup) ---
print("🔮 Loading Meta MMS-TTS Hindi model...")
mms_model = VitsModel.from_pretrained("facebook/mms-tts-hin")
mms_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-hin")
print("✅ Meta MMS-TTS Hindi loaded successfully!")


def generate_mms_tts(text, speed=1.0):
    """Generate speech using Facebook MMS-TTS Hindi (VITS model)"""
    if not text or not text.strip():
        raise gr.Error("⚠️ कृपया कुछ हिंदी टेक्स्ट दर्ज करें! (Please enter some Hindi text!)")

    try:
        inputs = mms_tokenizer(text.strip(), return_tensors="pt")

        with torch.no_grad():
            # Set speaking rate
            mms_model.speaking_rate = speed
            output = mms_model(**inputs)

        waveform = output.waveform.float().numpy().squeeze()
        sample_rate = mms_model.config.sampling_rate

        # Normalize audio
        if np.max(np.abs(waveform)) > 0:
            waveform = waveform / np.max(np.abs(waveform)) * 0.95

        # Save to temp file
        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        scipy.io.wavfile.write(tmp.name, rate=sample_rate, data=waveform)
        return tmp.name

    except Exception as e:
        raise gr.Error(f"❌ MMS-TTS Error: {str(e)}")


def generate_google_tts(text, slow=False):
    """Generate speech using Google TTS Hindi"""
    if not text or not text.strip():
        raise gr.Error("⚠️ कृपया कुछ हिंदी टेक्स्ट दर्ज करें! (Please enter some Hindi text!)")

    try:
        tts = gTTS(text=text.strip(), lang='hi', slow=slow)
        tmp = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
        tts.save(tmp.name)
        return tmp.name

    except Exception as e:
        raise gr.Error(f"❌ Google TTS Error: {str(e)}")


def generate_both(text, mms_speed, google_slow):
    """Generate speech from both models for comparison"""
    if not text or not text.strip():
        raise gr.Error("⚠️ कृपया कुछ हिंदी टेक्स्ट दर्ज करें! (Please enter some Hindi text!)")

    mms_audio = generate_mms_tts(text, mms_speed)
    google_audio = generate_google_tts(text, google_slow)
    return mms_audio, google_audio


# ============================================================
# 🎨 Custom CSS for Premium UI
# ============================================================
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&family=Noto+Sans+Devanagari:wght@300;400;500;600;700&display=swap');

* {
    font-family: 'Inter', 'Noto Sans Devanagari', sans-serif !important;
}

.gradio-container {
    max-width: 1100px !important;
    margin: 0 auto !important;
    background: linear-gradient(135deg, #0f0c29 0%, #302b63 50%, #24243e 100%) !important;
    min-height: 100vh;
}

.main-header {
    text-align: center;
    padding: 2rem 1rem;
    background: linear-gradient(135deg, rgba(139, 92, 246, 0.15), rgba(59, 130, 246, 0.1));
    border-radius: 20px;
    border: 1px solid rgba(139, 92, 246, 0.2);
    margin-bottom: 1.5rem;
    backdrop-filter: blur(10px);
}

.main-header h1 {
    font-size: 2.5rem !important;
    font-weight: 800 !important;
    background: linear-gradient(135deg, #a78bfa, #818cf8, #60a5fa) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    background-clip: text !important;
    margin-bottom: 0.5rem !important;
    letter-spacing: -0.02em;
}

.main-header p {
    color: #a5b4fc !important;
    font-size: 1.1rem !important;
    font-weight: 400;
    opacity: 0.9;
}

.model-badge {
    display: inline-block;
    padding: 0.3rem 0.8rem;
    border-radius: 20px;
    font-size: 0.8rem;
    font-weight: 600;
    margin: 0.2rem;
}

.badge-meta {
    background: linear-gradient(135deg, #7c3aed, #6d28d9);
    color: white;
}

.badge-google {
    background: linear-gradient(135deg, #2563eb, #1d4ed8);
    color: white;
}

/* Tab Styling */
.tab-nav {
    background: rgba(30, 27, 75, 0.6) !important;
    border-radius: 16px !important;
    padding: 0.4rem !important;
    border: 1px solid rgba(139, 92, 246, 0.2) !important;
}

.tab-nav button {
    border-radius: 12px !important;
    font-weight: 600 !important;
    padding: 0.8rem 1.5rem !important;
    transition: all 0.3s ease !important;
    color: #a5b4fc !important;
    font-size: 0.95rem !important;
}

.tab-nav button.selected {
    background: linear-gradient(135deg, #7c3aed, #6d28d9) !important;
    color: white !important;
    box-shadow: 0 4px 15px rgba(124, 58, 237, 0.4) !important;
}

/* Input/Output Containers */
.gr-block, .gr-box, .gr-panel {
    background: rgba(30, 27, 75, 0.5) !important;
    border: 1px solid rgba(139, 92, 246, 0.15) !important;
    border-radius: 16px !important;
    backdrop-filter: blur(10px) !important;
}

/* Textbox */
textarea, input[type="text"] {
    background: rgba(15, 12, 41, 0.7) !important;
    border: 1px solid rgba(139, 92, 246, 0.3) !important;
    border-radius: 12px !important;
    color: #e0e7ff !important;
    font-size: 1.05rem !important;
    padding: 1rem !important;
    transition: all 0.3s ease !important;
}

textarea:focus, input[type="text"]:focus {
    border-color: #8b5cf6 !important;
    box-shadow: 0 0 20px rgba(139, 92, 246, 0.2) !important;
}

/* Buttons */
.gr-button-primary {
    background: linear-gradient(135deg, #7c3aed, #6d28d9, #5b21b6) !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 0.8rem 2rem !important;
    font-weight: 700 !important;
    font-size: 1rem !important;
    color: white !important;
    box-shadow: 0 4px 15px rgba(124, 58, 237, 0.4) !important;
    transition: all 0.3s ease !important;
    text-transform: uppercase;
    letter-spacing: 0.05em;
}

.gr-button-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 25px rgba(124, 58, 237, 0.5) !important;
}

/* Slider */
input[type="range"] {
    accent-color: #8b5cf6 !important;
}

/* Labels */
label {
    color: #c4b5fd !important;
    font-weight: 600 !important;
    font-size: 0.9rem !important;
}

/* Audio Players */
audio {
    border-radius: 12px !important;
    width: 100% !important;
}

/* Checkbox */
input[type="checkbox"] {
    accent-color: #8b5cf6 !important;
}

/* Examples */
.gr-examples {
    background: rgba(30, 27, 75, 0.3) !important;
    border-radius: 16px !important;
    padding: 1rem !important;
}

/* Footer */
.footer-info {
    text-align: center;
    padding: 1.5rem;
    color: #6366f1;
    font-size: 0.85rem;
    opacity: 0.7;
}

/* Scrollbar */
::-webkit-scrollbar {
    width: 6px;
}
::-webkit-scrollbar-track {
    background: rgba(15, 12, 41, 0.3);
}
::-webkit-scrollbar-thumb {
    background: rgba(139, 92, 246, 0.4);
    border-radius: 3px;
}

/* Animations */
@keyframes pulse-glow {
    0%, 100% { box-shadow: 0 0 20px rgba(139, 92, 246, 0.2); }
    50% { box-shadow: 0 0 40px rgba(139, 92, 246, 0.4); }
}

.generating {
    animation: pulse-glow 2s ease-in-out infinite;
}
"""

# ============================================================
# 🏗️ Build the Gradio App
# ============================================================

# Example Hindi texts
example_texts = [
    "नमस्ते! आज का मौसम बहुत सुहावना है। चलिए बाहर घूमने चलते हैं।",
    "भारत एक विविधताओं से भरा देश है। यहाँ अनेक भाषाएँ, संस्कृतियाँ और परंपराएँ हैं।",
    "कृत्रिम बुद्धिमत्ता का क्षेत्र तेजी से विकसित हो रहा है और यह हमारे जीवन को बदल रहा है।",
    "हिमालय की ऊँची चोटियों से लेकर कन्याकुमारी के समुद्र तट तक, भारत अपनी प्राकृतिक सुंदरता के लिए प्रसिद्ध है।",
    "शिक्षा सबसे शक्तिशाली हथियार है जिसका उपयोग आप दुनिया को बदलने के लिए कर सकते हैं।",
    "प्रौद्योगिकी ने संचार को आसान बना दिया है। अब हम दुनिया के किसी भी कोने से जुड़ सकते हैं।",
]

with gr.Blocks(css=custom_css, title="🎙️ Hindi Text-to-Speech Studio", theme=gr.themes.Base()) as demo:

    # Header
    gr.HTML("""
    <div class="main-header">
        <h1>🎙️ हिंदी Text-to-Speech Studio</h1>
        <p>Premium Hindi speech synthesis powered by AI</p>
        <div style="margin-top: 1rem;">
            <span class="model-badge badge-meta">🔮 Meta MMS-TTS</span>
            <span class="model-badge badge-google">🌐 Google TTS</span>
        </div>
    </div>
    """)

    with gr.Tabs() as tabs:

        # ═══════════════════════════════════════════
        # TAB 1: Meta MMS-TTS Hindi
        # ═══════════════════════════════════════════
        with gr.TabItem("🔮 Meta MMS-TTS Hindi", id="mms"):
            gr.HTML("""
            <div style="padding: 1rem; text-align: center;">
                <p style="color: #c4b5fd; font-size: 0.95rem;">
                    <strong>Facebook's Massively Multilingual Speech</strong> — Neural TTS using VITS architecture.
                    Fast, lightweight, and produces natural-sounding Hindi speech.
                </p>
            </div>
            """)

            with gr.Row():
                with gr.Column(scale=3):
                    mms_text = gr.Textbox(
                        label="📝 हिंदी टेक्स्ट (Hindi Text)",
                        placeholder="यहाँ अपना हिंदी टेक्स्ट लिखें...\nExample: नमस्ते, आप कैसे हैं?",
                        lines=5,
                        max_lines=15,
                        elem_id="mms-text-input"
                    )
                with gr.Column(scale=1):
                    mms_speed = gr.Slider(
                        minimum=0.5,
                        maximum=2.0,
                        value=1.0,
                        step=0.1,
                        label="⚡ Speaking Speed",
                        info="0.5 = Slow, 1.0 = Normal, 2.0 = Fast"
                    )

            mms_btn = gr.Button(
                "🎤 Generate with MMS-TTS",
                variant="primary",
                size="lg",
                elem_id="mms-generate-btn"
            )

            mms_output = gr.Audio(
                label="🔊 Generated Audio (Meta MMS-TTS)",
                type="filepath",
                elem_id="mms-audio-output"
            )

            gr.Examples(
                examples=[[text, 1.0] for text in example_texts[:3]],
                inputs=[mms_text, mms_speed],
                label="📌 Example Hindi Texts"
            )

            mms_btn.click(
                fn=generate_mms_tts,
                inputs=[mms_text, mms_speed],
                outputs=mms_output,
                api_name="mms_tts"
            )

        # ═══════════════════════════════════════════
        # TAB 2: Google TTS Hindi
        # ═══════════════════════════════════════════
        with gr.TabItem("🌐 Google TTS Hindi", id="google"):
            gr.HTML("""
            <div style="padding: 1rem; text-align: center;">
                <p style="color: #93c5fd; font-size: 0.95rem;">
                    <strong>Google Text-to-Speech</strong> — Production-grade TTS engine.
                    Delivers clear, natural Hindi speech with professional voice quality.
                </p>
            </div>
            """)

            with gr.Row():
                with gr.Column(scale=3):
                    google_text = gr.Textbox(
                        label="📝 हिंदी टेक्स्ट (Hindi Text)",
                        placeholder="यहाँ अपना हिंदी टेक्स्ट लिखें...\nExample: भारत एक महान देश है।",
                        lines=5,
                        max_lines=15,
                        elem_id="google-text-input"
                    )
                with gr.Column(scale=1):
                    google_slow = gr.Checkbox(
                        label="🐢 Slow Speech Mode",
                        value=False,
                        info="Enable for slower, clearer pronunciation"
                    )

            google_btn = gr.Button(
                "🎤 Generate with Google TTS",
                variant="primary",
                size="lg",
                elem_id="google-generate-btn"
            )

            google_output = gr.Audio(
                label="🔊 Generated Audio (Google TTS)",
                type="filepath",
                elem_id="google-audio-output"
            )

            gr.Examples(
                examples=[[text, False] for text in example_texts[3:]],
                inputs=[google_text, google_slow],
                label="📌 Example Hindi Texts"
            )

            google_btn.click(
                fn=generate_google_tts,
                inputs=[google_text, google_slow],
                outputs=google_output,
                api_name="google_tts"
            )

        # ═══════════════════════════════════════════
        # TAB 3: Compare Both Models
        # ═══════════════════════════════════════════
        with gr.TabItem("⚔️ Compare Both Models", id="compare"):
            gr.HTML("""
            <div style="padding: 1rem; text-align: center;">
                <p style="color: #a5b4fc; font-size: 0.95rem;">
                    <strong>Side-by-Side Comparison</strong> — Enter text once and hear how both models synthesize it.
                    Compare quality, naturalness, and pronunciation.
                </p>
            </div>
            """)

            compare_text = gr.Textbox(
                label="📝 हिंदी टेक्स्ट (Hindi Text)",
                placeholder="दोनों मॉडल की तुलना के लिए यहाँ टेक्स्ट लिखें...",
                lines=5,
                max_lines=15,
                elem_id="compare-text-input"
            )

            with gr.Row():
                compare_mms_speed = gr.Slider(
                    minimum=0.5,
                    maximum=2.0,
                    value=1.0,
                    step=0.1,
                    label="⚡ MMS Speed"
                )
                compare_google_slow = gr.Checkbox(
                    label="🐢 Google Slow Mode",
                    value=False
                )

            compare_btn = gr.Button(
                "⚔️ Generate & Compare Both",
                variant="primary",
                size="lg",
                elem_id="compare-generate-btn"
            )

            with gr.Row():
                with gr.Column():
                    gr.HTML('<p style="text-align:center; color:#c4b5fd; font-weight:700; font-size:1.1rem;">🔮 Meta MMS-TTS</p>')
                    compare_mms_output = gr.Audio(
                        label="Meta MMS-TTS Output",
                        type="filepath",
                        elem_id="compare-mms-output"
                    )
                with gr.Column():
                    gr.HTML('<p style="text-align:center; color:#93c5fd; font-weight:700; font-size:1.1rem;">🌐 Google TTS</p>')
                    compare_google_output = gr.Audio(
                        label="Google TTS Output",
                        type="filepath",
                        elem_id="compare-google-output"
                    )

            gr.Examples(
                examples=[[text, 1.0, False] for text in example_texts[:4]],
                inputs=[compare_text, compare_mms_speed, compare_google_slow],
                label="📌 Try These Examples"
            )

            compare_btn.click(
                fn=generate_both,
                inputs=[compare_text, compare_mms_speed, compare_google_slow],
                outputs=[compare_mms_output, compare_google_output],
                api_name="compare_both"
            )

    # Footer
    gr.HTML("""
    <div class="footer-info">
        <p>🎙️ Hindi Text-to-Speech Studio | Powered by Meta MMS & Google TTS</p>
        <p style="margin-top: 0.3rem; font-size: 0.75rem;">
            Models: facebook/mms-tts-hin (VITS) • gTTS (Google)
        </p>
    </div>
    """)

# Launch
if __name__ == "__main__":
    demo.launch()