import os
import time
import numpy as np
import librosa
import soundfile as sf
import pyloudnorm as pyln
import scipy.signal
import gradio as gr

# ---- CONFIG ----
TARGET_LOUDNESS = -24.0
FINAL_LOUDNESS = -14.0

# ---- AUDIO PROCESSING FUNCTIONS ----
def normalize_loudness(audio, sr, target_lufs):
    meter = pyln.Meter(sr)
    return pyln.normalize.loudness(audio, meter.integrated_loudness(audio), target_lufs)

def highpass(audio, sr, cutoff=40):
    sos = scipy.signal.butter(2, cutoff, btype='highpass', fs=sr, output='sos')
    return scipy.signal.sosfilt(sos, audio)

def pan_stereo(mono_audio, pan):
    left = mono_audio * (1 - max(0, pan))
    right = mono_audio * (1 - max(0, -pan))
    return np.vstack((left, right))

def auto_mix(track1, track2, track3, track4, track5, track6, track7, track8,
             vol1, vol2, vol3, vol4, vol5, vol6, vol7, vol8,
             pan1, pan2, pan3, pan4, pan5, pan6, pan7, pan8,
             bass_boost=0.0, brightness=0.0, vocal_boost=0.0):

    tracks = [track1, track2, track3, track4, track5, track6, track7, track8]
    volumes = [vol1, vol2, vol3, vol4, vol5, vol6, vol7, vol8]
    pans = [pan1, pan2, pan3, pan4, pan5, pan6, pan7, pan8]

    tracks_info = [(t, v, p) for t, v, p in zip(tracks, volumes, pans) if t is not None]
    if not tracks_info:
        return None

    stems = []
    sr = None

    for t, vol, pan in tracks_info:
        fpath = t.name
        audio, sr = librosa.load(fpath, sr=None, mono=True)
        audio = normalize_loudness(audio, sr, TARGET_LOUDNESS)
        audio = highpass(audio, sr)

        if bass_boost > 0:
            sos = scipy.signal.butter(2, 150, btype='low', fs=sr, output='sos')
            audio += bass_boost * scipy.signal.sosfilt(sos, audio)
        if brightness > 0:
            sos = scipy.signal.butter(2, 4000, btype='high', fs=sr, output='sos')
            audio += brightness * scipy.signal.sosfilt(sos, audio)
        if vocal_boost > 0 and "voc" in os.path.basename(fpath).lower():
            audio *= (1.0 + vocal_boost)

        audio *= vol
        stems.append(pan_stereo(audio, pan))

    max_len = max(stem.shape[1] for stem in stems)
    mix = np.zeros((2, max_len))
    for stem in stems:
        padded = np.zeros((2, max_len))
        padded[:, :stem.shape[1]] = stem
        mix += padded

    meter = pyln.Meter(sr)
    mix = pyln.normalize.loudness(mix.T, meter.integrated_loudness(mix.T), FINAL_LOUDNESS).T

    timestamp = time.strftime("%Y%m%d_%H%M%S")
    out_file = f"mix_{timestamp}.wav"
    sf.write(out_file, mix.T, sr)
    return out_file

# ---- GRADIO UI ----
with gr.Blocks() as demo:
    gr.Markdown("## 🎚️ Mini Automatic Mixer (Online Version)")

    tracks = []
    vols = []
    pans = []

    # ---- First row: Track 1-4 ----
    with gr.Row():
        for i in range(4):
            with gr.Column():
                track = gr.File(label=f"Track {i+1}", file_types=[".wav", ".aiff", ".mp3"])
                vol = gr.Slider(0, 2.0, step=0.05, label="Volume", value=1.0)
                pan = gr.Slider(-1.0, 1.0, step=0.05, label="Pan", value=0.0)
                tracks.append(track)
                vols.append(vol)
                pans.append(pan)

    # ---- Second row: Track 5-8 ----
    with gr.Row():
        for i in range(4, 8):
            with gr.Column():
                track = gr.File(label=f"Track {i+1}", file_types=[".wav", ".aiff", ".mp3"])
                vol = gr.Slider(0, 2.0, step=0.05, label="Volume", value=1.0)
                pan = gr.Slider(-1.0, 1.0, step=0.05, label="Pan", value=0.0)
                tracks.append(track)
                vols.append(vol)
                pans.append(pan)

    # ---- General effects row ----
    with gr.Row():
        bass = gr.Slider(0, 1.0, step=0.1, label="Bass Boost")
        bright = gr.Slider(0, 1.0, step=0.1, label="Brightness")
        vocal = gr.Slider(0, 1.0, step=0.1, label="Vocal Boost (files with 'voc' in name)")

    # ---- Mix / Preview / Download ----
    mix_btn = gr.Button("🎛️ Auto Mix & Master")
    audio_out = gr.Audio(label="Preview Mix", type="filepath")
    download_btn = gr.File(label="Download Mix")

    mix_btn.click(
        auto_mix,
        inputs=tracks + vols + pans + [bass, bright, vocal],
        outputs=[audio_out]
    )

    audio_out.change(lambda f: f, inputs=audio_out, outputs=download_btn)

demo.launch(share=True)