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="Music recordings often suffer from audio quality issues such as excessive reverberation, distortion, clipping, tonal imbalances, and a narrowed stereo image, especially when created in non-professional settings without specialized equipment or expertise. These problems are typically corrected using separate specialized tools and manual adjustments. In this paper, we introduce SonicMaster, the firs",
    author="amaai-lab",
    tags=[],
)


def process_fn(audio_path, prompt):
    _raw = _backend_client().predict(
        handle_file(audio_path),
        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_enhanced_audio_output = _values[0] if len(_values) > 0 else None
    if not _out_enhanced_audio_output:
        raise gr.Error(_detail or "The backend Space returned no 'enhanced_audio_output' output. Check the backend Space's logs; if it uses ZeroGPU it may need a moment to warm up.")
    return _out_enhanced_audio_output


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio"),
        gr.Textbox(label="Text Prompt"),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Enhanced Audio (output)"),
    ]
    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)