import gradio as gr
import spaces
import torch
import torchaudio
import io
import base64
import uuid
import os
import time
import re
import threading
import gc
import random
import numpy as np
from einops import rearrange
from huggingface_hub import login
from stable_audio_tools import get_pretrained_model
from stable_audio_tools.inference.generation import generate_diffusion_cond
from gradio_client import Client, handle_file
from contextlib import contextmanager

# Global model storage
model_cache = {}
model_lock = threading.Lock()

@contextmanager
def resource_cleanup():
    """Lightweight context manager - let zerogpu handle memory management"""
    try:
        yield
    finally:
        # Minimal cleanup - let zerogpu handle the heavy lifting
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        # Removed aggressive empty_cache() and gc.collect() calls

def load_stable_audio_model():
    """Load stable-audio-open-small model if not already loaded."""
    with model_lock:
        if 'stable_audio_model' not in model_cache:
            print("🔄 Loading stable-audio-open-small model...")
            load_start = time.time()
            
            # Authenticate with HF
            hf_token = os.getenv('HF_TOKEN')
            if hf_token:
                login(token=hf_token)
                print(f"✅ HF authenticated")
            
            # Load model
            model, config = get_pretrained_model("stabilityai/stable-audio-open-small")
            device = "cuda" if torch.cuda.is_available() else "cpu"
            model = model.to(device)
            if device == "cuda":
                model = model.half()
            
            load_time = time.time() - load_start
            print(f"✅ Model loaded on {device} in {load_time:.2f}s")
            
            # Aggressive model persistence - warm up with dummy generation
            print("🔥 Warming up model...")
            warmup_start = time.time()
            try:
                dummy_conditioning = [{"prompt": "test", "seconds_total": 12}]
                with torch.no_grad():
                    _ = generate_diffusion_cond(
                        model,
                        steps=1,  # Minimal steps for warmup
                        cfg_scale=1.0,
                        conditioning=dummy_conditioning,
                        sample_size=config["sample_size"],
                        sampler_type="pingpong",
                        device=device,
                        seed=42
                    )
                warmup_time = time.time() - warmup_start
                print(f"🔥 Model warmed up in {warmup_time:.2f}s")
            except Exception as e:
                print(f"⚠️ Warmup failed (but continuing): {e}")
            
            model_cache['stable_audio_model'] = model
            model_cache['stable_audio_config'] = config
            model_cache['stable_audio_device'] = device
            print(f"✅ Stable Audio model ready for fast generation!")
        else:
            print("♻️ Using cached model (should be fast!)")
        
        return (model_cache['stable_audio_model'], 
                model_cache['stable_audio_config'], 
                model_cache['stable_audio_device'])

