import functools
import os
import time

# Disable Dynamo by default for Space stability; the CLI script does the same.
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import gradio as gr
import numpy as np
import torch

try:
    import spaces
except ImportError:
    class _SpacesFallback:
        @staticmethod
        def GPU(*_args, **_kwargs):
            def _decorator(func):
                return func

            return _decorator

    spaces = _SpacesFallback()

from moss_soundeffect_v2 import MossSoundEffectPipeline

MODEL_PATH = "OpenMOSS-Team/MOSS-SoundEffect-v2.0"
DEFAULT_DEVICE = "cuda"
MAX_INFERENCE_SECONDS = 30


@functools.lru_cache(maxsize=1)
def load_backend(model_dir: str, device_str: str):
    device = torch.device(device_str if torch.cuda.is_available() else "cpu")
    pipe = MossSoundEffectPipeline.from_pretrained(
        model_dir,
        torch_dtype=torch.bfloat16 if device.type == "cuda" else torch.float32,
        device=str(device),
    )
    return pipe, device


@spaces.GPU(duration=180)
def run_inference(prompt, seconds, steps, cfg_scale, sigma_shift, seed):
    if not (prompt or "").strip():
        raise ValueError("Please enter a prompt describing the audio you want to generate.")

    seconds = round(float(seconds), 1)
    if seconds <= 0:
        raise ValueError("Duration must be greater than 0.")
    if seconds > MAX_INFERENCE_SECONDS:
        raise ValueError(f"Duration must be no greater than {MAX_INFERENCE_SECONDS}s.")

    started_at = time.monotonic()
    pipe, _ = load_backend(MODEL_PATH, DEFAULT_DEVICE)

    audio = pipe(
        prompt=prompt,
        seconds=seconds,
        num_inference_steps=int(steps),
        cfg_scale=float(cfg_scale),
        sigma_shift=float(sigma_shift),
        seed=int(seed),
    )

    audio_np = audio[0].detach().float().cpu().numpy()
    if audio_np.ndim > 1 and audio_np.shape[0] == 1:
        audio_np = audio_np.squeeze(0)
    elif audio_np.ndim > 1:
        audio_np = audio_np.T
    audio_np = audio_np.astype(np.float32, copy=False)

    elapsed = time.monotonic() - started_at
    status = (
        f"Done | elapsed: {elapsed:.2f}s | "
        f"duration={seconds:.1f}s, steps={int(steps)}, "
        f"cfg_scale={float(cfg_scale):.2f}, sigma_shift={float(sigma_shift):.2f}, "
        f"seed={int(seed)}"
    )
    return (pipe.sample_rate, audio_np), status


with gr.Blocks(title="MOSS-SoundEffect v2.0") as demo:
    gr.Markdown(
        """
        # MOSS-SoundEffect v2.0
        Text-to-audio diffusion demo.
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="Prompt",
                lines=8,
                value="The crisp, rhythmic click-clack of fast typing on a mechanical keyboard.",
            )
            seconds = gr.Slider(1, MAX_INFERENCE_SECONDS, step=0.1, value=10, label="Duration (seconds)")
            with gr.Accordion("Sampling Parameters", open=True):
                steps = gr.Slider(10, 150, step=1, value=50, label="num_inference_steps")
                cfg_scale = gr.Slider(1.0, 8.0, step=0.1, value=4.0, label="cfg_scale")
                sigma_shift = gr.Slider(0.0, 10.0, step=0.1, value=5.0, label="sigma_shift")
                seed = gr.Number(value=0, label="seed", precision=0)
            run_btn = gr.Button("Generate Sound Effect", variant="primary")
        with gr.Column(scale=2):
            output_audio = gr.Audio(label="Output Audio", type="numpy")
            status = gr.Textbox(label="Status", lines=4, interactive=False)

    run_btn.click(
        fn=run_inference,
        inputs=[prompt, seconds, steps, cfg_scale, sigma_shift, seed],
        outputs=[output_audio, status],
    )

demo.queue(max_size=16, default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch()