import os

import gradio as gr
import numpy as np
import soundfile as sf
import spaces
import torch
from huggingface_hub import login
from stable_audio_3 import StableAudioModel


hf_token = os.environ.get("HF_TOKEN")

if not hf_token:
    raise RuntimeError("HF_TOKEN Space secret is missing.")

login(token=hf_token)

print("PyTorch:", torch.__version__)
print("PyTorch CUDA:", torch.version.cuda)
print("Compiled architectures:", torch.cuda.get_arch_list())

model = StableAudioModel.from_pretrained(
    "small-sfx",
    device="cuda",
)

def requested_gpu_duration(prompt: str, duration: float) -> int:
    """
    Reserve a modest amount of GPU time.

    Stable Audio 3 Small models are extremely fast on modern GPUs,
    but leave room for initial setup and audio decoding.
    """
    del prompt

    duration = float(duration)

    if duration <= 15:
        return 15
    if duration <= 60:
        return 20

    return 30


@spaces.GPU(duration=requested_gpu_duration)
@torch.inference_mode()
def generate_sound(prompt: str, duration: float) -> str:
    prompt = prompt.strip()
    duration = float(duration)

    if not prompt:
        raise gr.Error("Enter a description of the sound.")

    if not 0.5 <= duration <= 120:
        raise gr.Error("Duration must be between 0.5 and 120 seconds.")

    audio = model.generate(
        prompt=prompt,
        duration=duration,
    )

    if isinstance(audio, torch.Tensor):
        audio = audio.detach().float().cpu().numpy()

    audio = np.asarray(audio)

    # Remove batch dimension.
    if audio.ndim == 3:
        audio = audio[0]

    # Convert [channels, samples] to [samples, channels].
    if audio.ndim == 2 and audio.shape[0] <= 8:
        audio = audio.T

    with tempfile.NamedTemporaryFile(
        suffix=".wav",
        delete=False,
    ) as output_file:
        output_path = Path(output_file.name)

    sf.write(
        output_path,
        audio,
        samplerate=44_100,
        subtype="FLOAT",
    )

    return str(output_path)


with gr.Blocks(title="Stable Audio 3 Small SFX") as demo:
    gr.Markdown(
        """
        # Stable Audio 3 Small SFX

        Generate stereo, 44.1 kHz sound effects using Stable Audio 3
        on Hugging Face ZeroGPU.
        """
    )

    prompt = gr.Textbox(
        label="Sound description",
        value="Chugging train coming into a station with a loud horn",
        lines=3,
    )

    duration = gr.Slider(
        minimum=0.5,
        maximum=120,
        value=7,
        step=0.5,
        label="Duration in seconds",
    )

    generate_button = gr.Button(
        "Generate sound",
        variant="primary",
    )

    output_audio = gr.Audio(
        label="Generated sound",
        type="filepath",
    )

    generate_button.click(
        fn=generate_sound,
        inputs=[prompt, duration],
        outputs=output_audio,
        concurrency_limit=1,
    )


demo.queue(
    max_size=20,
    default_concurrency_limit=1,
).launch()