"""
SonicMaster on HARP: text-controlled music restoration and mastering.

The model runs here rather than on the authors' Space, so a request draws on
this Space's own allowance instead of theirs. The weights come from
amaai-lab/SonicMaster and the inference code is vendored from the authors'
repository, which LICENSE-SonicMaster carries.
"""

import tempfile
from pathlib import Path

import gradio as gr
import spaces
from pyharp import (Channels, ModelCard, SampleRate, Subcategory,
                    build_endpoint)

import sonicmaster_runtime

model_card = ModelCard(
    name="SonicMaster",
    description=(
        "Text-controlled all-in-one music restoration and mastering: one "
        "instruction in plain language fixes reverb, clipping, tonal balance, "
        "dynamics, or the stereo image. The model is a flow-matching "
        "transformer over Stable Audio's latent space, conditioned on a "
        "FLAN-T5 encoding of the instruction, so the wording of the prompt is "
        "what selects the treatment. Audio is resampled to 44.1 kHz stereo and "
        "processed in overlapping 30-second chunks, each conditioned on the "
        "tail of the one before it, which keeps a long track consistent across "
        "chunk boundaries. SonicMaster itself is Apache 2.0, but it decodes "
        "through Stable Audio Open's VAE, so results are also governed by the "
        "Stability AI Community License: "
        "https://huggingface.co/stabilityai/stable-audio-open-1.0/blob/main/LICENSE.md "
        "which permits non-commercial use and commercial use below an annual "
        "revenue threshold. Repository: "
        "https://github.com/AMAAI-Lab/SonicMaster Paper: "
        "https://arxiv.org/abs/2508.03448"
    ),
    author="Jan Melechovsky, Ambuj Mehrish, Dorien Herremans",
    tags=[Subcategory.POST_PROCESSING, SampleRate(sonicmaster_runtime.SAMPLE_RATE),
          Channels(sonicmaster_runtime.CHANNELS), "text-guided"],
)


@spaces.GPU(duration=240)
def process_fn(input_audio_path, prompt, num_inference_steps, guidance_scale,
               solver, seed):
    """
    Apply the instruction to the recording.

    Args:
        input_audio_path (str): The recording to edit.
        prompt (str): What to change about it.
        num_inference_steps (float): Flow solver steps per chunk.
        guidance_scale (float): How strongly the instruction is followed.
        solver (str): Which solver steps the flow.
        seed (float): Seed for the noise.

    Returns:
        output_path (str): The edited audio.
    """

    if not (prompt or "").strip():
        raise gr.Error("Describe what to change about the audio.")

    output_path = Path(tempfile.mkdtemp()) / "sonicmaster.wav"

    return sonicmaster_runtime.enhance(
        input_audio_path,
        prompt.strip(),
        output_path,
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale),
        seed=int(seed),
        solver=solver,
    )


if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input audio").harp_required(True).set_info((
                "The recording to treat. It is resampled to 44.1 kHz and "
                "forced to stereo, since that is what the model was trained "
                "on, and a mono input is duplicated across both channels."
            )),
            gr.Textbox(label="Instruction", value="Remove the reverb from this recording.",
                       info=(
                           "What to change, in plain language, for example "
                           "\"reduce the reverb\", \"brighten the high end\", "
                           "\"even out the loudness\", or \"widen the stereo "
                           "image\". One instruction at a time works better "
                           "than several combined."
                       )),
            gr.Slider(minimum=1, maximum=50, step=1, value=10,
                      label="Solver steps",
                      info=(
                          "How many steps the flow solver takes per chunk. "
                          "More steps follow the instruction more closely and "
                          "cost proportionally more time."
                      )),
            gr.Slider(minimum=1.0, maximum=7.0, step=0.1, value=1.0,
                      label="Instruction strength",
                      info=(
                          "How strongly the instruction is weighted against "
                          "leaving the audio alone. Raising it applies a more "
                          "pronounced treatment, at the risk of artifacts."
                      )),
            gr.Dropdown(choices=["Euler", "rk4"], value="Euler", label="Solver",
                        info=(
                            "Euler takes one evaluation per step. rk4 takes "
                            "four, which is steadier at a low step count and "
                            "roughly four times slower."
                        )),
            gr.Number(value=0, precision=0, label="Seed",
                      info=(
                          "Seeds the noise the solver starts from, so the same "
                          "seed and settings reproduce a result."
                      )),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Treated audio").set_info((
                "The result, at 44.1 kHz stereo."
            )),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )
    demo.queue().launch(server_name="0.0.0.0", server_port=7860,
                        show_error=True, pwa=True)