@spaces.GPU(duration=12)
def generate_stable_audio_loop(prompt, loop_type, bpm, bars, seed=-1):
    """Generate a BPM-aware loop using stable-audio-open-small"""
    try:
        total_start = time.time()
        
        # Model loading timing
        load_start = time.time()
        model, config, device = load_stable_audio_model()
        load_time = time.time() - load_start
        
        # Calculate loop duration based on BPM and bars
        seconds_per_beat = 60.0 / bpm
        seconds_per_bar = seconds_per_beat * 4  # 4/4 time
        target_loop_duration = seconds_per_bar * bars
        
        # Enhance prompt based on loop type and BPM
        if loop_type == "drums":
            enhanced_prompt = f"{prompt} drum loop {bpm}bpm"
            negative_prompt = "melody, harmony, pitched instruments, vocals, singing"
        else:  # instruments
            enhanced_prompt = f"{prompt} instrumental loop {bpm}bpm"
            negative_prompt = "drums, percussion, kick, snare, hi-hat"
        
        # Set seed
        if seed == -1:
            seed = random.randint(0, 2**32 - 1)
        
        torch.manual_seed(seed)
        if device == "cuda":
            torch.cuda.manual_seed(seed)
        
        print(f"🎵 Generating {loop_type} loop:")
        print(f"   Enhanced prompt: {enhanced_prompt}")
        print(f"   Target duration: {target_loop_duration:.2f}s ({bars} bars at {bpm}bpm)")
        print(f"   Seed: {seed}")
        
        # Prepare conditioning
        conditioning_start = time.time()
        conditioning = [{
            "prompt": enhanced_prompt,
            "seconds_total": 12  # Model generates 12s max
        }]
        
        negative_conditioning = [{
            "prompt": negative_prompt,
            "seconds_total": 12
        }]
        conditioning_time = time.time() - conditioning_start
        
        # Generation timing
        generation_start = time.time()
        
        # Removed aggressive resource cleanup wrapper
        # Clear GPU cache once before generation (not after)
        if device == "cuda":
            torch.cuda.empty_cache()
        
        with torch.cuda.amp.autocast(enabled=(device == "cuda")):
            output = generate_diffusion_cond(
                model,
                steps=8,  # Fast generation
                cfg_scale=1.0,  # Good balance for loops
                conditioning=conditioning,
                negative_conditioning=negative_conditioning,
                sample_size=config["sample_size"],
                sampler_type="pingpong",
                device=device,
                seed=seed
            )
        
        generation_time = time.time() - generation_start
        
        # Post-processing timing
        postproc_start = time.time()
        
        # Post-process audio
        output = rearrange(output, "b d n -> d (b n)")  # (2, N) stereo
        output = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1)
        
        # Extract the loop portion
        sample_rate = config["sample_rate"]
        loop_samples = int(target_loop_duration * sample_rate)
        available_samples = output.shape[1]
        
        if loop_samples > available_samples:
            loop_samples = available_samples
            actual_duration = available_samples / sample_rate
            print(f"⚠️ Requested {target_loop_duration:.2f}s, got {actual_duration:.2f}s")
        
        # Extract loop from beginning (cleanest beat alignment)
        loop_output = output[:, :loop_samples]
        loop_output_int16 = loop_output.mul(32767).to(torch.int16).cpu()
        
        # Save to temporary file
        loop_filename = f"loop_{loop_type}_{bpm}bpm_{bars}bars_{seed}.wav"
        torchaudio.save(loop_filename, loop_output_int16, sample_rate)
        
        postproc_time = time.time() - postproc_start
        total_time = time.time() - total_start
        actual_duration = loop_samples / sample_rate
        
        # Detailed timing breakdown
        print(f"⏱️ Timing breakdown:")
        print(f"   Model load: {load_time:.2f}s")
        print(f"   Conditioning: {conditioning_time:.3f}s") 
        print(f"   Generation: {generation_time:.2f}s")
        print(f"   Post-processing: {postproc_time:.3f}s")
        print(f"   Total: {total_time:.2f}s")
        print(f"✅ {loop_type.title()} loop: {actual_duration:.2f}s audio in {total_time:.2f}s")
        
        return loop_filename, f"Generated {actual_duration:.2f}s {loop_type} loop at {bpm}bpm ({bars} bars) in {total_time:.2f}s"
    
    except Exception as e:
        print(f"❌ Generation error: {str(e)}")
        return None, f"Error: {str(e)}"

