import gradio as gr
import os
import torch

# --- 1. Setup Logic (Lazy Loading) ---
# Hum model ko globally load nahi karenge taaki app turant start ho jaye
vf_model = None

def get_model():
    global vf_model
    if vf_model is None:
        print("Loading VoiceFixer Model...")
        from voicefixer import VoiceFixer
        vf_model = VoiceFixer()
    return vf_model

def fix_audio_ts(audio_path):
    if audio_path is None:
        return None
    
    # Model load karo (agar pehle se loaded nahi hai)
    vf = get_model()
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    output_file = "ts_studio_enhanced.wav"
    
    # Mode 0: Cleaning + Restoration
    vf.restore(input=audio_path, output=output_file, cuda=(device=='cuda'), mode=0)
    
    return output_file

# --- 2. Ultra-Premium TS Audio CSS ---
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;900&display=swap');

body, .gradio-container {
    font-family: 'Inter', sans-serif !important;
    background: #0a0a0a !important;
    color: #ffffff !important;
}

#main_card {
    background: rgba(30, 30, 40, 0.8);
    backdrop-filter: blur(30px);
    border: 1px solid rgba(255, 255, 255, 0.15);
    border-radius: 30px;
    padding: 40px;
    box-shadow: 0 0 60px rgba(0, 255, 136, 0.1), inset 0 0 20px rgba(255, 255, 255, 0.05);
    max-width: 800px;
    margin: 40px auto;
}

#logo_text h1 {
    color: #ffffff;
    font-size: 4.5rem !important;
    font-weight: 900;
    text-align: center;
    margin-bottom: 0px;
    letter-spacing: -2px;
    text-shadow: 0 0 25px rgba(0, 255, 136, 0.6);
}

#subtitle_text h3 {
    text-align: center;
    color: #bbbbbb;
    font-size: 1.2rem;
    font-weight: 400;
    margin-top: 10px;
    margin-bottom: 40px;
    letter-spacing: 1px;
}

.block, .svelte-12cmxck, .form {
    background: rgba(255,255,255,0.03) !important;
    border: 1px solid rgba(255,255,255,0.1) !important;
    border-radius: 20px !important;
    padding: 15px !important;
}

.primary-btn {
    background: linear-gradient(135deg, #00ff88, #00b8ff) !important;
    color: #000 !important;
    border: none !important;
    font-size: 1.2rem !important;
    font-weight: 900 !important;
    padding: 18px !important;
    border-radius: 15px !important;
    box-shadow: 0 0 30px rgba(0, 255, 136, 0.3);
    transition: all 0.3s ease;
    margin-top: 20px;
}
.primary-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 0 50px rgba(0, 255, 136, 0.6);
}
"""

# --- 3. Interface ---
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as app:
    
    with gr.Column(elem_id="main_card"):
        gr.Markdown("# TS", elem_id="logo_text")
        gr.Markdown("### PREMIUM VOCAL ENHANCER", elem_id="subtitle_text")
        
        gr.Markdown("**🎤 Upload Noisy Audio**")
        inp_audio = gr.Audio(type="filepath", label="Input Audio", sources=["upload", "microphone"])
        
        btn = gr.Button("🎹 RESTORE VOICE (Studio Quality)", variant="primary", elem_classes=["primary-btn"])
        
        gr.Markdown("**🎧 Enhanced Result (Clean Voice)**")
        out_audio = gr.Audio(type="filepath", label="Studio Output", interactive=False)

    btn.click(fix_audio_ts, inputs=inp_audio, outputs=out_audio)

app.launch()