import os, sys
import gradio as gr

from main import song_cover_pipeline
from audio_effects import add_audio_effects
from modules.model_management import ignore_files, update_models_list, download_from_url, upload_zip_model, upload_separate_files
from modules.ui_updates import show_hop_slider, update_f0_method, update_button_text_voc, update_button_text_inst, swap_visibility, swap_buttons
from modules.file_processing import process_file_upload

# Setup directories
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
rvc_models_dir = os.path.join(BASE_DIR, 'rvc_models')

# Check for warnings
warning = sys.argv[1] if len(sys.argv) > 1 else 'False'
warning = True if warning == 'True' else False

# Initialize voice models
voice_models = ignore_files(rvc_models_dir)

# Define the Gradio interface
with gr.Blocks(title="🎤 RVC Inference", css="footer{display:none !important}") as app:
    
    # Show warning message if applicable
    if warning:
        gr.HTML("<center><h2>⚠️ This space is running slowly due to limited resources. Use the faster <a href='https://colab.research.google.com/drive/1HzuPgICRrjqUWQWb5Zn-l07m099-n-Nr'>Google Colab Notebook</a>.</h2></center>")
    
    # Main Tab: Voice Conversion
    with gr.Tab("🎶 Voice Conversion"):
        with gr.Row():
            rvc_model = gr.Dropdown(voice_models, label='🎙️ Voice Models')
            output_format = gr.Dropdown(['mp3', 'flac', 'wav'], value='mp3', label='🎧 Output Format')
            pitch = gr.Slider(-12, 12, value=0, step=1, label='🎚️ Pitch Adjustment', info='-12 for deeper voice, +12 for higher voice')

        with gr.Row():
            refresh_btn = gr.Button('🔄 Refresh Models')
            generate_btn = gr.Button("🎼 Generate", variant="primary", scale=2)

        with gr.Column():
            audio_input = gr.Audio(label='🎵 Upload Audio', interactive=True)
            converted_audio = gr.Audio(label='🎶 Converted Voice')

        with gr.Accordion('⚙️ Voice Conversion Settings', open=False):
            use_hybrid_methods = gr.Checkbox(label="🧬 Use Hybrid Methods", value=False)
            f0_method = gr.Dropdown(['rmvpe+', 'fcpe', 'rmvpe', 'mangio-crepe', 'crepe'], value='rmvpe+', label='🔧 F0 Method')
            use_hybrid_methods.change(update_f0_method, inputs=use_hybrid_methods, outputs=f0_method)
            crepe_hop_length = gr.Slider(8, 512, value=128, step=8, visible=False, label='🎛️ Crepe Hop Length')
            f0_method.change(show_hop_slider, inputs=f0_method, outputs=crepe_hop_length)
        
        refresh_btn.click(update_models_list, None, outputs=rvc_model)
        generate_btn.click(song_cover_pipeline, inputs=[audio_input, rvc_model, pitch, f0_method, crepe_hop_length, output_format], outputs=[converted_audio])

    # Merge/Process Tab
    with gr.Tab('🔄 Merge/Process'):
        with gr.Row():
            vocal_input = gr.Audio(label='🎤 Vocals', interactive=True)
            instrumental_input = gr.Audio(label='🎹 Instrumental', interactive=True)
        
        with gr.Row():
            process_btn = gr.Button("🎬 Process")
            ai_cover_output = gr.Audio(label='🎧 AI Cover')

        process_btn.click(add_audio_effects, inputs=[vocal_input, instrumental_input], outputs=[ai_cover_output])

    # Model Options Tab
    with gr.Tab('📦 Model Options'):
        with gr.Tab('🔗 Download Models'):
            model_zip_link = gr.Text(label='🔗 Enter Model Download Link')
            model_name = gr.Text(label='📝 Model Name')
            download_btn = gr.Button('⬇️ Download Model')
            dl_output_message = gr.Text(label='📢 Output Message')
            download_btn.click(download_from_url, inputs=[model_zip_link, model_name], outputs=dl_output_message)

        with gr.Tab('📤 Upload ZIP Models'):
            zip_file = gr.File(label='📁 Upload ZIP File', file_types=['.zip'])
            zip_model_name = gr.Text(label='📝 Model Name')
            upload_btn = gr.Button('📤 Upload Model')
            upload_output_message = gr.Text(label='📢 Upload Status')
            upload_btn.click(upload_zip_model, inputs=[zip_file, zip_model_name], outputs=upload_output_message)

        with gr.Tab('📤 Upload Separate Files'):
            pth_file = gr.File(label='📂 Upload .pth File', file_types=['.pth'])
            index_file = gr.File(label='📂 Upload .index File', file_types=['.index'])
            sep_model_name = gr.Text(label='📝 Model Name')
            separate_upload_btn = gr.Button('📤 Upload Files')
            separate_upload_output = gr.Text(label='📢 Upload Status')
            separate_upload_btn.click(upload_separate_files, inputs=[pth_file, index_file, sep_model_name], outputs=separate_upload_output)

app.launch(share=True)