from __future__ import annotations

import gradio as gr

from pyharp import *


try:  # torch>=2.6 flipped torch.load(weights_only) to True; legacy ckpts need False
    import torch as _torch

    if getattr(_torch.load, "__harp_compat__", False) is False:
        _torch_load_orig = _torch.load

        def _torch_load_compat(*args, **kwargs):
            kwargs.setdefault("weights_only", False)
            return _torch_load_orig(*args, **kwargs)

        _torch_load_compat.__harp_compat__ = True
        _torch.load = _torch_load_compat
except Exception:  # torch not installed / unexpected API -- nothing to patch
    pass


import tempfile
import matchering as mg


model_card = ModelCard(
    name="Matchering",
    description="Match the RMS, Frequency Response, Peak Amplitude, and Stereo Width of a target track to a reference track for instant mastering.",
    author="sergree",
    tags=["audio", "mastering", "dsp", "matching"],
)


def process_fn(target_audio, ref_audio, bit_depth):
    out_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    out_file.close()
    out_path = out_file.name

    if bit_depth == "pcm24":
        result_config = mg.pcm24(out_path)
    else:
        result_config = mg.pcm16(out_path)

    mg.process(
        target=target_audio,
        reference=ref_audio,
        results=[result_config]
    )

    return out_path


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Target Audio").harp_required(True).set_info("The track you want to master (your mix)."),
        gr.Audio(type="filepath", label="Reference Audio").harp_required(True).set_info("The reference track you want your target to sound like."),
        gr.Dropdown(choices=["pcm16", "pcm24"], value="pcm16", label="Output Bit Depth", info="Choose 16-bit PCM (CD quality) or 24-bit PCM (Studio quality) for the output file."),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Mastered Audio"),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

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