server.py

#!/usr/bin/env python3
"""Drama 3 Space: a streaming web UI around Fish Audio Drama 3 on ZeroGPU.

This process serves Gradio and runs Drama 3 itself: zerogpu_backend loads the model
at startup and streams each take from one @spaces.GPU call, through the same
/v1/audio/speech request and events the SGLang Omni server used. Whisper
transcribes reference clips on CPU.

Voices come three ways: cloned from a clip, cast from a text description, or left to
the model. A take can be redone with a director's note; the earlier takes ride along
in the prompt so the note is read against what the listener heard
(drama3_engine.py builds those prompts).

There is no daily take allowance: ZeroGPU applies each visitor's own GPU quota to
their takes. Per-account concurrency and burst limits keep one caller from filling the
queue (accounts.py identifies callers).
"""

from __future__ import annotations

import asyncio
import base64
import collections
import contextlib
import hashlib
import io
import json
import math
import os
import re
import secrets
import sys
import threading
import time
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Literal

import httpx
import numpy as np
import soundfile as sf
import uvicorn
from fastapi import Depends, FastAPI, File, HTTPException, Request, UploadFile
from fastapi.responses import HTMLResponse, Response, StreamingResponse
from pydantic import BaseModel, Field
from stream_transport import StreamingGZipMiddleware

from accounts import (
    Identity,
    identify,
    public_user,
)

APP_DIR = Path(__file__).resolve().parent
INDEX_HTML = APP_DIR / "index.html"
VOICES_DIR = APP_DIR / "voices"


def env_int(name: str, default: int) -> int:
    return int(os.environ.get(name, default))


PORT = env_int("PORT", 7860)
# The in-process engine takes the speech server's request body; the URL is a label.
ENGINE_URL = "zerogpu://drama-3"
# Takes this server runs at once (each on its own ZeroGPU GPU); the rest wait in a
# queue that reports position.
MAX_ACTIVE = env_int("DEMO_MAX_ACTIVE", 12)
MAX_QUEUE = env_int("DEMO_MAX_QUEUE", 64)
# Per account: concurrent takes, and a burst cap. GPU time itself is bounded by each
# visitor's ZeroGPU quota.
PER_CLIENT_ACTIVE = env_int("DEMO_PER_CLIENT_ACTIVE", 2)
PER_CLIENT_REQUESTS = env_int("DEMO_PER_CLIENT_REQUESTS", 60)
RATE_WINDOW_S = env_int("DEMO_RATE_WINDOW_S", 600)

MAX_TEXT_CHARS = 600
MAX_TRANSCRIPT_CHARS = 600
MAX_NOTE_CHARS = 120
MAX_NAME_CHARS = 40
MAX_DESCRIPTION_CHARS = 300
MAX_SCENE_CHARS = 200
# Drama 3 frames are 2048 samples at 44.1 kHz, about 21.5 per second: 1200 frames is
# about 56 s, under the engine's 60 s limit for a take that returns as a reference.
MAX_NEW_TOKENS = 1200
# Drama 3 falls silent mid-take under top-k truncation: at a pause a few silence
# codes outrank the many ways speech can resume, and top-k 30 keeps only the former.
# On an H200, top-k 30 (the SGLang Omni default) left 11 of 40 takes silent to the
# length cap, and no cap left none. Requests send this many candidates instead: the
# end-of-speech token and all 4,096 semantic codes, which also covers the 4,096-entry
# codebook heads. A top_k of 0 asks for it.
ALL_CANDIDATES = 4097
SAMPLE_RATE = 44100
REF_MAX_SECONDS = 20.0
REF_MIN_SECONDS = 1.0
UPLOAD_MAX_BYTES = 25 * 1024 * 1024
# Earlier takes kept in the prompt of a retake.
HISTORY_TAKES = 2
# Keep short natural pauses, trim leading/trailing silence, and abort sustained
# silence. Keep the existing threshold so quiet speech is still passed through.
SILENCE_RMS = 0.003
MAX_PAUSE_S = 0.75
TAIL_S = 0.15
# Hold only the beginning so short failed fragments can be retried before playback.
# The first four engine windows total 64 frames (2.972 s). A 3 s threshold
# misses that boundary and waits for the next 64-frame window. Keep room for
# trimming quiet edges while still withholding short failed openings.
START_BUFFER_S = 2.5
# A whole opening buffer is a large SSE line; clients must receive that entire
# line before playing anything. Send half-second packets after validating it.
AUDIO_PACKET_S = 0.5
STALL_LIMIT_S = 4.0
KEEPALIVE_S = 5.0
# Until the first audio: a ZeroGPU queue, a fresh GPU worker (weights to the GPU,
# compile and CUDA-graph capture) and the prefill.
ENGINE_START_TIMEOUT_S = float(os.environ.get("ENGINE_START_TIMEOUT_S", 150))
ENGINE_IDLE_TIMEOUT_S = 15.0
TAKE_TIMEOUT_S = float(os.environ.get("TAKE_TIMEOUT_S", 180))
TAKE_STORE_BYTES = 512 * 1024 * 1024
VOICE_STORE_SIZE = 256
WHISPER_MODEL = os.environ.get("WHISPER_MODEL", "turbo")
WHISPER_DIR = os.environ.get("WHISPER_DIR", str(Path.home() / ".cache" / "whisper"))

