import os
import tempfile
from pathlib import Path

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


MODEL_ID = "small-sfx"
SAMPLE_RATE = 44_100

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

if not hf_token:
    raise RuntimeError(
        "HF_TOKEN is missing. Add it under Settings → Variables and secrets."
    )

login(
    token=hf_token,
    add_to_git_credential=False,
)

# These diagnostics do not create a CUDA context.
print("PyTorch:", torch.__version__)
print("PyTorch CUDA build:", torch.version.cuda)
print(f"Loading model: {MODEL_ID}")

# ZeroGPU intercepts CUDA placement here and packs the model for transfer
# to the dynamically allocated GPU worker.
model = StableAudioModel.from_pretrained(
    MODEL_ID,
    device="cuda",
)

print("Model loaded.")


def requested_gpu_duration(prompt: str, duration: float) -> int:
    """Return the number of ZeroGPU seconds to reserve."""

    del prompt

    duration = float(duration)

    if duration <= 15:
        return 15

    if duration <= 60:
        return 25

    return 40


def convert_audio(audio: object) -> np.ndarray:
    """Convert model output into a SoundFile-compatible NumPy array."""

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

    array = np.asarray(audio, dtype=np.float32)

    # Remove batch dimension: [batch, channels, samples].
    if array.ndim == 3:
        array = array[0]

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

    if array.ndim not in (1, 2):
        raise RuntimeError(
            f"Unexpected generated audio shape: {array.shape}"
        )

    if array.size == 0:
        raise RuntimeError("The model returned empty audio.")

    if not np.all(np.isfinite(array)):
        raise RuntimeError(
            "The generated audio contains invalid numeric values."
        )

    peak = float(np.max(np.abs(array)))

    if peak > 1.0:
        array = array / peak

    return array


@spaces.GPU(duration=requested_gpu_duration)
@torch.inference_mode()
def generate_sound(prompt: str, duration: float) -> str:
    """Generate a sound effect and return its WAV path."""

    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."
        )

    # CUDA inspection is safe here because this runs in the allocated
    # ZeroGPU worker, not the parent web process.
    print("Runtime Torch:", torch.__version__)
    print("Runtime CUDA:", torch.version.cuda)
    print("Runtime GPU:", torch.cuda.get_device_name(0))
    print(
        "Runtime compute capability:",
        torch.cuda.get_device_capability(0),
    )
    print(
        "Runtime Torch architectures:",
        torch.cuda.get_arch_list(),
    )

    try:
        generated_audio = model.generate(
            prompt=prompt,
            duration=duration,
        )

        audio = convert_audio(generated_audio)

        with tempfile.NamedTemporaryFile(
            prefix="stable-audio-3-",
            suffix=".wav",
            delete=False,
        ) as temporary_file:
            output_path = Path(temporary_file.name)

        sf.write(
            str(output_path),
            audio,
            samplerate=SAMPLE_RATE,
            subtype="FLOAT",
        )

        return str(output_path)

    except gr.Error:
        raise

    except Exception as error:
        print(
            "Generation error:",
            type(error).__name__,
            str(error),
        )

        raise gr.Error(
            f"Generation failed: {type(error).__name__}: {error}"
        ) from error


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_input = gr.Textbox(
        label="Sound description",
        value=(
            "A chugging steam train entering a station "
            "and sounding a loud horn"
        ),
        placeholder=(
            "Describe the sound, environment, timing, and perspective."
        ),
        lines=3,
    )

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

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

    audio_output = gr.Audio(
        label="Generated sound",
        type="filepath",
        format="wav",
    )

    generate_button.click(
        fn=generate_sound,
        inputs=[
            prompt_input,
            duration_input,
        ],
        outputs=audio_output,
        concurrency_limit=1,
        show_progress="full",
    )


if __name__ == "__main__":
    demo.queue(
        max_size=20,
        default_concurrency_limit=1,
    ).launch()