import os
import torch
import torchaudio
import torchaudio.transforms as T
import gradio as gr
from multi_stem_pipeline import MultiStemSeparationModel

# =====================================================================
# 1. ZEROGPU COMPATIBILITY LAYER
# =====================================================================
try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False
    class spaces:
        @staticmethod
        def GPU(fn):
            return fn

# Specify the stems we are extracting
STEMS_LIST = ["vocals", "backing_vocals", "drums", "bass", "guitar", "synth"]
cached_model = None

# =====================================================================
# 2. CORE MULTI-STEM INFERENCE (Top-Level for ZeroGPU Static Analysis)
# =====================================================================
@spaces.GPU
def separate_multi_stems(audio_path):
    global cached_model
    if not audio_path:
        return [None] * len(STEMS_LIST)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    sample_rate = 16000

    # Lazy load our model
    if cached_model is None:
        model_path = "./audio_separation_model"
        if os.path.exists(model_path):
            print(f"Loading local multi-stem model from {model_path}...")
            cached_model = MultiStemSeparationModel.from_pretrained(model_path)
        else:
            print("No trained model found. Launching with untrained model fallback...")
            cached_model = MultiStemSeparationModel(in_channels=1, num_stems=len(STEMS_LIST), base_features=16)

    model = cached_model.to(device)
    model.eval()

    # Load & Resample
    waveform, sr = torchaudio.load(audio_path)
    if sr != sample_rate:
        resampler = T.Resample(orig_freq=sr, new_freq=sample_rate)
        waveform = resampler(waveform)
    
    if waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)

    # Extract Spectrogram Magnitude & Phase
    complex_spec = torch.stft(
        waveform, n_fft=512, win_length=512, hop_length=128, normalized=True, return_complex=True
    )
    magnitude = torch.abs(complex_spec)
    phase = torch.angle(complex_spec)

    input_tensor = magnitude.unsqueeze(0).to(device)

    # Run Multi-Mask Inference
    with torch.no_grad():
        pred_stems_spec, masks = model(input_tensor)

    # Move tensors back to CPU for audio reconstruct
    pred_stems_spec = pred_stems_spec.squeeze(0).cpu() # [num_stems, Freqs, Time]
    phase = phase.cpu()

    saved_filepaths = []
    
    # Process each stem index
    for idx, stem_name in enumerate(STEMS_LIST):
        stem_spec = pred_stems_spec[idx]
        
        # Combine predicted magnitude with original mix phase
        stem_complex = torch.polar(stem_spec, phase)
        stem_wav = torch.istft(
            stem_complex, n_fft=512, win_length=512, hop_length=128, normalized=True
        )
        
        # Save stem to disk
        out_path = f"output_{stem_name}.wav"
        torchaudio.save(out_path, stem_wav, sample_rate)
        saved_filepaths.append(out_path)

    return saved_filepaths


# =====================================================================
# 3. DYNAMIC GRADIO INTERFACE GENERATION
# =====================================================================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        f"""
        # 🎙️ Multi-Stem AI Audio Splitter
        Upload a song to automatically separate it into **{len(STEMS_LIST)}** individual, isolated instrument tracks using advanced spectrogram masking!
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            input_audio = gr.Audio(type="filepath", label="Upload Mixed Song (.wav, .mp3)")
            submit_btn = gr.Button("Split Stems ⚡", variant="primary")
            
        with gr.Column(scale=2):
            # Dynamically create Audio Player components for each target stem
            audio_outputs = []
            for stem_name in STEMS_LIST:
                player = gr.Audio(label=f"Isolated {stem_name.replace('_', ' ').title()}", interactive=False)
                audio_outputs.append(player)

    # Wire button to trigger separation
    submit_btn.click(
        fn=separate_multi_stems,
        inputs=input_audio,
        outputs=audio_outputs
    )

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