import gradio as gr
from gradio_client import Client, handle_file
from transformers import pipeline

# Engines
classifier = pipeline("audio-classification", model="MIT/ast-finetuned-audioset")
mastering_client = Client("amaai-lab/SonicMaster")

def master_process(audio_path):
    if audio_path is None: return None, None, "Upload audio first."
    
    # 1. Analyze Genre
    results = classifier(audio_path)
    top_genre = results[0]['label']
    
    # 2. Process
    result = mastering_client.predict(
        audio=handle_file(audio_path),
        prompt=f"Professional {top_genre} master.",
        api_name="/predict"
    )
    mastered_path = result[1]
    
    # Return: Mastered File, Status, and update the internal State
    return mastered_path, mastered_path, f"Mode: Mastered ({top_genre})", mastered_path

def toggle_audio(choice, raw_path, mastered_path):
    # This function switches the audio source based on the toggle
    if choice == "Mastered ✨":
        return mastered_path
    return raw_path

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎚️ Pro Mastering Suite with A/B Bypass")
    
    # Internal storage for the audio paths
    raw_storage = gr.State()
    mastered_storage = gr.State()
    
    with gr.Row():
        with gr.Column():
            in_audio = gr.Audio(label="1. Upload Raw Demo", type="filepath")
            run_btn = gr.Button("🚀 MASTER AUDIO", variant="primary")
            
        with gr.Column():
            # The Monitor Player
            monitor_player = gr.Audio(label="2. Monitor Output", interactive=False)
            
            # The Bypass Toggle
            mode_toggle = gr.Radio(
                choices=["Original 🔈", "Mastered ✨"], 
                label="Bypass Toggle (A/B Test)",
                value="Mastered ✨",
                interactive=True
            )
            
            download_link = gr.File(label="3. Export Final WAV")

    # Mastering Trigger
    run_btn.click(
        fn=master_process,
        inputs=in_audio,
        outputs=[monitor_player, download_link, monitor_player.label, mastered_storage]
    ).then(lambda x: x, inputs=in_audio, outputs=raw_storage)

    # Toggle Trigger (Changes the player source instantly)
    mode_toggle.change(
        fn=toggle_audio,
        inputs=[mode_toggle, raw_storage, mastered_storage],
        outputs=monitor_player
    )

demo.launch()