def combine_loops(drums_audio, instruments_audio, bpm, bars, num_repeats):
    """Combine drum and instrument loops with specified repetitions"""
    try:
        if not drums_audio and not instruments_audio:
            return None, "No audio files to combine"
        
        # Calculate timing
        seconds_per_beat = 60.0 / bpm
        seconds_per_bar = seconds_per_beat * 4
        loop_duration = seconds_per_bar * bars
        total_duration = loop_duration * num_repeats
        
        print(f"🎛️ Combining loops:")
        print(f"   Loop duration: {loop_duration:.2f}s ({bars} bars)")
        print(f"   Repeats: {num_repeats}")
        print(f"   Total duration: {total_duration:.2f}s")
        
        combined_audio = None
        sample_rate = None
        
        # Process each audio file
        for audio_path, audio_type in [(drums_audio, "drums"), (instruments_audio, "instruments")]:
            if audio_path:
                # Load audio
                waveform, sr = torchaudio.load(audio_path)
                if sample_rate is None:
                    sample_rate = sr
                
                # Ensure we have the exact loop duration
                target_samples = int(loop_duration * sr)
                if waveform.shape[1] > target_samples:
                    waveform = waveform[:, :target_samples]
                elif waveform.shape[1] < target_samples:
                    # Pad if necessary
                    padding = target_samples - waveform.shape[1]
                    waveform = torch.cat([waveform, torch.zeros(waveform.shape[0], padding)], dim=1)
                
                # Repeat the loop
                repeated_waveform = waveform.repeat(1, num_repeats)
                
                print(f"   {audio_type}: {waveform.shape[1]/sr:.2f}s repeated {num_repeats}x = {repeated_waveform.shape[1]/sr:.2f}s")
                
                # Add to combined audio
                if combined_audio is None:
                    combined_audio = repeated_waveform
                else:
                    combined_audio = combined_audio + repeated_waveform
        
        if combined_audio is None:
            return None, "No valid audio to combine"
        
        # Normalize to prevent clipping
        combined_audio = combined_audio / torch.max(torch.abs(combined_audio))
        combined_audio = combined_audio.clamp(-1, 1)
        
        # Convert to int16 and save
        combined_audio_int16 = combined_audio.mul(32767).to(torch.int16)
        combined_filename = f"combined_{bpm}bpm_{bars}bars_{num_repeats}loops_{random.randint(1000, 9999)}.wav"
        torchaudio.save(combined_filename, combined_audio_int16, sample_rate)
        
        actual_duration = combined_audio.shape[1] / sample_rate
        status = f"Combined into {actual_duration:.2f}s audio ({num_repeats} × {bars} bars at {bpm}bpm)"
        
        print(f"✅ {status}")
        return combined_filename, status
    
    except Exception as e:
        print(f"❌ Combine error: {str(e)}")
        return None, f"Combine error: {str(e)}"

def transform_with_melodyflow_api(audio_path, prompt, solver="euler", flowstep=0.12):
    """Transform audio using Facebook/MelodyFlow space API"""
    if audio_path is None:
        return None, "❌ No audio file provided"
    
    try:
        # Initialize client for Facebook MelodyFlow space
        client = Client("facebook/MelodyFlow")
        
        # Set steps based on solver
        if solver == "midpoint":
            base_steps = 128
            effective_steps = base_steps // 2  # 64 effective steps
        else:  # euler
            base_steps = 125  
            effective_steps = base_steps // 5  # 25 effective steps
        
        print(f"🎛️ MelodyFlow transformation:")
        print(f"   Prompt: {prompt}")
        print(f"   Solver: {solver} ({effective_steps} effective steps)")
        print(f"   Flowstep: {flowstep}")
        
        # Call the MelodyFlow API
        result = client.predict(
            model="facebook/melodyflow-t24-30secs",
            text=prompt,
            solver=solver,
            steps=base_steps,
            target_flowstep=flowstep,
            regularize=solver == "euler",
            regularization_strength=0.2,
            duration=30,
            melody=handle_file(audio_path),
            api_name="/predict"
        )
        
        if result and len(result) > 0 and result[0]:
            # Save the result locally
            output_filename = f"melodyflow_transformed_{random.randint(1000, 9999)}.wav"
            import shutil
            shutil.copy2(result[0], output_filename)
            
            status_msg = f"✅ Transformed with prompt: '{prompt}' (flowstep: {flowstep}, {effective_steps} steps)"
            return output_filename, status_msg
        else:
            return None, "❌ MelodyFlow API returned no results"
            
    except Exception as e:
        return None, f"❌ MelodyFlow API error: {str(e)}"

def calculate_optimal_bars(bpm):
    """Calculate optimal bar count for given BPM to fit in ~10s"""
    seconds_per_beat = 60.0 / bpm
    seconds_per_bar = seconds_per_beat * 4
    max_duration = 10.0
    
    for bars in [8, 4, 2, 1]:
        if seconds_per_bar * bars <= max_duration:
            return bars
    return 1

# ========== GRADIO INTERFACE ==========

