import gradio as gr
import os
import urllib.request
from pydub import AudioSegment

# तीनों ओरिजिनल Piper मॉडल्स डाउनलोड और चेक करने का फंक्शन
def ensure_model_exists(voice_choice):
    base_url = "https://huggingface.co/rhasspy/piper-voices/resolve/main/hi/hi_IN"
    
    if voice_choice == "Female (Priyamvada)":
        model_name = "hi_IN-priyamvada-medium"
        url_path = f"{base_url}/priyamvada/medium"
    elif voice_choice == "Male (Pratham)":
        model_name = "hi_IN-pratham-medium"
        url_path = f"{base_url}/pratham/medium"
    else:  # Male (Rohan)
        model_name = "hi_IN-rohan-medium"
        url_path = f"{base_url}/rohan/medium"
        
    onnx_file = f"{model_name}.onnx"
    json_file = f"{model_name}.onnx.json"
    
    if not os.path.exists(onnx_file):
        print(f"Downloading {onnx_file}...")
        urllib.request.urlretrieve(f"{url_path}/{onnx_file}", onnx_file)
    if not os.path.exists(json_file):
        print(f"Downloading {json_file}...")
        urllib.request.urlretrieve(f"{url_path}/{json_file}", json_file)
        
    return onnx_file, model_name

# कैरेक्टर काउंट दिखाने का फंक्शन
def count_characters(text):
    return f"📝 **अक्षरों की संख्या (Character Count):** {len(text)}"

def generate_news_audio(text, voice_choice, speed, apply_enhance):
    # 1. टेक्स्ट क्लीनिंग: स्टार (**) और गैर-जरूरी शब्द ऑटोमैटिक हटाना
    text = text.replace("मैं हूँ आपका न्यूज़ होस्ट", "")
    text = text.replace("**", "")  
    text = text.replace("*", "")   
    text = text.strip()
    
    if not text:
        return None, "⚠️ कृपया टेक्स्ट दर्ज करें।"
        
    try:
        # 2. मॉडल सेटअप
        onnx_file, model_name = ensure_model_exists(voice_choice)
        raw_output = "raw_audio.wav"
        final_output = "news_audio.wav"
        
        # 3. स्पीड सेट करना (1.0 नॉर्मल है)
        length_scale = 1.0 / float(speed)
        
        # 4. वाक्य के बीच रुकने का समय (0.4 सेकंड)
        sentence_silence = 0.4 
        
        # 5. Piper से Raw ऑडियो जनरेट करना
        command = f'echo "{text}" | piper --model ./{onnx_file} --length_scale {length_scale} --sentence_silence {sentence_silence} --output_file {raw_output}'
        os.system(command)
        
        if os.path.exists(raw_output) and os.path.getsize(raw_output) > 0:
            
            # --- स्टूडियो एनहांसमेंट लॉजिक (Studio Enhancement) ---
            if apply_enhance:
                # Raw ऑडियो को लोड करें
                audio = AudioSegment.from_wav(raw_output)
                
                # A. वॉल्यूम को लाउड और बैलेंस करें (Normalization)
                audio = audio.normalize()
                
                # B. रोबोटिक बेस और Hiss हटाने के लिए High-Pass Filter (80Hz)
                audio = audio.high_pass_filter(80)
                
                # एनहांस्ड ऑडियो को सेव करें
                audio.export(final_output, format="wav")
            else:
                # अगर एनहांसमेंट बंद है, तो बिना बदलाव के फाइल का नाम बदल दें
                os.rename(raw_output, final_output)
                
            status_msg = f"✅ ऑडियो तैयार है! (Active Model: **{model_name}** | Enhanced: **{'Yes' if apply_enhance else 'No'}**)"
            return final_output, status_msg
        else:
            return None, "❌ एरर: सर्वर ऑडियो जनरेट नहीं कर पाया।"
            
    except Exception as e:
        return None, f"❌ एरर: {str(e)}"

# --- UI Setup ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as app:
    gr.Markdown("# 🎙️ Open-Source Hindi News TTS (Pro Studio)")
    
    gr.Markdown("### ⚙️ **Available Models:** Piper TTS (`Priyamvada`, `Pratham`, `Rohan`)")
    gr.Markdown("💡 **सुझाव:** मॉडल को सही से रुकने का निर्देश देने के लिए स्क्रिप्ट में **कॉमा (,)** और **फुल स्टॉप (। या .)** का इस्तेमाल ज़रूर करें। (स्क्रिप्ट पेस्ट करते ही स्टार ** ऑटोमैटिक हट जाएंगे)")
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="न्यूज़ स्क्रिप्ट डालें", 
                lines=7, 
                placeholder="आज की मुख्य खबरें इस प्रकार हैं। भारत निर्वाचन आयोग ने..."
            )
            
            # यहाँ लाइव कैरेक्टर काउंट दिखेगा
            char_count_display = gr.Markdown("📝 **अक्षरों की संख्या (Character Count):** 0")
            
            voice_dropdown = gr.Dropdown(
                choices=["Female (Priyamvada)", "Male (Pratham)", "Male (Rohan)"], 
                value="Female (Priyamvada)", 
                label="आवाज़ चुनें (Voice)"
            )
            
            speed_slider = gr.Slider(minimum=0.5, maximum=1.5, value=0.9, step=0.1, label="स्पीड (Speed) - न्यूज़ के लिए 0.9 या 1.0 सबसे सही है")
            
            # --- नया एनहांसमेंट चेकबॉक्स ---
            enhance_checkbox = gr.Checkbox(
                label="🎙️ स्टूडियो एनहांसमेंट (Studio Enhancement) - आवाज़ को न्यूज़-रूम जैसी लाउड और क्लियर बनाने के लिए इसे चालू रखें", 
                value=True
            )
            
            generate_btn = gr.Button("Generate Audio", variant="primary")
            
        with gr.Column():
            status_output = gr.Markdown("स्टेटस यहाँ दिखेगा...")
            audio_output = gr.Audio(label="जनरेटेड ऑडियो (यहाँ से डाउनलोड करें)", type="filepath")
            
    # टेक्स्ट टाइप या पेस्ट करते ही कैरेक्टर काउंट अपडेट होगा
    text_input.change(fn=count_characters, inputs=text_input, outputs=char_count_display)
            
    # यहाँ स्पेसिंग की दिक्कत थी, जिसे अब ठीक कर दिया गया है
    generate_btn.click(
        fn=generate_news_audio, 
        inputs=[text_input, voice_dropdown, speed_slider, enhance_checkbox], 
        outputs=[audio_output, status_output],
        api_name="predict"
    )

if __name__ == "__main__":
    app.launch()