from __future__ import annotations

import os

import gradio as gr

from pyharp import *
from gradio_client import Client, handle_file


_BACKEND_SPACE = "amaai-lab/SonicMaster"
_BACKEND_API_NAME = "/enhance_audio_ui"
_BACKEND_TOKEN_ENV = "HF_TOKEN"
_client = None


def _backend_client():
    # Lazily create and cache one warm connection to the backend Space.
    global _client
    if _client is None:
        _token = os.environ.get(_BACKEND_TOKEN_ENV) or None
        _client = Client(_BACKEND_SPACE, hf_token=_token)
    return _client


model_card = ModelCard(
    name="SonicMaster",
    description="Text-controlled all-in-one music restoration and mastering (fix reverb, clipping, EQ/tonal balance, dynamics, and stereo image) from a natural-language instruction. This is a thin HARP frontend that proxies to the authors' official SonicMaster Gradio Space via gradio_client, so none of its heavy/gated dependencies (torch, diffusers, the gated stable-audio VAE) are installed here.",
    author="AMAAI-Lab",
    tags=["audio-to-audio", "music-restoration", "mastering", "text-guided"],
)


def process_fn(input_audio, prompt):
    _raw = _backend_client().predict(
        handle_file(input_audio),
        prompt,
        api_name="/enhance_audio_ui",
    )
    _values = list(_raw) if isinstance(_raw, (list, tuple)) else [_raw]
    _detail = " | ".join(str(_v) for _v in _values if isinstance(_v, str) and _v.strip())
    _out_output_audio = _values[0] if len(_values) > 0 else None
    if not _out_output_audio:
        raise gr.Error(_detail or "The backend Space returned no 'output_audio' output. Check the backend Space's logs; if it uses ZeroGPU it may need a moment to warm up.")
    return _out_output_audio


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input audio").harp_required(True).set_info("The music/audio to restore or master."),
        gr.Textbox(label="Instruction", value="Enhance the input audio", info="Natural-language edit, e.g. 'reduce reverb and brighten the vocals', 'make it louder', 'dereverb', 'widen the stereo image'. Leave as-is for general restoration."),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Enhanced audio").set_info("SonicMaster output returned by the backend Space."),
    ]
    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)