import os
import torch
import numpy as np
import soundfile as sf
import gradio as gr
from pydub import AudioSegment
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

STEMS_LIST = ["vocals", "backing_vocals", "drums", "bass", "guitar", "synth"]
cached_model = None

# =====================================================================
# 2. AUDIO LOADER & RESAMPLER (No torchaudio)
# =====================================================================
def load_audio_as_tensor(audio_path, target_sr=16000):
    """
    Loads any audio format directly to a normalized mono PyTorch tensor
    using pydub in memory (no external CLI conversion needed!).
    """
    # Load audio (works on mp3, wav, flac, m4a, aiff, mp4 video, etc.)
    audio = AudioSegment.from_file(audio_path)
    
    # Resample to 16kHz and convert to Mono
    audio = audio.set_frame_rate(target_sr).set_channels(1)
    
    # Get raw samples as numpy array
    samples = np.array(audio.get_array_of_samples(), dtype=np.float32)
    
    # Normalize based on source bit depth
    if audio.sample_width == 1:
        samples = (samples - 128.0) / 128.0
    elif audio.sample_width == 2:
        samples = samples / 32768.0
    elif audio.sample_width == 4:
        samples = samples / 2147483648.0
        
    # Convert directly to PyTorch tensor [1, Samples]
    waveform = torch.from_numpy(samples).unsqueeze(0)
    return waveform


# =====================================================================
# 3. CORE MULTI-STEM INFERENCE
# =====================================================================
@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

    # 1. Load waveform in-memory using our pydub function
    try:
        waveform = load_audio_as_tensor(audio_path, sample_rate)
    except Exception as e:
        print(f"Error decoding audio file: {e}")
        return [None] * len(STEMS_LIST)

    # 2. 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()

    # Move audio tensors to computation device
    waveform = waveform.to(device)
    window = torch.hann_window(512).to(device)

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

    input_tensor = magnitude.unsqueeze(0) # Shape [1, 1, Freq, Time]

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

    # Move tensors back to CPU for audio reconstruction
    pred_stems_spec = pred_stems_spec.squeeze(0) # [num_stems, Freqs, Time]
    phase = phase.squeeze(0)                      # [Freqs, Time]
    window_cpu = torch.hann_window(512).cpu()

    saved_filepaths = []
    
    # 5. Reconstruct and write each isolated stem (No torchaudio)
    for idx, stem_name in enumerate(STEMS_LIST):
        stem_spec = pred_stems_spec[idx].cpu()
        stem_complex = torch.polar(stem_spec, phase.cpu())
        
        # Native PyTorch inverse Fourier transform to reconstruct waveform
        stem_wav = torch.istft(
            stem_complex, n_fft=512, win_length=512, hop_length=128, window=window_cpu, normalized=True
        )
        
        # Output clean WAV stem using soundfile
        out_path = f"output_{stem_name}.wav"
        sf.write(out_path, stem_wav.numpy(), sample_rate)
        saved_filepaths.append(out_path)

    return saved_filepaths


# =====================================================================
# 4. GRADIO INTERFACE LAYOUT
# =====================================================================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        f"""
        # 🎙️ Universal Multi-Stem AI Audio Splitter
        Upload **any** audio file format (.wav, .mp3, .m4a, .flac, .aif, or even .mp4 video) to separate it into **{len(STEMS_LIST)}** individual tracks!
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            input_audio = gr.Audio(type="filepath", label="Upload Audio/Video File")
            submit_btn = gr.Button("Split Stems ⚡", variant="primary")
            
        with gr.Column(scale=2):
            audio_outputs = []
            for stem_name in STEMS_LIST:
                player = gr.Audio(label=f"Isolated {stem_name.replace('_', ' ').title()}", interactive=False)
                audio_outputs.append(player)

    submit_btn.click(
        fn=separate_multi_stems,
        inputs=input_audio,
        outputs=audio_outputs
    )

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