with gr.Blocks(title="stable-melodyflow") as iface:
    gr.Markdown("# stable-melodyflow (aka jerry and terry)")
    gr.Markdown("**generate synchronized drum and instrument loops with stable-audio-open-small (jerry), then transform with melodyflow (terry)!**")
    
    # ========== MODELS & PROJECT INFO ==========
    with gr.Accordion(" some info about these models", open=False):
        
        with gr.Accordion("🚀 stable-audio-open-small", open=False):
            gr.Markdown("""
            **stable-audio-open-small** is an incredibly fast model from the zachs and friends at Stability AI. It's capable of generating 12 seconds of audio in under a second, which gives rise to a lot of very interesting kinds of UX.
            
            **note about generation speed in this zerogpu space:** you'll notice generation times are a little slower here than if you were to use the model on a local gpu. that's just a result of the way zerogpu spaces work i think... let me know if there's a way to keep the model loaded in a zerogpu space!
            
            **links:**
            - 🤗 [model on HuggingFace](https://huggingface.co/stabilityai/stable-audio-open-small)
                        there's a docker container at this repo that can be spun up as a standalone api specifically for stable-audio-open-small:
            -  [stable-audio-api](https://github.com/betweentwomidnights/stable-audio-api)
            """)
        
        with gr.Accordion("🎛️ melodyflow", open=False):
            gr.Markdown("""
            **MelodyFlow** is a model by meta that can use regularized latent inversion to do transformations of input audio.
            
            It's not officially a part of the audiocraft repo yet, but we use it as a docker container in the backend for gary4live. i really enjoy turning my guitar riffs into orchestra.
            
            **links:**
            - 🤗 [Official MelodyFlow Space](https://huggingface.co/spaces/Facebook/MelodyFlow)  
                        
            -  [our melodyflow api](https://github.com/betweentwomidnights/melodyflow)
            """)
        
        with gr.Accordion("gary4live Project", open=False):
            gr.Markdown("""
            **gary4live** is a free/open source project that uses these models, along with musicGen, inside of ableton live to iterate on your projects with you. i run a backend myself so that we can all experiment with it, but you can also spin the backend up locally using docker-compose with our repo.
            
            **project Links:**
            -  [frontend repo](https://github.com/betweentwomidnights/gary4live)
            -  [backend repo](https://github.com/betweentwomidnights/gary-backend-combined)
            
            **installers:**
            -  [p.c. & mac installers on gumroad](https://thepatch.gumroad.com/l/gary4live)
            """)
    
    with gr.Accordion("how this works", open=False):
        gr.Markdown("""
        **workflow:**
        1. **set global bpm and bars** - affects both drum and instrument generation
        2. **generate drum loop** - creates BPM-aware percussion with negative prompting to attempt to get rid of instruments
        3. **generate instrument loop** - creates melodic/harmonic content with negative prompting to attempt to get rid of drums 
        4. **combine loops** - layer them together with repetitions (up to 30s)
        5. **transform** - use melodyflow to stylistically transform the combined result
        
        **features:**
        - bpm-aware generation ensures perfect sync between loops (most the time lol)
        - negative prompting separates drums from instruments (most the time)
        - smart bar calculation optimizes loop length for the BPM
        """)
    
    # ========== GLOBAL CONTROLS ==========
    gr.Markdown("## 🎛️ global settings")
    
    with gr.Row():
        global_bpm = gr.Dropdown(
            label="global bpm",
            choices=[90, 100, 110, 120, 130, 140, 150],
            value=120,
            info="bpm applied to both drum and instrument generation. keep this the same for the combine step to work correctly"
        )
        
        global_bars = gr.Dropdown(
            label="loop length (bars)",
            choices=[1, 2, 4],
            value=4,
            info="number of bars for each loop. keep this the same for both pieces of audio"
        )
        
        base_prompt = gr.Textbox(
            label="base prompt",
            value="lofi hiphop with pianos",
            placeholder="e.g., 'aggressive techno', 'lofi hiphop', 'chillwave', 'liquid drum and bass'",
            info="prompt applied to either loop. make it more drum/instrument specific for best results"
        )
    
    # Auto-suggest optimal bars based on BPM
    def update_suggested_bars(bpm):
        optimal = calculate_optimal_bars(bpm)
        return gr.update(info=f"Suggested: {optimal} bars for {bpm}bpm (≤10s)")
    
    global_bpm.change(update_suggested_bars, inputs=[global_bpm], outputs=[global_bars])
    
    # ========== LOOP GENERATION ==========
    gr.Markdown("## step one: generate individual loops")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### drums")
            generate_drums_btn = gr.Button("generate drums", variant="primary", size="lg")
            drums_audio = gr.Audio(label="drum loop", type="filepath")
            drums_status = gr.Textbox(label="status", value="ready to generate")
            
        with gr.Column():
            gr.Markdown("### instruments")
            generate_instruments_btn = gr.Button("generate instruments", variant="secondary", size="lg")
            instruments_audio = gr.Audio(label="instrument loop", type="filepath")
            instruments_status = gr.Textbox(label="status", value="Ready to generate")
    
    # Seed controls
    with gr.Row():
        drums_seed = gr.Number(label="drums seed", value=-1, info="-1 for random")
        instruments_seed = gr.Number(label="instruments seed", value=-1, info="-1 for random")
    
    # ========== COMBINATION ==========
    gr.Markdown("## step two: combine loops")
    
    with gr.Row():
        num_repeats = gr.Slider(
            label="number of repetitions",
            minimum=1,
            maximum=5,
            step=1,
            value=2,
            info="how many times to repeat each loop (creates longer audio). aim for 30 seconds max"
        )
        combine_btn = gr.Button("combine", variant="primary", size="lg")
    
    combined_audio = gr.Audio(label="combined loops", type="filepath")
    combine_status = gr.Textbox(label="status", value="Generate loops first")
    
    # ========== MELODYFLOW TRANSFORMATION ==========
    gr.Markdown("## step three: transform with melodyflow")
    
    with gr.Row():
        with gr.Column():
            transform_prompt = gr.Textbox(
                label="transformation prompt",
                value="aggressive industrial techno with distorted sounds",
                placeholder="describe the style of transformation",
                lines=2
            )
            
        with gr.Column():
            transform_solver = gr.Dropdown(
                label="solver",
                choices=["euler", "midpoint"],
                value="euler",
                info="EULER: faster (25 steps), MIDPOINT: slower (64 steps)"
            )
            transform_flowstep = gr.Slider(
                label="transform intensity",
                minimum=0.0,
                maximum=0.15,
                step=0.01,
                value=0.12,
                info="Lower = more dramatic transformation"
            )
    
    transform_btn = gr.Button("transform audio", variant="secondary", size="lg")
    transformed_audio = gr.Audio(label="transformed audio", type="filepath")
    transform_status = gr.Textbox(label="status", value="Combine audio first")
    
    # ========== EVENT HANDLERS ==========
    
    # Generate drums
    generate_drums_btn.click(
        generate_stable_audio_loop,
        inputs=[base_prompt, gr.State("drums"), global_bpm, global_bars, drums_seed],
        outputs=[drums_audio, drums_status]
    )
    
    # Generate instruments  
    generate_instruments_btn.click(
        generate_stable_audio_loop,
        inputs=[base_prompt, gr.State("instruments"), global_bpm, global_bars, instruments_seed],
        outputs=[instruments_audio, instruments_status]
    )
    
    # Combine loops
    combine_btn.click(
        combine_loops,
        inputs=[drums_audio, instruments_audio, global_bpm, global_bars, num_repeats],
        outputs=[combined_audio, combine_status]
    )
    
    # Transform with MelodyFlow
    transform_btn.click(
        transform_with_melodyflow_api,
        inputs=[combined_audio, transform_prompt, transform_solver, transform_flowstep],
        outputs=[transformed_audio, transform_status]
    )
    
    # # ========== EXAMPLES ==========
    # gr.Markdown("## 🎯 Example Workflows")
    
    # examples = gr.Examples(
    #     examples=[
    #         ["techno", 128, 4, "aggressive industrial techno"],
    #         ["jazz", 110, 2, "smooth lo-fi jazz with vinyl crackle"],
    #         ["ambient", 90, 8, "ethereal ambient soundscape"],
    #         ["hip-hop", 100, 4, "classic boom bap hip-hop"],
    #         ["drum and bass", 140, 4, "liquid drum and bass"],
    #     ],
    #     inputs=[base_prompt, global_bpm, global_bars, transform_prompt],
    # )

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