import gradio as gr
from gradio_client import Client, handle_file
from transformers import pipeline
from moviepy import ImageClip, AudioFileClip 
import os

# --- ENGINES ---
classifier = pipeline("audio-classification", model="dima806/music_genres_classification")
mastering_client = Client("amaai-lab/SonicMaster")
image_client = Client("black-forest-labs/FLUX.1-schnell")

# --- LOGIC ---

def master_logic(audio_path, brightness, loudness, reverb, steps, cfg):
    if not audio_path: return None, None, "⚠️ Upload audio!", None
    
    results = classifier(audio_path)
    genre = results[0]['label']
    
    # NEW: Professional Engineering Keywords
    # We add "analog warmth" and "high headroom" to fight that "cheap" digital sound.
    quality_boosters = "High-fidelity, analog warmth, 24-bit depth, wide stereo image, no digital clipping, no artifacts."
    
    bright_txt = "shimmering top-end clarity" if brightness > 7 else ("vintage tube warmth" if brightness < 3 else "transparent EQ")
    loud_txt = "commercial loudness with preserved transients" if loudness > 7 else ("audiophile dynamic range" if loudness < 3 else "balanced RMS levels")
    verb_txt = "lush cinematic space" if reverb > 7 else ("dry precision studio" if reverb < 3 else "natural room")
    
    # The "Secret Sauce" Prompt
    full_prompt = f"Professional {genre} master: {bright_txt}, {loud_txt}, {verb_txt}. {quality_boosters}"
    
    result = mastering_client.predict(
        handle_file(audio_path),
        full_prompt,
        42,
        int(steps), 
        float(cfg)
    )
    
    return result[0], result[0], f"✨ Premium {genre} Master Complete", result[0]

def art_logic(status_text, vibe, detail_steps, creativity_cfg):
    if not status_text or "Ready" in status_text: return None
    
    genre = status_text.split("**")[1] if "**" in status_text else "Music"
    prompt = f"Professional album cover art, {genre} style, {vibe}, 4k, ultra-detailed, no text."
    
    img_result = image_client.predict(
        prompt=prompt,
        seed=0,
        width=1024,
        height=1024,
        guidance_scale=float(creativity_cfg),
        num_inference_steps=int(detail_steps),
        api_name="/infer"
    )
    return img_result[0]

def video_logic(audio_path, image_path, volume_level):
    if not audio_path or not image_path: return None
    
    # Load and adjust levels
    audio = AudioFileClip(audio_path).subclipped(0, 30).volumex(volume_level)
    img = ImageClip(image_path).with_duration(audio.duration).resized(width=1080)
    
    # Composition
    video = img.on_color(size=(1080, 1920), color=(15, 15, 15), pos="center").with_audio(audio)
    
    out_path = "final_promo.mp4"
    video.write_videofile(out_path, fps=24, codec="libx264", audio_codec="aac")
    return out_path

# --- UI LAYOUT ---
with gr.Blocks(theme=gr.themes.Base()) as demo:
    gr.HTML("<h1 style='text-align: center;'>🎚️ AI Studio: Effect Console</h1>")
    
    raw_storage = gr.State()
    master_storage = gr.State()

    with gr.Tabs():
        # TAB 1: MASTERING
        with gr.TabItem("🎧 Mastering Console"):
            with gr.Row():
                with gr.Column(scale=1):
                    in_audio = gr.Audio(label="Raw Upload", type="filepath")
                    master_btn = gr.Button("🚀 APPLY MASTERING", variant="primary")
                
                with gr.Column(scale=1):
                    gr.Markdown("### FX Levels")
                    bright_sl = gr.Slider(0, 10, value=5, label="Brightness (EQ)")
                    loud_sl = gr.Slider(0, 10, value=5, label="Loudness (Limiters)")
                    verb_sl = gr.Slider(0, 10, value=2, label="Space (Reverb)")
                    
                    with gr.Accordion("Advanced Engine Settings", open=False):
                        m_steps = gr.Slider(10, 50, value=25, step=1, label="Refinement Steps")
                        m_cfg = gr.Slider(1.0, 15.0, value=3.5, step=0.5, label="AI Guidance Scale")

            status = gr.Markdown("Ready.")
            with gr.Row():
                monitor = gr.Audio(label="Studio Monitor")
                export_wav = gr.File(label="Download Master")

        # TAB 2: ARTWORK
        with gr.TabItem("🎨 Art Direction"):
            with gr.Row():
                with gr.Column():
                    vibe_in = gr.Textbox(label="Visual Theme", placeholder="e.g. Neon Cyberpunk, 70s Vinyl...")
                    art_btn = gr.Button("🎨 GENERATE ART")
                    
                    gr.Markdown("### Image 'Levels'")
                    art_steps = gr.Slider(1, 4, value=4, step=1, label="Detail Iterations")
                    art_cfg = gr.Slider(0.0, 5.0, value=3.5, label="Prompt Strictness")
                
                art_out = gr.Image(label="Cover Art")

        # TAB 3: PROMO EXPORT
        with gr.TabItem("🎬 Video Export"):
            with gr.Row():
                with gr.Column():
                    vol_sl = gr.Slider(0.0, 2.0, value=1.0, label="Final Video Volume (Gain)")
                    promo_btn = gr.Button("🎬 RENDER TIKTOK CLIP", variant="primary")
                video_out = gr.Video(label="9:16 Social Preview")

    # --- WIRING ---
    master_btn.click(
        master_logic, 
        [in_audio, bright_sl, loud_sl, verb_sl, m_steps, m_cfg], 
        [monitor, export_wav, status, master_storage]
    ).then(lambda x: x, in_audio, raw_storage)

    art_btn.click(art_logic, [status, vibe_in, art_steps, art_cfg], art_out)
    
    promo_btn.click(video_logic, [master_storage, art_out, vol_sl], video_out)

demo.launch()