# Must match drama3_engine.ENVELOPE_MARKER (kept separate so the UI imports no torch).
ENVELOPE_MARKER = "@@drama3-space@@"
STRUCTURAL = re.compile(r"<\|")
DIRECTION_SPAN = re.compile(r"\[[^\[\]\n]*\]|</?[A-Za-z][^<>\n]{0,40}>")


def log(message: str) -> None:
    print(f"[space] {message}", flush=True)


# ─── Engine ───────────────────────────────────────────────────────────────────


class Engine:
    """Status of the in-process engine (zerogpu_backend), for the boot screen."""

    def __init__(self) -> None:
        self.state: Literal["starting", "ready", "error"] = "starting"
        self.phase = "Loading 41 GB of weights"
        self.started_at = time.time()
        self.ready_at: float | None = None
        self.last_line = ""
        # The prompt layouts for described characters and directing history.
        self.plugin = False

    def use_in_process(self) -> None:
        """The weights loaded before the server started; GPUs attach per take."""
        self.state = "ready"
        self.phase = "Ready"
        self.ready_at = time.time()
        self.plugin = True

    def status(self) -> dict:
        return {
            "state": self.state,
            "phase": self.phase,
            "elapsed_s": round(time.time() - self.started_at),
            "ready_s": round(self.ready_at - self.started_at) if self.ready_at else None,
            "log": self.last_line,
            "plugin": self.plugin,
        }


# ─── Transcription ────────────────────────────────────────────────────────────


class Transcriber:
    """Whisper on CPU, so transcribing a clip never takes a visitor's GPU time."""

    def __init__(self) -> None:
        self.model = None
        self.state: Literal["waiting", "loading", "ready", "error"] = "waiting"
        self.lock = threading.Lock()

    def start(self) -> None:
        if self.state == "waiting":
            self.state = "loading"
            threading.Thread(target=self.load, daemon=True).start()

    def load(self) -> None:
        try:
            from faster_whisper import WhisperModel
            self.model = WhisperModel(WHISPER_MODEL, device="cpu", compute_type="int8",
                                      cpu_threads=env_int("WHISPER_CPU_THREADS", 2),
                                      download_root=WHISPER_DIR)
            self.state = "ready"
            log(f"whisper {WHISPER_MODEL} ready")
        except Exception as error:  # transcription is optional; typing still works
            self.state = "error"
            log(f"whisper failed to load: {error!r}")

    def transcribe(self, wav: np.ndarray, sample_rate: int) -> tuple[str, str | None]:
        import librosa

        audio = librosa.resample(wav, orig_sr=sample_rate, target_sr=16000).astype(np.float32)
        with self.lock:
            segments, info = self.model.transcribe(
                audio, beam_size=1, temperature=0.0, condition_on_previous_text=False
            )
            return " ".join(segment.text.strip() for segment in segments).strip(), info.language


# ─── Audio helpers ────────────────────────────────────────────────────────────


def decode_audio(data: bytes) -> tuple[np.ndarray, int]:
    """Mono float32 audio from any common container (WAV, MP3, M4A, OGG, WebM...)."""
    try:
        wav, sample_rate = sf.read(io.BytesIO(data), dtype="float32", always_2d=True)
        return wav.mean(axis=1), sample_rate
    except Exception:
        pass
    import av

    parts = []
    with av.open(io.BytesIO(data)) as container:
        stream = container.streams.audio[0]
        resampler = av.AudioResampler(format="flt", layout="mono", rate=SAMPLE_RATE)
        for frame in container.decode(stream):
            parts += [out.to_ndarray().reshape(-1) for out in resampler.resample(frame)]
        parts += [out.to_ndarray().reshape(-1) for out in resampler.resample(None)]
    return (np.concatenate(parts) if parts else np.zeros(0, np.float32)), SAMPLE_RATE


