"""Woosh-DFlow text-to-audio BACKEND Space (plain Gradio 6, no pyharp).

This is the *backend* half of the HARP two-Space workflow for Sony AI's Woosh
sound-effect model. It reuses Woosh's own inference code (the same calls the
upstream ``gradio_Woosh-DFlow.py`` demo makes) and exposes a single clean
``/generate`` API endpoint that a thin pyharp HARP frontend proxies to via
``gradio_client``.

Why a separate backend Space at all: Woosh needs ``python>=3.12``,
``torch==2.8.0`` and ``gradio>=6.9.0``, which cannot coexist in one process with
pyharp v0.3.0 (it hard-pins ``gradio==5.28.0``). Isolating Woosh in its own
Space sidesteps that conflict entirely.

Weights (CC-BY-NC, from the official v1.0.0 GitHub release) are baked into the
Docker image at build time (see the Dockerfile), so this file just loads them.
"""

from __future__ import annotations

import logging
import os
import tempfile

import gradio as gr
import numpy as np
import soundfile as sf
import torch

from woosh.components.base import LoadConfig
from woosh.inference.flowmap_sampler import sample_euler
from woosh.model.flowmap_from_pretrained import FlowMapFromPretrained

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("woosh-backend")

# Matches the upstream demo: 128-channel latents, ~5s at 48 kHz after the AE.
SAMPLE_RATE = 48000
LATENT_CHANNELS = 128
LATENT_FRAMES = 501
NUM_STEPS = 4  # Woosh-DFlow is the *distilled* model: 4 sampling steps.
RENOISE = [0, 0.5, 0.5, 0.3]
CHECKPOINT = os.environ.get("WOOSH_CHECKPOINT", "checkpoints/Woosh-DFlow")

if torch.cuda.is_available():
    DEVICE = "cuda"
elif getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
    DEVICE = "mps"
else:
    DEVICE = "cpu"

log.info("Loading Woosh-DFlow from %s on %s ...", CHECKPOINT, DEVICE)
LDM = FlowMapFromPretrained(LoadConfig(path=CHECKPOINT)).eval().to(DEVICE)
log.info("Model loaded on %s.", DEVICE)


@torch.inference_mode()
def generate(prompt: str, cfg_scale: float = 4.5, seed: int = -1) -> str:
    """Generate one sound effect from a text prompt; return a WAV file path."""

    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a text prompt describing the sound effect.")

    if seed is None or int(seed) < 0:
        seed = int.from_bytes(os.urandom(4), "big") % (2**31)
    torch.manual_seed(int(seed))

    noise = torch.randn(1, LATENT_CHANNELS, LATENT_FRAMES).to(DEVICE)
    cond = LDM.get_cond(
        {"audio": None, "description": [prompt]}, no_dropout=True, device=DEVICE
    )
    x_fake = sample_euler(
        model=LDM,
        noise=noise,
        cond=cond,
        num_steps=NUM_STEPS,
        renoise=RENOISE,
        cfg=float(cfg_scale),
    )
    audio = LDM.autoencoder.inverse(x_fake).cpu().float()

    # Peak-normalize to avoid clipping, then write a float32 WAV.
    peak = audio.abs().amax(dim=-1, keepdim=True).clamp(min=1.0)
    audio = (audio / peak).clamp(-1.0, 1.0)
    wav = np.ascontiguousarray(audio[0].squeeze().numpy())

    out_path = os.path.join(tempfile.mkdtemp(), "woosh_dflow.wav")
    sf.write(out_path, wav, SAMPLE_RATE)
    return out_path


with gr.Blocks(title="Woosh-DFlow backend") as demo:
    gr.Markdown(
        "# Woosh-DFlow \u2014 text-to-audio backend\n"
        "Plain-Gradio **backend** for the HARP two-Space workflow. Point a pyharp "
        "remote frontend at the `/generate` API endpoint."
    )
    prompt = gr.Textbox(
        label="Prompt",
        placeholder="e.g. sportscar engine revving and driving away quickly",
    )
    cfg = gr.Slider(0.0, 15.0, value=4.5, step=0.1, label="CFG scale")
    seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
    out = gr.Audio(label="Generated audio", type="filepath")
    gr.Button("Generate", variant="primary").click(
        generate, inputs=[prompt, cfg, seed], outputs=out, api_name="generate"
    )

demo.queue(max_size=8).launch(
    server_name="0.0.0.0",
    server_port=int(os.environ.get("PORT", "7860")),
    show_error=True,
)