from __future__ import annotations

import gradio as gr
import soundfile as sf

try:
    import spaces
except ImportError:
    class spaces:
        class GPU:
            def __init__(self, func=None, duration=180):
                self.func = func

            def __call__(self, *args, **kwargs):
                if self.func is not None:
                    return self.func(*args, **kwargs)
                return args[0]

from pyharp import ModelCard, build_endpoint

from stable_audio_runtime import edit_audio

MIN_AUDIO_SECONDS = 5
MAX_AUDIO_SECONDS = 30

model_card = ModelCard(
    name="Stable Audio 3",
    description=(
        "Restyle, inpaint, or continue music audio using a text prompt."
    ),
    author="Stability AI",
    tags=[
        "audio-generation",
        "music-editing",
        "audio-inpainting",
        "audio-continuation",
    ],
)


def _audio_duration(path: str | None) -> float:
    if not path:
        raise gr.Error("Please upload a music clip.")
    try:
        duration = float(sf.info(path).duration)
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if duration < MIN_AUDIO_SECONDS:
        raise gr.Error(f"Audio must be at least {MIN_AUDIO_SECONDS} seconds.")
    if duration > MAX_AUDIO_SECONDS:
        raise gr.Error(f"Audio must be no longer than {MAX_AUDIO_SECONDS} seconds.")
    return duration


@spaces.GPU(duration=180)
def process_fn(
    audio_path: str | None,
    prompt: str,
    mode: str,
    edit_start: float,
    edit_end: float,
    continuation_length: float,
    strength: float,
    seed: int,
) -> str:
    duration = _audio_duration(audio_path)
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter an editing prompt.")

    if mode == "Inpaint" and (
        edit_start < 0 or edit_end > duration or edit_end <= edit_start
    ):
        raise gr.Error(
            "The inpaint region must be inside the uploaded clip, "
            "with the end after the start."
        )

    try:
        return edit_audio(
            audio_path=audio_path,
            prompt=prompt,
            mode=mode,
            edit_start=float(edit_start),
            edit_end=float(edit_end),
            continuation_length=float(continuation_length),
            strength=float(strength),
            seed=int(seed),
        )
    except gr.Error:
        raise
    except Exception as exc:
        raise gr.Error(f"Stable Audio 3 inference failed: {exc}") from exc


with gr.Blocks(title="Stable Audio 3") as demo:
    input_components = [
        gr.Audio(type="filepath", label="Music Audio")
        .harp_required(True)
        .set_info("Music clip between 5 and 30 seconds."),
        gr.Textbox(
            label="Prompt",
            placeholder="A warm synthwave groove with punchy drums",
        ).harp_required(True),
        gr.Dropdown(
            choices=["Restyle", "Inpaint", "Continue"],
            value="Restyle",
            label="Edit Mode",
        ),
        gr.Slider(
            minimum=0,
            maximum=30,
            value=4,
            step=0.1,
            label="Edit Start (seconds)",
        ).set_info("Used by Inpaint mode."),
        gr.Slider(
            minimum=0,
            maximum=30,
            value=8,
            step=0.1,
            label="Edit End (seconds)",
        ).set_info("Used by Inpaint mode."),
        gr.Slider(
            minimum=1,
            maximum=15,
            value=8,
            step=1,
            label="Continuation Length (seconds)",
        ).set_info("Used by Continue mode."),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.75,
            step=0.05,
            label="Transformation Strength",
        ).set_info("Used by Restyle mode."),
        gr.Slider(
            minimum=0,
            maximum=99999,
            value=0,
            step=1,
            label="Seed",
        ).set_info("Use 0 for a random seed."),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Edited Audio"),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(show_error=True, pwa=True)