from audiotools import AudioSignal
from pyharp.core import ModelCard, build_endpoint
from pyharp.media.audio import load_audio, save_audio
from pyharp.tags import Subcategory

import gradio as gr
import librosa
import soundfile as sf
import torch

import tempfile
from pathlib import Path

# ModelCard

model_card = ModelCard(
    name="Harmonic-Percussive Source Separation",
    description=(
        "Splits a track into its harmonic and percussive components and remixes them at an "
        "adjustable balance, using librosa's harmonic-percussive source separation. Each channel "
        "is separated on its own, at the rate of the input, and a remix that would exceed full "
        "scale is turned down to fit rather than clipped. Repository: "
        "https://github.com/librosa/librosa"
    ),
    author="Brian McFee, librosa developers",
    tags=[Subcategory.STEM_SEPARATION, 'example', 'domain:music']
)

# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the result
# lossily again would expose the artifacts of the first encode and add a second set.
LOSSLESS = {".wav", ".flac"}


# Model Logic
def hpss(signal: AudioSignal, **kwargs):
    h, p = librosa.effects.hpss(signal.audio_data.squeeze().numpy(), **kwargs)

    if h.ndim == 1:
        h = h[None, None, :]
        p = p[None, None, :]
    elif h.ndim == 2:
        h = h[None, :, :]
        p = p[None, :, :]
    else:
        raise ValueError("Unexpected audio shape")

    harmonic_signal = signal.clone()
    harmonic_signal.audio_data = torch.from_numpy(h)

    percussive_signal = signal.clone()
    percussive_signal.audio_data = torch.from_numpy(p)

    return harmonic_signal, percussive_signal

MIN_DB = -120

def process_fn(audio_file_path,
               harmonic_db: float,
               percussive_db: float,
               kernel_size: int = 31,
               margin: float = 1.0):
    if not audio_file_path:
        raise gr.Error("Provide the audio to split.")
    try:
        info = sf.info(audio_file_path)
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if info.frames == 0:
        raise gr.Error("The audio file is empty.")

    sig = load_audio(audio_file_path)

    harmonic, percussive = hpss(sig, kernel_size=int(kernel_size), margin=margin)

    def clip(db):
        return -float("inf") if db == MIN_DB else db

    # Mix signals with gain
    output_sig = (
        harmonic.volume_change(clip(harmonic_db)) +
        percussive.volume_change(clip(percussive_db))
    )

    # Boosting either part can carry the remix past full scale, which the file writer would clip,
    # so the whole remix is turned down instead, keeping the requested balance
    peak = output_sig.audio_data.abs().max().item()
    if peak > 1.0:
        output_sig.audio_data = output_sig.audio_data / peak

    input_ext = Path(audio_file_path).suffix.lower()
    output_ext = input_ext if input_ext in LOSSLESS else ".wav"
    output_audio_path = str(Path(tempfile.mkdtemp()) / f"{Path(audio_file_path).stem}_hpss{output_ext}")
    save_audio(output_sig, output_audio_path)
    return output_audio_path


# Launch the App
if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_audio = gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
            "The audio to split. Harmonic content is the sustained, pitched "
            "part and percussive content is the transient part."
        ))

        harmonic_slider = gr.Slider(
            minimum=MIN_DB, maximum=24,
            step=1, value=0,
            label="Harmonic Level (dB)",
            info=(
                "Boost or reduce tonal parts like chords and pads. A boost that would push the "
                "remix past full scale turns the whole remix down instead."
            )
        )

        percussive_slider = gr.Slider(
            minimum=MIN_DB, maximum=24,
            step=1, value=0,
            label="Percussive Level (dB)",
            info=(
                "Control the loudness of drums and other rhythmic hits. A boost that would push "
                "the remix past full scale turns the whole remix down instead."
            )
        )

        kernel_slider = gr.Slider(
            minimum=1, maximum=101,
            step=1, value=31,
            label="Time Resolution",
            info="How much time/spectral context is used to classify energy as harmonic or percussive. \n Lower = more detailed, reacts faster to rhythm, higher = smoother and slower"
        )

        margin_slider = gr.Slider(
            minimum=0.5, maximum=5.0,
            step=0.1, value=1.0,
            label="Separation Strength",
            info="How strongly harmonic & percussive components must differ to be separated. \n Higher = cleaner isolation, lower = more natural blend."
        )

        output_audio = gr.Audio(type="filepath", label="Remixed Audio").set_info((
            "The harmonic and percussive parts, remixed by the balance control."
        ))

        endpoint = build_endpoint(
            model_card=model_card,
            input_components=[
                input_audio,
                harmonic_slider,
                percussive_slider,
                kernel_slider,
                margin_slider
            ],
            output_components=[output_audio],
            process_fn=process_fn
        )

    demo.queue()
    demo.launch(pwa=True, show_error=True, ssr_mode=False)