import io
import json
from pydub import AudioSegment
from pydub.effects import normalize, compress_dynamic_range, low_pass_filter, high_pass_filter
import numpy as np
import gradio as gr

# Default presets for common genres
PRESETS = {
    "rock": {"gain": 6, "compress_threshold": -20, "compress_ratio": 4, "low_pass": 12000, "high_pass": 80, "stereo_width": 100, "pan": 0},
    "pop": {"gain": 4, "compress_threshold": -18, "compress_ratio": 3, "low_pass": 15000, "high_pass": 100, "stereo_width": 100, "pan": 0},
    "jazz": {"gain": 3, "compress_threshold": -22, "compress_ratio": 2, "low_pass": 14000, "high_pass": 60, "stereo_width": 90, "pan": 0},
    "electronic": {"gain": 8, "compress_threshold": -16, "compress_ratio": 5, "low_pass": 20000, "high_pass": 120, "stereo_width": 120, "pan": 0},
}

# JSON export extension
EXPORT_EXT = ".master"

# Custom CSS for Gradio
CUSTOM_CSS = """
body { background-color: #1e1e1e; color: #f0f0f0; }
.gradio-container { border-radius: 12px; box-shadow: 0 4px 12px rgba(0,0,0,0.5); }
#title { font-size: 2rem; margin-bottom: 1rem; }
"""

def apply_mastering(audio: AudioSegment, params: dict) -> AudioSegment:
    # Gain
    processed = audio + params["gain"]
    # Compression
    processed = compress_dynamic_range(
        processed,
        threshold=params["compress_threshold"],
        ratio=params["compress_ratio"]
    )
    # EQ filters
    processed = high_pass_filter(processed, params["high_pass"])
    processed = low_pass_filter(processed, params["low_pass"])
    # Stereo Imaging (Width)
    processed = apply_stereo_width(processed, params.get("stereo_width", 100))
    # Neural Pan
    processed = apply_pan(processed, params.get("pan", 0))
    # Normalization
    processed = normalize(processed)
    return processed


def apply_stereo_width(audio: AudioSegment, width: float) -> AudioSegment:
    # Width in percent: 0 = mono, 100 = original, >100 = enhanced
    left, right = audio.split_to_mono()
    mid = (left + right) / 2
    side = (left - right) / 2
    # scale side for width control
    side = side * (width / 100)
    new_left = mid + side
    new_right = mid - side
    return AudioSegment.from_mono_audiosegments(new_left, new_right)


def apply_pan(audio: AudioSegment, pan: float) -> AudioSegment:
    # Pan from -100 (left) to 100 (right)
    left, right = audio.split_to_mono()
    pan = max(-100, min(100, pan)) / 100
    if pan > 0:
        new_left = left * (1 - pan)
        new_right = right
    else:
        new_left = left
        new_right = right * (1 + pan)
    return AudioSegment.from_mono_audiosegments(new_left, new_right)


def process_and_export(audio_file, preset_name, gain, threshold, ratio, lp, hp, width, pan):
    audio = AudioSegment.from_file(audio_file)
    if preset_name != "custom":
        params = PRESETS[preset_name]
    else:
        params = {"gain": gain, "compress_threshold": threshold, "compress_ratio": ratio,
                  "low_pass": lp, "high_pass": hp, "stereo_width": width, "pan": pan}
    out = apply_mastering(audio, params)
    buf = io.BytesIO()
    out.export(buf, format="wav")
    buf.seek(0)
    preset = {"preset": preset_name, "params": params}
    preset_buf = io.BytesIO(json.dumps(preset, indent=2).encode())
    preset_buf.name = "preset" + EXPORT_EXT
    return (buf, "processed.wav"), (preset_buf, preset_buf.name)


def load_preset_json(file_obj):
    data = json.load(file_obj)
    name = data.get("preset", "custom")
    params = data.get("params", {})
    return name, params.get("gain", 0), params.get("compress_threshold", 0), \
           params.get("compress_ratio", 1), params.get("low_pass", 20000), \
           params.get("high_pass", 20), params.get("stereo_width", 100), params.get("pan", 0)

with gr.Blocks(css=CUSTOM_CSS) as demo:
    gr.Markdown("<div id='title'>Advanced Browser-based Audio Mastering Suite</div>")
    with gr.Row():
        with gr.Column():
            # Use 'filepath' to get the file path for AudioSegment
            audio_input = gr.Audio(label="Upload Audio", type="filepath")
            preset_dropdown = gr.Dropdown(choices=list(PRESETS.keys()) + ["custom"], value="rock", label="Preset")
            gain = gr.Slider(-10, 20, value=PRESETS["rock"]["gain"], label="Gain (dB)")
            threshold = gr.Slider(-60, 0, value=PRESETS["rock"]["compress_threshold"], label="Compress Threshold (dB)")
            ratio = gr.Slider(1, 10, value=PRESETS["rock"]["compress_ratio"], label="Compress Ratio")
            lp = gr.Slider(1000, 20000, value=PRESETS["rock"]["low_pass"], label="Low-pass Frequency (Hz)")
            hp = gr.Slider(20, 500, value=PRESETS["rock"]["high_pass"], label="High-pass Frequency (Hz)")
            width = gr.Slider(0, 200, value=PRESETS["rock"]["stereo_width"], label="Stereo Width (%)")
            pan = gr.Slider(-100, 100, value=PRESETS["rock"]["pan"], label="Pan (-100 Left to 100 Right)")
            load_preset = gr.File(label="Load .master Preset", file_types=[".master"])
            export_button = gr.Button("Process & Export")
        with gr.Column():
            # Set output to filepath to serve downloadable file
            output_audio = gr.Audio(label="Processed Audio", type="filepath")
            export_preset_file = gr.File(label="Download Preset (.master)")
    def sync_sliders(preset):
        if preset != "custom":
            p = PRESETS[preset]
            return (gr.update(value=p["gain"]), gr.update(value=p["compress_threshold"]), \
                    gr.update(value=p["compress_ratio"]), gr.update(value=p["low_pass"]), \
                    gr.update(value=p["high_pass"]), gr.update(value=p["stereo_width"]), \
                    gr.update(value=p["pan"]))
        return None
    preset_dropdown.change(sync_sliders, inputs=[preset_dropdown], outputs=[gain, threshold, ratio, lp, hp, width, pan])
    load_preset.upload(load_preset_json, inputs=[load_preset], outputs=[preset_dropdown, gain, threshold, ratio, lp, hp, width, pan])
    export_button.click(process_and_export, inputs=[audio_input, preset_dropdown, gain, threshold, ratio, lp, hp, width, pan], outputs=[output_audio, export_preset_file])

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