import gradio as gr
import json
import os
import shutil
import subprocess
import sys

def install_and_import(package):
    try:
        __import__(package)
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])

def generate_song(lyrics, audio_file):
    if not lyrics.strip():
        return None, "❌ Please enter some lyrics!"
    
    if audio_file is None:
        return None, "❌ Please upload an audio prompt file!"
    
    try:
        yield None, "🎵 Installing dependencies..."
        
        # Install required packages on the fly
        packages = ["torch", "torchaudio", "transformers", "accelerate", "omegaconf", "lightning", "einops", "numba", "librosa", "soundfile"]
        for package in packages:
            try:
                install_and_import(package)
            except:
                pass
        
        yield None, "🎵 Setting up SongBloom environment..."
        
        # Create input files
        audio_filename = "prompt.wav"
        shutil.copy(audio_file, audio_filename)
        
        jsonl_filename = "input.jsonl"
        input_data = {
            "idx": "1",
            "lyrics": lyrics.strip(),
            "prompt_wav": audio_filename
        }
        
        with open(jsonl_filename, 'w') as f:
            json.dump(input_data, f)
        
        yield None, "🎵 Downloading SongBloom model (first time only - 5-10 minutes)..."
        
        # Download SongBloom if not present
        if not os.path.exists("songbloom_repo"):
            subprocess.run(["git", "clone", "https://github.com/cypress-yang/SongBloom.git", "songbloom_repo"], check=True)
        
        yield None, "🎵 Running AI music generation (3-8 minutes)..."
        
        # Run SongBloom
        cmd = ["python3", "songbloom_repo/infer.py", "--input-jsonl", jsonl_filename, "--dtype", "bfloat16"]
        
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, universal_newlines=True)
        
        for line in process.stdout:
            if "it [" in line:
                yield None, f"🎵 Generating music... {line.strip()}"
        
        process.wait()
        
        
        # Look for output files more thoroughly
        output_dirs = ["output", "results", "generated", "songbloom_repo/output", "songbloom_repo"]
        output_file = None
        
        for output_dir in output_dirs:
            if os.path.exists(output_dir):
                for file in os.listdir(output_dir):
                    if file.endswith(('.wav', '.mp3', '.flac')) and not file.startswith('prompt'):
                        output_file = os.path.join(output_dir, file)
                        break
                if output_file:
                    break
                    
        if output_file and os.path.exists(output_file):
            yield output_file, "✅ Your AI song is ready! 🎵"
        else:
            yield None, "❌ Generation completed but output file not found. Check model setup."
    except Exception as e:
        yield None, f"❌ Error: {str(e)}"


# Create interface
with gr.Blocks(title="SongBloom AI Music Generator") as demo:
    gr.Markdown("""
    # 🎵 SongBloom AI Music Generator
    
    **Full AI Music Generation - From Lyrics to Complete Songs**
    
    Upload an audio style prompt and enter lyrics to generate professional AI music!
    """)
    
    with gr.Row():
        with gr.Column():
            lyrics_input = gr.Textbox(
                label="Song Lyrics",
                placeholder="Enter your song lyrics here...\n\nExample:\nHere comes the sun, doo-doo-doo-doo\nHere comes the sun, and I say it's all right",
                lines=8,
                value="Here comes the sun, doo-doo-doo-doo\nHere comes the sun, and I say it's all right"
            )
            
            audio_input = gr.File(
                label="Audio Style Prompt (Any Audio Format)",
                type="filepath"
            )
            
            generate_btn = gr.Button("🎵 Generate Full Song", variant="primary", size="lg")
        
        with gr.Column():
            status_output = gr.Textbox(
                label="Generation Progress",
                lines=8,
                interactive=False
            )
            
            audio_output = gr.Audio(label="🎵 Your AI-Generated Song")
    
    generate_btn.click(
        fn=generate_song,
        inputs=[lyrics_input, audio_input],
        outputs=[audio_output, status_output],
        show_progress=True
    )
    
    gr.Markdown("""
    ---
    **⚡ Powered by SongBloom AI - Professional Music Generation**
    
    First generation may take 10-15 minutes (downloading models). Subsequent generations: 3-8 minutes.
    """)

if __name__ == "__main__":
    demo.launch()