import os
import torch
import gradio as gr
import edge_tts
import asyncio
import subprocess
from openvoice import se_extractor
from openvoice.api import ToneColorConverter

# ==========================================
# ၁။ Model Load လုပ်ခြင်း
# ==========================================
checkpoint_path = "checkpoint.pth" 
config_path = "config.json"
tone_color_converter = None
model_error = None

if os.path.exists(checkpoint_path) and os.path.exists(config_path):
    try:
        device = "cpu"
        tone_color_converter = ToneColorConverter(config_path, device=device)
        tone_color_converter.load_ckpt(checkpoint_path)
        print("✅ Model Loaded!")
    except Exception as e:
        model_error = str(e)
else:
    model_error = "Files not found!"

os.makedirs("processed", exist_ok=True)

# ==========================================
# ၂။ Professional Mastering Engine (FFmpeg)
# ==========================================
def apply_mastering(input_wav, volume_gain, auto_master, treble_boost):
    output_wav = "output_mastered.wav"
    
    # FFmpeg Filters Chain တည်ဆောက်ခြင်း
    filters = []
    
    # 1. Volume Gain (အသံ အတိုးအကျယ်)
    # 0dB က ပုံမှန်၊ +5dB ဆိုရင် ပိုကျယ်မယ်
    if volume_gain != 0:
        filters.append(f"volume={volume_gain}dB")
    
    # 2. Treble Boost (အသံကြည်လင်စေရန် High Frequency တင်ခြင်း)
    if treble_boost:
        filters.append("treble=g=5")

    # 3. Auto Mastering (Loudness Normalization)
    # ဒါက Studio Quality ရအောင် အသံကို Compression လုပ်ပေးပါတယ်
    if auto_master:
        filters.append("loudnorm=I=-16:TP=-1.5:LRA=11")

    # Filter မရှိရင် မူရင်းအတိုင်းပြန်ပေး
    if not filters:
        return input_wav

    # FFmpeg Command Run ခြင်း
    filter_str = ",".join(filters)
    cmd = [
        "ffmpeg", "-y", "-i", input_wav,
        "-af", filter_str,
        output_wav
    ]
    
    print(f"Applying Mastering Filters: {filter_str}")
    subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    
    if os.path.exists(output_wav):
        return output_wav
    else:
        return input_wav # Error တက်ရင် မူရင်းပဲ ပြန်ပေးမယ်

# ==========================================
# ၃။ Main Process
# ==========================================
async def generate_base_audio(text, voice, speed, pitch):
    output_file = "temp_base.mp3"
    rate = f"{int(speed)}%" if speed < 0 else f"+{int(speed)}%"
    pitch_str = f"{int(pitch)}Hz" if pitch < 0 else f"+{int(pitch)}Hz"
    communicate = edge_tts.Communicate(text, voice, rate=rate, pitch=pitch_str)
    await communicate.save(output_file)
    return output_file

def voice_clone(text, gender, ref_audio, speed, pitch, volume, auto_master, treble):
    if model_error: return None, f"❌ {model_error}"
    if tone_color_converter is None: return None, "❌ Model Error"

    try:
        base_voice = "my-MM-NilarNeural" if gender == "Female (မ)" else "my-MM-ThihaNeural"
        
        # Step 1: Base Audio
        base_audio_path = asyncio.run(generate_base_audio(text, base_voice, speed, pitch))

        # Step 2: OpenVoice Conversion (Raw Audio)
        source_se, _ = se_extractor.get_se(base_audio_path, tone_color_converter, target_dir='processed', vad=True)
        target_se, _ = se_extractor.get_se(ref_audio, tone_color_converter, target_dir='processed', vad=True)
        
        raw_output = "output_raw.wav"
        tone_color_converter.convert(
            audio_src_path=base_audio_path, src_se=source_se, tgt_se=target_se, 
            output_path=raw_output
        )

        # Step 3: Mastering & Boosting
        final_output = apply_mastering(raw_output, volume, auto_master, treble)
        
        return final_output, "✅ အောင်မြင်ပါတယ် (Mastered)"
        
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# ==========================================
# ၄။ UI Design (Pro Version)
# ==========================================
with gr.Blocks(title="Myanmar Voice Master") as demo:
    gr.Markdown("# 🇲🇲 Myanmar Voice Cloning (Mastering Edition)")
    
    if model_error: gr.Warning(model_error)
    
    with gr.Row():
        with gr.Column():
            # Input Section
            inp_text = gr.Textbox(label="စာသား (Text)", lines=3, value="မင်္ဂလာပါ")
            with gr.Row():
                inp_gender = gr.Dropdown(["Male (ကျား)", "Female (မ)"], value="Female (မ)", label="Gender")
                inp_ref = gr.Audio(label="Reference Audio", type="filepath")
            
            # Basic Controls
            with gr.Accordion("⚙️ Basic Settings (အခြေခံ)", open=False):
                slider_speed = gr.Slider(-50, 50, value=0, step=5, label="Speed % (အမြန်နှုန်း)")
                slider_pitch = gr.Slider(-20, 20, value=0, step=1, label="Pitch Hz (အသံနိမ့်မြင့်)")

            # Mastering Controls (Highlight)
            gr.Markdown("### 🎚️ Studio Mastering (အသံပြုပြင်ရန်)")
            with gr.Group():
                # Volume Booster
                slider_vol = gr.Slider(0, 20, value=5, step=1, label="Volume Booster (+dB) - အသံကျယ်ရန်")
                
                # Checkboxes
                with gr.Row():
                    chk_master = gr.Checkbox(label="Auto Mastering (အသံအရည်အသွေး အလိုအလျောက်ညှိမည်)", value=True)
                    chk_treble = gr.Checkbox(label="Treble Boost (အသံပိုကြည်အောင်လုပ်မည်)", value=False)

            btn = gr.Button("Generate Mastered Audio", variant="primary")
        
        with gr.Column():
            out_audio = gr.Audio(label="Final Output")
            out_status = gr.Textbox(label="Status")

    btn.click(
        fn=voice_clone, 
        inputs=[inp_text, inp_gender, inp_ref, slider_speed, slider_pitch, slider_vol, chk_master, chk_treble], 
        outputs=[out_audio, out_status]
    )

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)