from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

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


from bark import generate_audio, preload_models, SAMPLE_RATE
from scipy.io.wavfile import write as write_wav
import tempfile


model_card = ModelCard(
    name="Bark",
    description=(
        "Bark is a transformer-based text-to-audio model created by Suno. Bark "
        "can generate highly realistic, multilingual speech as well as other "
        "audio - including music, background noise and simple sound effects. "
        "The model can also produce nonverbal communications like laughing, "
        "sighing and crying. It writes 24 kHz mono audio. Repository: "
        "https://github.com/suno-ai/bark"
    ),
    author="suno-ai",
    tags=[Subcategory.TEXT_TO_SPEECH, SampleRate(SAMPLE_RATE), Channels(1), 'multilingual', 'domain:speech'],
)

# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so every model is downloaded and loaded there alone
if __name__ != "__main__":
    preload_models()


def process_fn(text_prompt, history_prompt, text_temp, waveform_temp):
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not (text_prompt or "").strip():
        raise gr.Error("Enter a text prompt.")
    return generate(text_prompt, history_prompt, text_temp, waveform_temp)


@spaces.GPU
def generate(text_prompt, history_prompt, text_temp, waveform_temp):
    history_prompt_val = history_prompt if history_prompt != "None" else None

    # Generate audio from text
    audio_array = generate_audio(
        text_prompt,
        history_prompt=history_prompt_val,
        text_temp=text_temp,
        waveform_temp=waveform_temp,
        silent=True, # Disable progress bar for Gradio app
        output_full=False # Only return the audio array
    )

    # Save audio to a temporary WAV file
    output_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    output_file.close()
    write_wav(output_file.name, SAMPLE_RATE, audio_array)

    return output_file.name


if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Textbox(label="Text Prompt", info="The text to be converted into audio. Bark supports multilingual speech, music, and nonverbal communications.").harp_required(True),
            gr.Dropdown(choices=["None", "announcer", "en_speaker_0", "en_speaker_1", "en_speaker_2", "en_speaker_3", "en_speaker_4", "en_speaker_5", "en_speaker_6", "en_speaker_7", "en_speaker_8", "en_speaker_9"], value="en_speaker_0", label="Voice Prompt (Speaker)", info="Choose a speaker to use as a voice prompt for audio cloning. 'None' will use a default voice."),
            gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.7, label="Text Generation Temperature", info="Controls the diversity of text token generation (1.0 more diverse, 0.0 more conservative)."),
            gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.7, label="Waveform Generation Temperature", info="Controls the diversity of waveform token generation (1.0 more diverse, 0.0 more conservative)."),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Generated Audio").set_info((
                "The spoken or sung result, at 24 kHz."
            )),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )

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