def tidy_reference(wav: np.ndarray, sample_rate: int) -> tuple[np.ndarray, bool]:
    """44.1 kHz, edge silence trimmed, at most REF_MAX_SECONDS cut at a quiet spot."""
    import librosa

    if sample_rate != SAMPLE_RATE:
        wav = librosa.resample(wav, orig_sr=sample_rate, target_sr=SAMPLE_RATE)
    wav = np.nan_to_num(wav.astype(np.float32))
    trimmed, _ = librosa.effects.trim(wav, top_db=40, frame_length=1024, hop_length=256)
    if trimmed.size:
        wav = trimmed
    cut = False
    limit = int(REF_MAX_SECONDS * SAMPLE_RATE)
    if wav.size > limit:
        # End at the quietest 150 ms in the last 40% of the window, so the cut lands
        # between words and the transcript of the kept part stays exact.
        start = int(0.6 * limit)
        window, hop = int(0.15 * SAMPLE_RATE), int(0.01 * SAMPLE_RATE)
        frames = np.lib.stride_tricks.sliding_window_view(wav[start:limit], window)[::hop]
        quietest = int(np.argmin(np.sqrt((frames**2).mean(axis=1))))
        wav = wav[: start + quietest * hop + window // 2].copy()
        fade = min(int(0.03 * SAMPLE_RATE), wav.size)
        wav[-fade:] *= np.linspace(1.0, 0.0, fade, dtype=np.float32)
        cut = True
    peak = float(np.abs(wav).max()) if wav.size else 0.0
    if peak > 0:
        wav = wav * min(10.0, 0.9 / peak)
    return wav.astype(np.float32), cut


def wav_bytes(wav: np.ndarray, sample_rate: int) -> bytes:
    buffer = io.BytesIO()
    sf.write(buffer, np.clip(wav, -1.0, 1.0), sample_rate, format="WAV", subtype="PCM_16")
    return buffer.getvalue()


def pcm_to_float(pcm: bytes) -> np.ndarray:
    return np.frombuffer(pcm, dtype="<i2").astype(np.float32) / 32768.0


def flac_bytes(samples: np.ndarray, sample_rate: int) -> bytes:
    """Lossless PCM16 references: the GPU process decodes the exact samples."""
    buffer = io.BytesIO()
    sf.write(buffer, samples, sample_rate, format="FLAC", subtype="PCM_16")
    return buffer.getvalue()


def spoken_text(line: str) -> str:
    """The words of a line, without direction: what a listener hears."""
    match = re.search(r"^\s*Line:[ \t]*(.+)$", line, re.M | re.S)
    if match and re.match(r"\s*Direction:", line):
        line = match.group(1)
    return re.sub(r"\s+", " ", DIRECTION_SPAN.sub(" ", line)).strip()


def token_budget(body: GenerateBody, choice: VoiceChoice) -> int:
    """A generous speech duration limit, with extra room for slow directions.

    This bounds runaway *audible* output as well as silence. It is a safety limit,
    not an estimate that the model necessarily completes the script within it.
    """
    text = spoken_text(body.text)
    # Include languages without spaces; a word-count alone underbudgets them.
    cjk = len(re.findall(r"[\u3040-\u30ff\u3400-\u9fff\uac00-\ud7af]", text))
    other = re.sub(r"[\u3040-\u30ff\u3400-\u9fff\uac00-\ud7af]", " ", text)
    units = sum(max(1, len(word) / 5) for word in other.split()) + cjk / 2
    seconds = max(18.0, 10.0 + 1.2 * units)
    direction = body.text + " " + (body.note or "")
    if choice.character:
        direction += " " + choice.character.description + " " + (choice.character.scene or "")
    if re.search(r"\b(slow|slowly|drawn[- ]out|unhurried|long pause)\b", direction, re.I):
        seconds *= 1.6
    return min(MAX_NEW_TOKENS, math.ceil(seconds * SAMPLE_RATE / 2048))


class SilenceGuard:
    """Buffer quiet frames so empty takes and trailing silence never go on air."""

    def __init__(self, sample_rate: int) -> None:
        self.frame = max(1, int(0.05 * sample_rate))
        self.max_pause = int(MAX_PAUSE_S * sample_rate)
        self.tail = int(TAIL_S * sample_rate)
        self.stall_limit = int(STALL_LIMIT_S * sample_rate)
        self.pending = b""
        self.quiet = bytearray()
        self.silent_run = 0
        self.dropped = 0
        self.speech_detected = False
        self.stalled = False

    def feed(self, pcm: bytes, *, final: bool = False) -> bytes:
        data = self.pending + pcm
        samples = np.frombuffer(data[: len(data) // 2 * 2], dtype="<i2")
        whole = len(samples) if final else len(samples) // self.frame * self.frame
        self.pending = b"" if final else data[whole * 2 :]
        kept = bytearray()
        for start in range(0, whole, self.frame):
            frame = samples[start : start + self.frame]
            if self.stalled:
                self.dropped += len(frame)
                continue
            if np.sqrt(np.mean((frame.astype(np.float32) / 32768) ** 2)) < SILENCE_RMS:
                self.silent_run += len(frame)
                room = max(0, self.max_pause - len(self.quiet) // 2) if self.speech_detected else 0
                hold = min(room, len(frame))
                self.quiet.extend(frame[:hold].tobytes())
                self.dropped += len(frame) - hold
                if self.silent_run >= self.stall_limit:
                    self.stalled = True
            else:
                self.speech_detected = True
                self.silent_run = 0
                kept.extend(self.quiet)
                self.quiet.clear()
                kept.extend(frame.tobytes())
        if final:
            hold = min(self.tail, len(self.quiet) // 2)
            kept.extend(self.quiet[:hold * 2])
            self.dropped += len(self.quiet) // 2 - hold
            self.quiet.clear()
        return bytes(kept)


async def engine_lines(response, started: float, progress: dict):
    """Bound time without audio even when the server keeps sending heartbeats."""
    iterator = response.aiter_lines().__aiter__()
    pending = None
    try:
        while True:
            now = time.perf_counter()
            idle = ENGINE_IDLE_TIMEOUT_S if progress["has_audio"] else ENGINE_START_TIMEOUT_S
            remaining = min(TAKE_TIMEOUT_S - (now - started), idle - (now - progress["audio_at"]))
            if remaining <= 0:
                raise asyncio.TimeoutError
            if pending is None:
                pending = asyncio.ensure_future(anext(iterator))
            if not (await asyncio.wait({pending}, timeout=min(KEEPALIVE_S, remaining)))[0]:
                yield None
                continue
            try:
                raw = pending.result()
            except StopAsyncIteration:
                break
            pending = None
            yield raw
    finally:
        if pending is not None:
            pending.cancel()
            await asyncio.gather(pending, return_exceptions=True)
        if hasattr(iterator, "aclose"):
            await iterator.aclose()


# ─── Voices and takes ─────────────────────────────────────────────────────────


@dataclass
class Voice:
    id: str
    label: str
    wav: bytes
    transcript: str
    duration: float
    preset: bool = False
    language: str | None = None
    b64: str = field(init=False, repr=False)

    def __post_init__(self) -> None:
        # Keep WAV for browser downloads, and cache the lossless engine reference.
        samples, rate = sf.read(io.BytesIO(self.wav), dtype="int16")
        self.b64 = base64.b64encode(flac_bytes(samples, rate)).decode("ascii")

    def public(self) -> dict:
        return {
            "id": self.id,
            "label": self.label,
            "transcript": self.transcript,
            "duration": round(self.duration, 2),
            "language": self.language,
            "preset": self.preset,
        }


@dataclass
class Character:
    name: str
    description: str
    scene: str | None


@dataclass
class VoiceChoice:
    mode: Literal["clone", "design", "random"]
    voice: Voice | None = None
    transcript: str | None = None
    character: Character | None = None
    # A Surprise/Design voice becomes concrete after its first generated take.
    # Keep that reference across retakes without changing the visible voice mode.
    directing_voice: Voice | None = None


@dataclass
class Take:
    id: str
    choice: VoiceChoice
    line: str
    turn_text: str
    history: list[str]
    pcm: bytes
    sample_rate: int

    @property
    def seconds(self) -> float:
        return len(self.pcm) / 2 / self.sample_rate


class Store:
    def __init__(self) -> None:
        self.voices: collections.OrderedDict[str, Voice] = collections.OrderedDict()
        self.presets: dict[str, Voice] = {}
        self.takes: collections.OrderedDict[str, Take] = collections.OrderedDict()
        self.take_bytes = 0

    def load_presets(self) -> None:
        for entry in json.loads((VOICES_DIR / "voices.json").read_text()):
            wav, sample_rate = decode_audio((VOICES_DIR / entry["file"]).read_bytes())
            wav, _ = tidy_reference(wav, sample_rate)
            self.presets[entry["id"]] = Voice(
                id=entry["id"],
                label=entry["label"],
                wav=wav_bytes(wav, SAMPLE_RATE),
                transcript=entry["transcript"],
                duration=wav.size / SAMPLE_RATE,
                preset=True,
                language="en",
            )
        log(f"loaded {len(self.presets)} preset voices")

    def voice(self, voice_id: str) -> Voice | None:
        if voice_id in self.presets:
            return self.presets[voice_id]
        voice = self.voices.get(voice_id)
        if voice is not None:
            self.voices.move_to_end(voice_id)
        return voice

    def add_voice(self, voice: Voice) -> None:
        self.voices[voice.id] = voice
        self.voices.move_to_end(voice.id)
        while len(self.voices) > VOICE_STORE_SIZE:
            self.voices.popitem(last=False)

    def take(self, take_id: str) -> Take | None:
        return self.takes.get(take_id)

    def add_take(self, take: Take) -> None:
        self.takes[take.id] = take
        self.take_bytes += len(take.pcm)
        while self.take_bytes > TAKE_STORE_BYTES and len(self.takes) > 1:
            _, old = self.takes.popitem(last=False)
            self.take_bytes -= len(old.pcm)


# ─── Admission ────────────────────────────────────────────────────────────────


class Ticket:
    def __init__(self) -> None:
        self.granted = asyncio.Event()


class Gate:
    """FIFO admission to MAX_ACTIVE engine streams; waiters learn their position."""

    def __init__(self, slots: int) -> None:
        self.slots = slots
        self.active = 0
        self.waiting: collections.deque[Ticket] = collections.deque()

    def enqueue(self) -> Ticket:
        ticket = Ticket()
        if self.active < self.slots and not self.waiting:
            self.active += 1
            ticket.granted.set()
        else:
            self.waiting.append(ticket)
        return ticket

    def position(self, ticket: Ticket) -> int:
        try:
            return self.waiting.index(ticket) + 1
        except ValueError:
            return 0

    def leave(self, ticket: Ticket) -> None:
        if ticket.granted.is_set():
            self.active -= 1
        else:
            self.waiting.remove(ticket)
        while self.waiting and self.active < self.slots:
            self.active += 1
            self.waiting.popleft().granted.set()


class ClientLimits:
    """Per-account concurrency and burst limits.

    Expiration is a final recovery guard. Streams reserve their slot only once
    consumption starts, and release it when their iterator closes.
    """

    SLOT_TTL_S = 300

    def __init__(self) -> None:
        self.active: dict[str, list[float]] = collections.defaultdict(list)
        self.recent: dict[str, collections.deque[float]] = collections.defaultdict(collections.deque)

    def check(self, client: str) -> str | None:
        """Check limits without reserving a slot or recording a new request."""
        now = time.time()
        recent = self.recent[client]
        while recent and now - recent[0] > RATE_WINDOW_S:
            recent.popleft()
        slots = [t for t in self.active.get(client, []) if now - t < self.SLOT_TTL_S]
        if slots:
            self.active[client] = slots
        else:
            self.active.pop(client, None)
        if len(slots) >= PER_CLIENT_ACTIVE:
            return f"You already have {PER_CLIENT_ACTIVE} takes in progress. Wait for one to finish."
        if len(recent) >= PER_CLIENT_REQUESTS:
            wait = int(RATE_WINDOW_S - (now - recent[0])) + 1
            return f"That's a lot of takes! Give the GPU a breather and try again in {wait}s."
        return None

    def admit(self, client: str) -> float | str:
        """Reserve atomically when a response actually begins streaming."""
        refused = self.check(client)
        if refused:
            return refused
        now = time.time()
        self.recent[client].append(now)
        self.active[client].append(now)
        return now

    def release(self, client: str, slot: float) -> None:
        slots = self.active.get(client, [])
        if slot in slots:
            slots.remove(slot)
        if not slots:
            self.active.pop(client, None)


# ─── App state ────────────────────────────────────────────────────────────────

engine = Engine()
transcriber = Transcriber()
store = Store()
gate = Gate(MAX_ACTIVE)
limits = ClientLimits()
stats = {"takes": 0, "audio_s": 0.0}
recent_ttfa: collections.deque[float] = collections.deque(maxlen=50)
http = None  # zerogpu_backend.EngineClient once the model is loaded


def gpu_info() -> dict | None:
    # GPUs attach per take, so this process has no device to report on.
    return {"name": "ZeroGPU", "remote": True}


def shared_memory_report() -> str:
    try:
        stat = os.statvfs("/dev/shm")
        return f"{stat.f_blocks * stat.f_frsize / 2**30:.1f} GiB"
    except OSError:
        return "unavailable"


@contextlib.asynccontextmanager
async def lifespan(_: FastAPI):
    global http
    log(f"/dev/shm: {shared_memory_report()}")
    await asyncio.to_thread(store.load_presets)
    if os.environ.get("DEMO_PREVIEW") == "1":
        engine.phase = "UI preview · GPU engine offline"
    else:
        import zerogpu_backend  # app.py imports it first, so the weights are loaded

        http = zerogpu_backend.EngineClient()
        engine.use_in_process()
        transcriber.start()
        log("generation runs in this process, one ZeroGPU call per take")
        if os.environ.get("DRAMA3_WARM_UP", "1") == "1" and hasattr(zerogpu_backend, "warm_up"):
            voice = next(iter(store.presets.values()), None)
            reference = {"data": voice.b64, "media_type": "audio/flac", "text": voice.transcript} if voice else None
            threading.Thread(target=zerogpu_backend.warm_up, args=(reference,), daemon=True).start()
    try:
        yield
    finally:
        if http is not None:
            await http.aclose()


app = FastAPI(title="Drama 3", lifespan=lifespan)
from starlette.middleware.sessions import SessionMiddleware
app.add_middleware(SessionMiddleware, secret_key=secrets.token_hex(32))


# ─── Routes ───────────────────────────────────────────────────────────────────


@app.get("/legacy", include_in_schema=False)
async def index(request: Request) -> HTMLResponse:
    return HTMLResponse(INDEX_HTML.read_text(encoding="utf-8"), headers={"Cache-Control": "no-store"})


@app.get("/favicon.ico", include_in_schema=False)
async def favicon() -> Response:
    return Response(status_code=204)


@app.get("/api/status")
async def status(identity: Identity = Depends(identify)) -> dict:
    ttfas = sorted(recent_ttfa)
    return {
        "user": public_user(identity),
        "engine": engine.status(),
        "transcriber": transcriber.state,
        "gpu": await asyncio.to_thread(gpu_info),
        "load": {"active": gate.active, "waiting": len(gate.waiting), "slots": gate.slots},
        "stats": {
            "takes": stats["takes"],
            "audio_s": round(stats["audio_s"]),
            "ttfa_ms_p50": round(ttfas[len(ttfas) // 2]) if ttfas else None,
        },
        "limits": {
            "max_text_chars": MAX_TEXT_CHARS,
            "max_reference_s": REF_MAX_SECONDS,
            "max_audio_s": round(MAX_NEW_TOKENS * 2048 / SAMPLE_RATE),
        },
        "voices": [voice.public() for voice in store.presets.values()],
    }


@app.get("/api/voices/{voice_id}.wav", dependencies=[Depends(identify)])
async def voice_audio(voice_id: str) -> Response:
    voice = store.voice(voice_id)
    if voice is None:
        raise HTTPException(404, "That voice has expired. Add it again.")
    return Response(voice.wav, media_type="audio/wav", headers={"Cache-Control": "max-age=3600"})


@app.get("/api/takes/{take_id}.wav", dependencies=[Depends(identify)])
async def take_audio(take_id: str) -> Response:
    take = store.take(take_id)
    if take is None:
        raise HTTPException(404, "That take has expired.")
    return Response(
        wav_bytes(pcm_to_float(take.pcm), take.sample_rate),
        media_type="audio/wav",
        headers={"Content-Disposition": f'attachment; filename="drama3-{take_id}.wav"'},
    )


async def register_voice(wav: np.ndarray, sample_rate: int, label: str, fallback_text: str = "") -> dict:
    wav, cut = await asyncio.to_thread(tidy_reference, wav, sample_rate)
    duration = wav.size / SAMPLE_RATE
    if duration < REF_MIN_SECONDS:
        raise HTTPException(400, "That clip is too short. Record at least a second of speech.")
    transcript, language, transcribed = fallback_text, None, False
    if transcriber.state == "ready":
        try:
            transcript, language = await asyncio.to_thread(transcriber.transcribe, wav, SAMPLE_RATE)
            transcribed = True
        except Exception as error:
            log(f"transcription failed: {error!r}")
    audio = wav_bytes(wav, SAMPLE_RATE)
    voice = Voice(
        id=hashlib.sha1(audio).hexdigest()[:16],
        label=label,
        wav=audio,
        transcript=transcript,
        duration=duration,
        language=language,
    )
    store.add_voice(voice)
    return {**voice.public(), "trimmed": cut, "transcribed": transcribed}


@app.post("/api/voices", dependencies=[Depends(identify)])
async def add_voice(audio: UploadFile = File(...)) -> dict:
    data = await audio.read(UPLOAD_MAX_BYTES + 1)
    if len(data) > UPLOAD_MAX_BYTES:
        raise HTTPException(400, "That file is too big. Use a clip under 25 MB.")
    try:
        wav, sample_rate = await asyncio.to_thread(decode_audio, data)
    except Exception:
        raise HTTPException(400, "Couldn't read that audio. Try a WAV or MP3 file.")
    if not wav.size or not np.isfinite(wav).all() or float(np.abs(wav).max()) == 0.0:
        raise HTTPException(400, "That recording is silent. Check the microphone and try again.")
    label = Path(audio.filename or "").stem[:24] or "My voice"
    return await register_voice(wav, sample_rate, label)


class FromTake(BaseModel):
    take_id: str
    label: str | None = Field(default=None, max_length=40)


@app.post("/api/voices/from-take", dependencies=[Depends(identify)])
async def voice_from_take(body: FromTake) -> dict:
    take = store.take(body.take_id)
    if take is None:
        raise HTTPException(404, "That take has expired. Generate it again.")
    label = body.label or (take.choice.character.name if take.choice.character else "Found voice")
    return await register_voice(
        pcm_to_float(take.pcm), take.sample_rate, label, fallback_text=spoken_text(take.line)
    )


class VoiceBody(BaseModel):
    mode: Literal["clone", "design", "random"]
    id: str | None = None
    transcript: str | None = None
    name: str | None = None
    description: str | None = None
    scene: str | None = None


class GenerateBody(BaseModel):
    text: str
    voice: VoiceBody
    retake_of: str | None = None
    note: str | None = None
    temperature: float = Field(default=1.0, ge=0.1, le=1.5)
    top_p: float = Field(default=0.9, ge=0.1, le=1.0)
    top_k: int = Field(default=0, ge=0, le=ALL_CANDIDATES)  # 0: no cap


def clean_text(value: str | None, what: str, limit: int, *, single_line: bool = False) -> str:
    value = (value or "").replace("\r\n", "\n").strip()
    if single_line:
        value = re.sub(r"\s+", " ", value)
    if not value:
        raise HTTPException(400, f"Add the {what} first.")
    if len(value) > limit:
        raise HTTPException(400, f"Keep the {what} under {limit} characters.")
    if STRUCTURAL.search(value) or ENVELOPE_MARKER in value:
        raise HTTPException(400, f"Remove the <|…|> tokens from the {what}; this page writes those itself.")
    return value


def resolve_voice(body: VoiceBody) -> VoiceChoice:
    if body.mode == "clone":
        voice = store.voice(body.id or "")
        if voice is None:
            raise HTTPException(404, "That voice has expired. Upload or record it again.")
        transcript = clean_text(
            body.transcript if body.transcript is not None else voice.transcript,
            "reference transcript", MAX_TRANSCRIPT_CHARS, single_line=True,
        )
        return VoiceChoice(mode="clone", voice=voice, transcript=transcript)
    if body.mode == "design":
        name = clean_text(body.name or "Narrator", "character name", MAX_NAME_CHARS, single_line=True)
        if ":" in name:
            raise HTTPException(400, "Leave ':' out of the character name.")
        description = clean_text(body.description, "voice description", MAX_DESCRIPTION_CHARS, single_line=True)
        scene = (body.scene or "").strip()
        scene = clean_text(scene, "scene", MAX_SCENE_CHARS, single_line=True) if scene else None
        return VoiceChoice(mode="design", character=Character(name, description, scene))
    return VoiceChoice(mode="random")


def retake_text(line: str, note: str) -> str:
    """Give the new performance its own direction; history holds the old one.

    Repeating the original script's acting tags makes them compete with the
    director's new note. Preserve the spoken words, replace the acting cues,
    and let the user decide whether their note is absolute or comparative.
    """
    note = clean_text(re.sub(r"[\[\]<>|]", "", note), "note", MAX_NOTE_CHARS, single_line=True)
    words = clean_text(spoken_text(line), "spoken script", MAX_TEXT_CHARS)
    # A single short prefix is often drowned out by the reference performance.
    # Keep the note active through each sentence rather than only at the opening.
    sentences = re.split(r"(?<=[.!?。！？])\s+", words)
    directed = " ".join(f"[{note}] {sentence}" for sentence in sentences)
    return f"Direction: {note}\nLine: {directed}"


def engine_request(choice: VoiceChoice, turn_text: str, history: list[Take], body: GenerateBody) -> dict:
    references = []
    if choice.mode == "clone":
        references.append({"data": choice.voice.b64, "media_type": "audio/flac", "text": choice.transcript})
    elif history:
        voice = choice.directing_voice
        if voice is None:
            raise HTTPException(500, "The generated voice reference is missing for this retake.")
        references.append({"data": voice.b64, "media_type": "audio/flac", "text": voice.transcript})
    for take in history:
        references.append({
            "data": base64.b64encode(flac_bytes(np.frombuffer(take.pcm, dtype="<i2"), take.sample_rate)).decode("ascii"),
            "media_type": "audio/flac",
            "text": spoken_text(take.line) or take.line,
        })
    if choice.mode == "design" or history:
        # Directing is the voice-conditioned conversation of Figure 12. Casting
        # a described character (Figure 13) is only the initial generation task.
        character = choice.character if not history else None
        envelope = {
            "text": turn_text,
            "voice_references": 1 if choice.mode == "clone" or history else 0,
            "history": [take.turn_text for take in history],
            "character": (
                {"name": character.name, "description": character.description, "scene": character.scene}
                if character else None
            ),
        }
        text = ENVELOPE_MARKER + json.dumps(envelope, ensure_ascii=False)
    else:
        text = turn_text
    return {
        "input": text,
        "references": references or None,
        "stream": True,
        "response_format": "pcm",
        "stream_format": "sse",
        "temperature": body.temperature,
        "top_p": body.top_p,
        "top_k": body.top_k or ALL_CANDIDATES,
        "max_new_tokens": token_budget(body, choice),
    }


def engine_error_message(raw: bytes | str) -> str:
    try:
        payload = json.loads(raw)
        message = payload.get("error", {}).get("message") or payload.get("message") or payload.get("detail")
    except (ValueError, AttributeError):
        message = None
    message = str(message or raw or "The engine refused this request.")
    return message.removeprefix("Drama 3 request:").strip()[:400]


def sse(payload: dict) -> str:
    return f"data: {json.dumps(payload)}\n\n"


@app.post("/api/generate")
async def generate(
    body: GenerateBody, identity: Identity = Depends(identify), request: Request = None
) -> StreamingResponse:
    if request is not None:
        # API callers: let ZeroGPU see this request's headers (Gradio sets this
        # itself for the Studio and Workflow).
        from gradio import Request as GradioRequest
        from gradio.context import LocalContext

        if LocalContext.request.get(None) is None:
            LocalContext.request.set(GradioRequest(request=request))
    if engine.state != "ready":
        raise HTTPException(503, "Drama 3 is still warming up. It will be ready in a moment.")
    line = clean_text(body.text, "script", MAX_TEXT_CHARS)

    history: list[Take] = []
    if body.retake_of:
        previous = store.take(body.retake_of)
        if previous is None:
            raise HTTPException(404, "That take has expired. Generate a fresh one.")
        choice = previous.choice
        if choice.mode != "clone" and choice.directing_voice is None:
            choice = replace(choice, directing_voice=Voice(
                id=f"take-{previous.id}",
                label=choice.character.name if choice.character else "Surprise",
                wav=wav_bytes(pcm_to_float(previous.pcm), previous.sample_rate),
                transcript=spoken_text(previous.line) or previous.line,
                duration=previous.seconds,
            ))
        chain = [take for take in map(store.take, previous.history) if take is not None]
        history = chain[max(0, len(chain) - (HISTORY_TAKES - 1)):] + [previous]
        note = clean_text(body.note, "note", MAX_NOTE_CHARS, single_line=True)
        turn_text = retake_text(line, note)
    else:
        choice = resolve_voice(body.voice)
        turn_text = line

    if len(gate.waiting) >= MAX_QUEUE:
        raise HTTPException(503, "The stage is packed right now. Try again in a minute.")
    refusal = limits.check(identity.key)
    if refusal:
        raise HTTPException(429, refusal)

    payload = engine_request(choice, turn_text, history, body)
    take_id = secrets.token_urlsafe(9)

    async def stream():
        slot = limits.admit(identity.key)
        if isinstance(slot, str):
            yield sse({"type": "error", "message": slot})
            return
        ticket = gate.enqueue()
        queued_at = time.perf_counter()
        try:
            while not ticket.granted.is_set():
                yield sse({"type": "queued", "position": gate.position(ticket)})
                try:
                    await asyncio.wait_for(ticket.granted.wait(), timeout=1.0)
                except asyncio.TimeoutError:
                    pass
            admitted = time.perf_counter()
            queue_ms = (admitted - queued_at) * 1000
            started = time.perf_counter()
            timings = {"queue": round(queue_ms)}
            pcm = bytearray()
            ttfa_ms = None
            finish = None
            sent_audio = False

            async def audio_events(kept, *, flush=False):
                nonlocal ttfa_ms, sent_audio
                buffered.extend(kept)
                if not buffered:
                    return
                if not sent_audio and not flush and len(buffered) < START_BUFFER_S * sample_rate * 2:
                    return
                kept = bytes(buffered)
                buffered.clear()
                if not sent_audio:
                    ttfa_ms = (time.perf_counter() - started) * 1000
                    timings["opening_validated"] = round(ttfa_ms)
                    yield sse({"type": "start", "take_id": take_id, "sample_rate": sample_rate,
                               "queue_ms": round(queue_ms)})
                pcm.extend(kept)
                packet_bytes = max(2, round(AUDIO_PACKET_S * sample_rate) * 2)
                for offset in range(0, len(kept), packet_bytes):
                    packet = kept[offset:offset + packet_bytes]
                    chunk = {"type": "audio", "pcm": base64.b64encode(packet).decode("ascii")}
                    if not sent_audio:
                        chunk["ttfa_ms"] = round(ttfa_ms)
                        timings["first_packet_after_admission"] = round((time.perf_counter() - admitted) * 1000)
                        sent_audio = True
                    yield sse(chunk)

            for attempt in range(2):
                buffered = bytearray()
                guard = None
                response = None
                attempt_started = time.perf_counter()
                opening = asyncio.ensure_future(http.send(
                    http.build_request("POST", f"{ENGINE_URL}/v1/audio/speech", json=payload), stream=True
                ))
                try:
                    while not opening.done():
                        remaining = min(ENGINE_START_TIMEOUT_S - (time.perf_counter() - attempt_started),
                                        TAKE_TIMEOUT_S - (time.perf_counter() - started))
                        if remaining <= 0:
                            raise asyncio.TimeoutError
                        if not (await asyncio.wait({opening}, timeout=min(KEEPALIVE_S, remaining)))[0]:
                            yield ": keepalive\n\n"
                    response = opening.result()
                    if response.status_code != 200:
                        message = engine_error_message(await asyncio.wait_for(response.aread(), ENGINE_IDLE_TIMEOUT_S))
                        log(f"engine refused {take_id}: {message}")
                        yield sse({"type": "error", "message": message})
                        return
                    sample_rate = int(response.headers.get("x-sample-rate", SAMPLE_RATE))
                    guard = SilenceGuard(sample_rate)
                    progress = {"audio_at": attempt_started, "has_audio": False}
                    async with contextlib.aclosing(engine_lines(response, started, progress)) as lines:
                        async for raw in lines:
                            if raw is None:
                                yield ": keepalive\n\n"
                                continue
                            if not raw.startswith("data: "):
                                continue
                            event = json.loads(raw[6:])
                            kind = event.get("type")
                            if kind == "speech.audio.delta":
                                timings.setdefault("engine_first_audio", round((time.perf_counter() - started) * 1000))
                                progress.update(audio_at=time.perf_counter(), has_audio=True)
                                kept = guard.feed(base64.b64decode(event["audio"]))
                                async for chunk in audio_events(kept):
                                    yield chunk
                                if guard.stalled:
                                    finish = "stall"
                                    break
                            elif kind == "speech.audio.done":
                                finish = event.get("finish_reason") or "stop"
                                if event.get("timings"):
                                    timings["engine"] = event["timings"]
                                break  # A done event is final; do not wait for socket EOF.
                            elif kind == "error":
                                message = engine_error_message(json.dumps(event))
                                log(f"engine failed {take_id}: {message}")
                                yield sse({"type": "error", "message": message})
                                return
                except (asyncio.TimeoutError, httpx.TimeoutException):
                    finish = "timeout"
                finally:
                    # Cancellation can race with response headers arriving.
                    if not opening.done():
                        opening.cancel()
                    await asyncio.gather(opening, return_exceptions=True)
                    if response is None and not opening.cancelled() and opening.exception() is None:
                        response = opening.result()
                    if response is not None:
                        await response.aclose()
                if guard is not None:
                    tail = guard.feed(b"", final=True)
                    # A failed opening's tail must not push it over the playback
                    # threshold just before it is discarded and retried.
                    if sent_audio or finish in {"stop", "timeout"}:
                        async for chunk in audio_events(tail):
                            yield chunk
                # Normal short lines finish without filling the startup buffer.
                # A stalled fragment is discarded only if nothing has gone on air.
                if finish in {"stop", "timeout"} or sent_audio:
                    async for chunk in audio_events(b"", flush=True):
                        yield chunk
                if pcm:
                    break
                if attempt == 0 and finish != "timeout":
                    log(f"take {take_id}: silent or incomplete opening, retrying once")
                    yield sse({"type": "retrying", "message": "The actor did not finish the opening. Trying once more…"})
                    finish = None
                else:
                    message = ("The actor did not respond in time. Try again." if finish == "timeout"
                               else "No complete opening came out after two attempts. Try a different voice or direction.")
                    yield sse({"type": "error", "message": message})
                    return
            warning = {
                "stall": "The actor went quiet before ending the take. Part of the line may be missing; try a retake.",
                "length": "The actor reached the time limit. The line may be incomplete; try a retake or a shorter script.",
                "timeout": "The actor stopped responding. This take may be incomplete; try a retake.",
            }.get(finish)
            if finish is None:
                finish = "incomplete"
                warning = "The stream ended before the actor finished. This take may be incomplete; try a retake."
            total_ms = (time.perf_counter() - started) * 1000
            seconds = len(pcm) / 2 / sample_rate
            trimmed_s = guard.dropped / sample_rate
            store.add_take(Take(
                id=take_id,
                choice=choice,
                line=line,
                turn_text=turn_text,
                history=[take.id for take in history],
                pcm=bytes(pcm),
                sample_rate=sample_rate,
            ))
            stats["takes"] += 1
            stats["audio_s"] += seconds
            recent_ttfa.append(ttfa_ms or total_ms)
            log(f"take {take_id}: {choice.mode}, {len(history)} earlier, {seconds:.1f}s audio, "
                f"{trimmed_s:.1f}s silence trimmed, ttfa {ttfa_ms or 0:.0f}ms, "
                f"total {total_ms:.0f}ms, queue {queue_ms:.0f}ms, {finish}")
            yield sse({
                "type": "done",
                "take_id": take_id,
                "ttfa_ms": round(ttfa_ms or total_ms),
                "total_ms": round(total_ms),
                "queue_ms": round(queue_ms),
                "timings_ms": timings,
                "audio_s": round(seconds, 2),
                "speed": round(seconds / (total_ms / 1000), 2) if total_ms > 0 else None,
                "finish_reason": finish,
                "warning": warning,
                "retry_count": attempt,
                "trimmed_s": round(trimmed_s, 1),
                "turn_text": turn_text,
            })
        except httpx.HTTPError as error:
            log(f"engine connection failed for {take_id}: {error!r}")
            yield sse({"type": "error", "message": "Lost the connection to the engine. Try again."})
        finally:
            gate.leave(ticket)
            limits.release(identity.key, slot)

    return StreamingResponse(
        stream(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


from gradio_ui import mount_apps

app = mount_apps(sys.modules[__name__])
app.add_middleware(StreamingGZipMiddleware)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info", access_log=False)