from __future__ import annotations

import os

import gradio as gr

from pyharp import *
from gradio_client import Client, handle_file


_BACKEND_SPACE = "Soul-AILab/SoulX-Singer"
_BACKEND_API_NAME = "/synthesis_function"
_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="SoulX-Singer",
    description="SoulX-Singer is a high-fidelity, zero-shot singing voice synthesis model that enables users to generate realistic singing voices for unseen singers. It supports melody-conditioned (F0 contour) and score-conditioned (MIDI notes) control for precise pitch, rhythm, and expression.",
    author="Soul-AILab",
    tags=["text-to-audio", "music", "singing-voice-synthesis", "svs", "zero-shot", "text-to-speech", "en", "zh"],
)


def process_fn(prompt_audio, target_audio, control, auto_shift, pitch_shift, seed, prompt_lyric_lang, target_lyric_lang, prompt_vocal_sep, target_vocal_sep):
    _raw = _backend_client().predict(
        handle_file(prompt_audio),
        handle_file(target_audio),
        None,
        None,
        control,
        auto_shift,
        pitch_shift,
        seed,
        prompt_lyric_lang,
        target_lyric_lang,
        prompt_vocal_sep,
        target_vocal_sep,
        api_name="/synthesis_function",
    )
    _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_generated_audio = _values[0] if len(_values) > 0 else None
    if not _out_generated_audio:
        raise gr.Error(_detail or "The backend Space returned no 'generated_audio' output. Check the backend Space's logs; if it uses ZeroGPU it may need a moment to warm up.")
    _out_processed_prompt_metadata = _values[1] if len(_values) > 1 else None
    if not _out_processed_prompt_metadata:
        raise gr.Error(_detail or "The backend Space returned no 'processed_prompt_metadata' output. Check the backend Space's logs; if it uses ZeroGPU it may need a moment to warm up.")
    _out_processed_target_metadata = _values[2] if len(_values) > 2 else None
    if not _out_processed_target_metadata:
        raise gr.Error(_detail or "The backend Space returned no 'processed_target_metadata' output. Check the backend Space's logs; if it uses ZeroGPU it may need a moment to warm up.")
    return _out_generated_audio, _out_processed_prompt_metadata, _out_processed_target_metadata


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Prompt audio (reference voice), max 30s").set_info("Upload an audio file (max 30 seconds) to provide the reference voice for synthesis."),
        gr.Audio(type="filepath", label="Target audio (melody / lyrics source), max 60s").set_info("Upload an audio file (max 60 seconds) to provide the melody or lyrics source."),
        gr.Dropdown(choices=["melody", "score"], value="melody", label="Control type", info="Choose the control type for synthesis: 'melody' for F0 contour or 'score' for MIDI notes."),
        gr.Checkbox(value=True, label="Auto pitch shift", info="Automatically adjust pitch shift to match the target audio's range."),
        gr.Slider(minimum=-12, maximum=12, step=1, value=0, label="Pitch shift (semitones)", info="Manually adjust the pitch shift in semitones. Auto pitch shift will be ignored if a non-zero value is set."),
        gr.Number(value=12306, label="Seed", info="Random seed for reproducibility."),
        gr.Dropdown(choices=["English", "Chinese"], value="English", label="Prompt lyric language", info="Select the language of the lyrics in the prompt audio."),
        gr.Dropdown(choices=["English", "Chinese"], value="English", label="Target lyric language", info="Select the language of the lyrics in the target audio."),
        gr.Checkbox(value=False, label="Prompt vocal separation", info="Enable vocal separation for the prompt audio if it contains accompaniment."),
        gr.Checkbox(value=True, label="Target vocal separation", info="Enable vocal separation for the target audio if it contains accompaniment."),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Generated audio").set_info("The synthesized singing voice."),
        gr.File(type="filepath", label="Processed Prompt Metadata", file_types=[".mid", ".midi"]).set_info("Metadata file generated from the prompt audio."),
        gr.File(type="filepath", label="Processed Target Metadata", file_types=[".mid", ".midi"]).set_info("Metadata file generated from the target 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)