# FastAPI shim around diffusers' StableAudioPipeline.
# Single endpoint POST /v1/audio/sfx returns a WAV blob.
# Model is loaded once on startup and held in process memory.
import asyncio
import io
import os
import time
from contextlib import asynccontextmanager
from typing import Optional

import soundfile as sf
import torch
from diffusers import StableAudioPipeline
from fastapi import FastAPI, HTTPException, Response
from pydantic import BaseModel, Field

MODEL_ID = os.environ.get("SAO_MODEL", "stabilityai/stable-audio-open-1.0")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

state: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    print(f"[sao] loading {MODEL_ID} on {DEVICE} ({DTYPE})", flush=True)
    t0 = time.time()
    pipe = StableAudioPipeline.from_pretrained(MODEL_ID, torch_dtype=DTYPE)
    pipe = pipe.to(DEVICE)
    state["pipe"] = pipe
    # Single-GPU diffusers pipelines aren't reentrant — concurrent calls
    # share the same scheduler step_index counter and corrupt each
    # other (observed: IndexError in cosine_dpmsolver_multistep when
    # two requests overlap mid-run). Serialize at the request boundary.
    state["lock"] = asyncio.Lock()
    print(f"[sao] loaded in {time.time() - t0:.1f}s", flush=True)
    yield
    state.clear()


app = FastAPI(lifespan=lifespan)


class SfxRequest(BaseModel):
    prompt: str = Field(..., min_length=1)
    negative_prompt: Optional[str] = "Low quality."
    duration: float = Field(10.0, gt=0.5, le=47.0)
    steps: int = Field(100, ge=10, le=300)
    seed: Optional[int] = None
    cfg_scale: float = Field(7.0, gt=0.0, le=20.0)


@app.get("/health")
def health():
    return {
        "status": "ok",
        "model": MODEL_ID,
        "device": DEVICE,
        "loaded": "pipe" in state,
    }


@app.post("/v1/audio/sfx")
async def sfx(req: SfxRequest):
    pipe = state.get("pipe")
    lock = state.get("lock")
    if pipe is None or lock is None:
        raise HTTPException(503, "model not loaded yet")

    generator = None
    if req.seed is not None:
        generator = torch.Generator(DEVICE).manual_seed(req.seed)

    # Hold the lock for the whole inference + encode. Concurrent
    # requests queue cleanly instead of racing the scheduler's
    # step_index counter into an IndexError.
    async with lock:
        audio = await asyncio.to_thread(
            lambda: pipe(
                req.prompt,
                negative_prompt=req.negative_prompt,
                num_inference_steps=req.steps,
                audio_end_in_s=req.duration,
                num_waveforms_per_prompt=1,
                guidance_scale=req.cfg_scale,
                generator=generator,
            ).audios
        )
        waveform = audio[0].T.float().cpu().numpy()

    buf = io.BytesIO()
    sf.write(buf, waveform, pipe.vae.sampling_rate, format="WAV")
    buf.seek(0)
    return Response(content=buf.read(), media_type="audio/wav")