#!/usr/bin/env bash
[ -n "$BASH_VERSION" ] || exec bash "$0" "$@"
set -uo pipefail          # no -e: every step has an explicit || fail
fail() { printf '\033[31mx %s\033[0m\n' "$1" >&2; printf '\033[31mRun this file again to retry.\033[0m\n' >&2; exit 1; }
# In a terminal the banner stays in place: clear the screen, print it on rows 2-6,
# then limit scrolling to row 8 and below. The EXIT trap gives the whole
# window back (cursor kept) however the script ends.
if [ -t 1 ]; then
  trap 'printf "\0337\033[r\0338"' EXIT
  printf '\033[2J\033[H'
fi
printf '\n'
printf '\033[48;2;85;138;255m \033[49m   \033[48;2;96;133;255m \033[49m \033[48;2;107;129;255m   \033[48;2;119;124;255m  \033[49m \033[48;2;130;120;255m \033[49m   \033[48;2;141;115;255m \033[49m  \033[48;2;152;111;255m  \033[48;2;163;107;250m \033[49m    \033[48;2;173;106;234m \033[48;2;183;104;218m   \033[49m  \033[48;2;193;102;202m \033[49m   \033[48;2;214;99;169m \033[49m \033[48;2;214;99;169m \033[48;2;224;97;153m   \033[49m \033[0m\n'
printf '\033[48;2;85;138;255m  \033[49m  \033[48;2;96;133;255m \033[49m \033[48;2;107;129;255m \033[49m     \033[48;2;130;120;255m \033[49m  \033[48;2;141;115;255m \033[49m  \033[48;2;152;111;255m \033[49m   \033[48;2;163;107;250m \033[49m   \033[48;2;173;106;234m \033[49m   \033[48;2;193;102;202m \033[49m \033[48;2;193;102;202m \033[49m   \033[48;2;214;99;169m \033[49m \033[48;2;214;99;169m \033[49m   \033[48;2;230;96;143m \033[0m\n'
printf '\033[48;2;85;138;255m \033[49m \033[48;2;85;138;255m \033[49m \033[48;2;96;133;255m \033[49m \033[48;2;107;129;255m   \033[48;2;119;124;255m \033[49m  \033[48;2;130;120;255m   \033[49m   \033[48;2;152;111;255m \033[49m   \033[48;2;163;107;250m \033[49m   \033[48;2;173;106;234m \033[49m   \033[48;2;193;102;202m \033[49m \033[48;2;193;102;202m \033[49m   \033[48;2;214;99;169m \033[49m \033[48;2;214;99;169m \033[48;2;224;97;153m   \033[49m \033[0m\n'
printf '\033[48;2;85;138;255m \033[49m  \033[48;2;96;133;255m  \033[49m \033[48;2;107;129;255m \033[49m     \033[48;2;130;120;255m \033[49m  \033[48;2;141;115;255m \033[49m  \033[48;2;152;111;255m \033[49m   \033[48;2;163;107;250m \033[49m   \033[48;2;173;106;234m \033[49m   \033[48;2;193;102;202m \033[49m \033[48;2;193;102;202m \033[49m   \033[48;2;214;99;169m \033[49m \033[48;2;214;99;169m \033[49m   \033[48;2;230;96;143m \033[0m\n'
printf '\033[48;2;85;138;255m \033[49m   \033[48;2;96;133;255m \033[49m \033[48;2;107;129;255m   \033[48;2;119;124;255m  \033[49m \033[48;2;130;120;255m \033[49m   \033[48;2;141;115;255m \033[49m  \033[48;2;152;111;255m  \033[48;2;163;107;250m \033[49m    \033[48;2;173;106;234m \033[48;2;183;104;218m   \033[49m   \033[48;2;204;101;185m   \033[49m  \033[48;2;214;99;169m \033[48;2;224;97;153m   \033[49m \033[0m\n'
printf '\n'
if [ -t 1 ]; then
  printf '\033[8r\033[8;1H'
  export NEKO_UI_TOP=8
fi
printf '\033[90m%s\033[0m\n' 'Tips for getting started:'
printf '\033[90m%s\033[0m\n' '1. First run downloads ~9 GB (PyTorch + CUDA libs + model). Later runs start in seconds.'
printf '\033[90m%s\033[0m\n' '2. When you see READY, Neko Dub uses this server automatically.'
printf '\033[90m%s\033[0m\n' '3. Press Ctrl+C to stop. The model unloads after 3 idle min to free VRAM.'
printf '\n'
[ "$(uname -m)" = x86_64 ] || fail "Only x86_64 Linux is supported."
command -v curl >/dev/null || fail "curl is required (e.g. sudo apt install curl)."
nvidia-smi >/dev/null 2>&1 || fail "nvidia-smi did not run. An NVIDIA GPU and a recent NVIDIA driver are required. Nothing has been downloaded."
SELF="${BASH_SOURCE:-}"
[ -f "$SELF" ] || fail "Download NekoDub-GPU.sh first, then run: bash NekoDub-GPU.sh (piping it into bash is not supported)."
DATA="$(cd "$(dirname "$SELF")" && pwd)/neko-omnivoice-data"
mkdir -p "$DATA" && cd "$DATA" || fail "Could not create the neko-omnivoice-data folder next to this file."
export UV_PYTHON_INSTALL_DIR="$DATA/python" UV_CACHE_DIR="$DATA/uv-cache" HF_HOME="$DATA/hf-cache" TMPDIR="$DATA/tmp"
# The voice registry is in TMPDIR and is lost on restart, so start clean.
rm -rf "$DATA/tmp"; mkdir -p "$DATA/tmp"
cat > requirements.txt <<'NEKO_REQ'
accelerate==1.15.0
annotated-doc==0.0.5
annotated-types==0.8.0
anyio==4.16.0
audioread==3.1.0
braceexpand==0.1.7
brotli==1.2.1
certifi==2026.10.23
cffi==2.1.1
charset-normalizer==3.5.2
click==8.5.1
cloudpickle==3.1.2
colorama==0.4.6
decorator==5.3.1
fastapi==0.162.0
filelock==4.2.9
fsspec==2026.9.0
gradio==6.31.0
gradio-client==2.7.2
groovy==0.1.2
h11==0.16.0
hf-gradio==0.4.1
hf-xet==1.8.0
httpcore==1.0.9
httpcore2==2.16.0
httpx==0.28.1
httpx2==2.16.0
huggingface-hub==2.4.0
idna==3.20
jinja2==3.1.6
joblib==1.6.0
lazy-loader==0.6
librosa==0.11.0
llvmlite==0.50.0
markdown-it-py==4.2.0
markupsafe==3.0.4
mdurl==0.1.2
mpmath==1.3.0
msgpack==1.2.4
narwhals==2.28.0
networkx==3.6.1
numba==0.68.1
numpy==2.4.6
omnivoice==0.2.1
opentelemetry-api==1.45.2
orjson==3.13.0
packaging==26.3
pandas==3.1.1
pillow==12.3.0
platformdirs==4.14.5
pooch==1.9.0
protobuf==7.36.2
psutil==7.2.2
pycparser==3.12
pydantic==2.14.3
pydantic-core==2.50.4
pydub==0.25.1
pygments==2.21.0
python-dateutil==4.0.0
python-multipart==0.0.33
pytz==2026.6
pyyaml==6.0.3
regex==2026.11.7
requests==2.34.2
rich==15.0.0
safehttpx==0.1.7
safetensors==0.8.0
scikit-learn==1.9.1
scipy==1.17.1
semantic-version==2.10.0
setuptools==81.0.0
shellingham==1.5.4
soundfile==0.14.0
soxr==1.1.0
starlette==1.10.0
sympy==1.14.0
tensorboardx==2.6.5
threadpoolctl==3.8.0
tokenizers==0.23.3
tomlkit==0.14.0
torch==2.11.0+cu128
torchaudio==2.11.0+cu128
tqdm==4.70.1
transformers==5.22.2
truststore==0.10.5
typer==0.28.0
typing-extensions==4.16.0
typing-inspection==0.4.5
tzdata==2026.6
urllib3==2.8.0
uvicorn==0.54.0
webdataset==1.0.2
NEKO_REQ
if [ ! -f installed.ok ]; then
  # [1/4] to [3/4] run only when the install has not finished before.
  if [ ! -x ./uv ]; then
    printf '\n\033[36m[1/4] Downloading uv (~15 MB)\033[0m\n'
    printf '\033[2m%s\033[0m\n' 'Do not close this terminal.'
    { curl -L --fail --progress-bar -o uv.tar.gz https://github.com/astral-sh/uv/releases/download/0.13.7/uv-x86_64-unknown-linux-gnu.tar.gz \
        && tar -xzf uv.tar.gz --strip-components=1 && rm -f uv.tar.gz && [ -x ./uv ]; } || fail "Failed to download uv."
  fi
  if [ ! -x venv/bin/python ]; then
    printf '\n\033[36m[2/4] Creating a Python 3.11 environment in this folder (~25 MB)\033[0m\n'
    printf '\033[2m%s\033[0m\n' 'This can take a few minutes. Do not close this terminal.'
    ./uv venv --managed-python --python 3.11 --clear venv || fail "Failed to create the Python environment."
  fi
  printf '\n\033[36m[3/4] Installing PyTorch + packages (~5 GB)\033[0m\n'
  printf '\033[2m%s\033[0m\n' 'This can take a few minutes. Do not close this terminal.'
  ./uv pip install --python venv/bin/python --torch-backend cu128 -r requirements.txt || fail "Failed to install Python packages."
  ./uv cache clean
  echo installed > installed.ok
  printf '\033[32mOK Installed. Python packages are ready.\033[0m\n'
fi
cat > omnivoice_local.py <<'NEKO_OMNIVOICE_PY'
# Reference: omnivoice/omnivoice_modal.py is the server this file must stay in step with.
# When the API changes there, port the change here.
"""
OmniVoice local server for Neko Dub (standalone, runs outside Colab and Modal).

Requirements:
  - Python 3.10+
  - NVIDIA GPU + CUDA-enabled PyTorch (install torch/torchaudio from the CUDA index,
    e.g. --index-url https://download.pytorch.org/whl/cu128; a plain `pip install torch`
    is a CPU build and this server exits at startup)

Install (versions pinned in omnivoice/omnivoice-windows-requirements.txt):
  pip install omnivoice==0.2.1 fastapi uvicorn soundfile python-multipart transformers huggingface_hub

First start downloads about 6.5 GB (PyTorch is installed separately; the model and Whisper
download on first run and resume if interrupted).

Run (this machine only):
  python omnivoice_local.py
  -> paste http://127.0.0.1:7860 into Neko Dub -> OmniVoice API URL

Run (reachable from other machines on your LAN; there is no password):
  python omnivoice_local.py --host 0.0.0.0
  -> paste http://<this-machine-LAN-IP>:7860 into Neko Dub -> OmniVoice API URL

Options:
  --concurrency N   simultaneous GPU generations (default 1). Neko Dub sends one item at a time
                    to local servers, so the default matches what it uses.
  --idle-minutes N  unload the model after N idle minutes to free VRAM (default 3; 0 = never).
                    The next request reloads it, normally within the extension's 90 s timeout.

Only extension origins (chrome-extension://, moz-extension://, safari-web-extension://) may
call the API. Requests from ordinary web pages get 403. Requests with no Origin header, such
as opening /api/health in a browser tab, are allowed.

Note: the "Auto setup" toggle doesn't matter for local URLs (the Colab automation
never runs for loopback/LAN servers).

API contract mirrors omnivoice/Colab-block-2.py.
"""

import os, sys

# The output contains non-ASCII text; Windows consoles often start in a legacy code page.
sys.stdout.reconfigure(encoding="utf-8")
sys.stderr.reconfigure(encoding="utf-8")
if sys.platform == "win32":
    os.system("")  # makes the Windows console handle ANSI escape codes

import argparse

parser = argparse.ArgumentParser(description="OmniVoice local TTS server for Neko Dub")
parser.add_argument("--host", default="127.0.0.1", help="bind address (default 127.0.0.1; 0.0.0.0 = LAN)")
parser.add_argument("--port", type=int, default=7860, help="port (default 7860)")
parser.add_argument("--concurrency", type=int, default=1, help="simultaneous GPU generations (default 1)")
parser.add_argument("--idle-minutes", type=int, default=3,
                    help="unload the model after N idle minutes to free VRAM (default 3; 0 = never)")
args = parser.parse_args()
if args.concurrency < 1:
    parser.error("--concurrency must be at least 1")
if args.idle_minutes < 0:
    parser.error("--idle-minutes must be 0 or more")

# Heavy imports only after argparse so --help stays fast.
import re, io, gc, atexit, base64, asyncio, threading, tempfile, time
import torch

if not torch.cuda.is_available():
    print("ERROR: CUDA GPU not available. OmniVoice local server needs an NVIDIA GPU + CUDA PyTorch.\n"
          "  If the computer has an NVIDIA GPU, update the NVIDIA driver to the latest version and restart.",
          file=sys.stderr)
    sys.exit(1)

import soundfile as sf
from fastapi import FastAPI, UploadFile, File, Request, Form
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
from rich import print as rich_print
from rich.console import Group
from rich.live import Live
from rich.panel import Panel
from rich.table import Table

# ── 0. Console UI ──
# The launchers (NekoDub-GPU.bat/.sh) keep their banner fixed at the top and pass the first
# row below it as NEKO_UI_TOP. When the server is ready, everything from that row down (the
# install and download log) is cleared and replaced by one status panel that updates in place.
# Without a terminal the panel is printed once and only errors are printed after it.
IS_TTY = sys.stdout.isatty()
_ui_top = os.environ.get("NEKO_UI_TOP", "")
UI_TOP = int(_ui_top) if IS_TTY and _ui_top.isdigit() and int(_ui_top) > 0 else 0
if UI_TOP:
    # The launcher limited scrolling to the rows below its banner. Give the whole window back
    # on exit, keeping the cursor where it is.
    atexit.register(lambda: (sys.stdout.write("\x1b7\x1b[r\x1b8"), sys.stdout.flush()))

_ui_lock = threading.Lock()
_ui = {"last_seen": None, "ok": 0, "failed": 0, "last_error": None}
_live_on = False   # True while the live panel owns the screen
_loading = False   # True while acquire() reloads the model

def _seen():
    with _ui_lock:
        _ui["last_seen"] = time.strftime("%H:%M:%S")

def _count(ok: bool):
    with _ui_lock:
        _ui["ok" if ok else "failed"] += 1

def _error(message):
    """Record an error for the panel. Without the live panel, print it instead."""
    text = " ".join(str(message).split())[:200]
    with _ui_lock:
        _ui["last_error"] = f'{time.strftime("%H:%M:%S")}  {text}'
    if not _live_on:
        print(f"ERROR: {text}", file=sys.stderr, flush=True)

# ── 1. Load model ──
# Download first with progress bars, so the model never downloads inside a request.
from huggingface_hub import snapshot_download
# Whisper needs only the config files and model.safetensors (290 MB of its 1.17 GB repo).
WHISPER_FILES = ["*.json", "*.txt", "model.safetensors"]

def _in_cache(repo, **kwargs):
    try:
        snapshot_download(repo, local_files_only=True, **kwargs)
        return True
    except Exception:
        return False

# The cache check only picks the message. The real snapshot_download calls below always run,
# so an interrupted download still resumes.
if _in_cache("k2-fsa/OmniVoice") and _in_cache("openai/whisper-base", allow_patterns=WHISPER_FILES):
    print("Checking model files...", flush=True)
else:
    print("[5/5] Downloading model (~3.6 GB, first run only)", flush=True)
snapshot_download("k2-fsa/OmniVoice")
snapshot_download("openai/whisper-base", allow_patterns=WHISPER_FILES)

print("Loading OmniVoice model...", flush=True)
from omnivoice import OmniVoice
_model = OmniVoice.from_pretrained("k2-fsa/OmniVoice", device_map="cuda:0", dtype=torch.float16)

# Later reloads happen under the live panel, so they must not draw progress bars.
try:
    from huggingface_hub.utils import disable_progress_bars
    disable_progress_bars()
    from transformers.utils import logging as _hf_logging
    _hf_logging.disable_progress_bar()
except Exception:
    pass

# ── 1b. Idle VRAM release ──
# After --idle-minutes with no generation, the watcher drops the model and frees VRAM. The next
# generation reloads it inside acquire(). /api/health never calls acquire(), so the extension's
# keep-alive pings do not wake the model.
_model_lock = threading.Lock()
_busy = 0                   # generations in flight; the watcher unloads only when this is 0
_last = time.monotonic()    # last release(); the idle clock

def acquire():
    """Return the loaded model, reloading it if it was idle-unloaded. Runs in a worker thread.
    Every successful call must be paired with release()."""
    global _model, _busy, _loading
    with _model_lock:
        _busy += 1  # MUST be inside the lock (prevents unload/double-load race)
        try:
            if _model is None:
                _loading = True  # the panel shows "Reloading model..."
                # A reload from the disk cache normally finishes inside the extension's 90 s
                # request timeout (tts.ts:28). A cold re-download would not.
                _model = OmniVoice.from_pretrained("k2-fsa/OmniVoice", device_map="cuda:0", dtype=torch.float16)
            return _model
        except BaseException as e:
            _busy -= 1  # a failed acquire holds no slot, so release() must not run for it
            _error(f"Could not reload the model: {e}")
            raise
        finally:
            _loading = False

def release():
    global _busy, _last
    with _model_lock:
        _busy -= 1
        _last = time.monotonic()

async def acquire_async():
    """acquire() in a worker thread, safe against request cancellation. If the request is
    cancelled while acquire() runs (e.g. the client times out during a reload), the thread
    still finishes and takes a slot; release it then, or _busy stays >0 and idle unload
    never fires again."""
    fut = asyncio.ensure_future(asyncio.to_thread(acquire))
    try:
        return await asyncio.shield(fut)
    except asyncio.CancelledError:
        fut.add_done_callback(lambda f: release() if not f.cancelled() and f.exception() is None else None)
        raise

def _idle_watcher():
    global _model
    while True:
        time.sleep(30)
        with _model_lock:
            if (args.idle_minutes > 0 and _model is not None and _busy == 0
                    and time.monotonic() - _last >= args.idle_minutes * 60):
                _model = None
                whisper_cache["pipe"] = None
                gc.collect()
                torch.cuda.empty_cache()  # the panel now shows "Sleeping"

# ── 2. Voice Registry ──
# Replaces the single ref-audio slot. The OmniVoice model is stateless
# (ref_audio + ref_text passed per generate() call), so multi-voice only
# needs a registry lookup here.
DEFAULT_VOICE_ID = '__default__'
# LRU cap on the registry. Entries are temp files only (the model is stateless).
MAX_REGISTERED_VOICES = 128
registry: dict = {}                   # voice_id -> {"path": str, "ref_text": str|None}

# Batch limits. Concurrency is --concurrency (default 1). It is not detected from VRAM.
MAX_BATCH_ITEMS = 24
MAX_BATCH_CONCURRENCY = args.concurrency
generate_semaphore = asyncio.Semaphore(MAX_BATCH_CONCURRENCY)

# Whisper pipeline is built lazily once and reused. HF pipelines are not thread-safe,
# so calls are serialized with a lock and run in a worker thread (not on the event loop).
whisper_cache: dict = {"pipe": None}
_whisper_lock = threading.Lock()

def _transcribe_with_whisper(path):
    """Runs in a worker thread. Raises on failure: returning None here would make generate()
    auto-load whisper-large-v3-turbo, an unannounced large download with no lock."""
    from omnivoice.utils.audio import load_audio  # reads WAV/OGG/MP3 without ffmpeg
    with _whisper_lock:
        asr = whisper_cache["pipe"]
        if asr is None:
            from transformers import pipeline as hf_pipeline
            asr = hf_pipeline('automatic-speech-recognition', model='openai/whisper-base', device='cuda:0', dtype=torch.float16)
            whisper_cache["pipe"] = asr
        audio = load_audio(path, 16000)
        text = asr({"raw": audio[0], "sampling_rate": 16000})['text'].strip()
    return text

def _safe_unlink(path):
    """Windows refuses to delete a file that a thread or antivirus still has open.
    A failed delete leaves a stale temp file; it must not turn into an HTTP 500."""
    try:
        os.unlink(path)
    except OSError:
        pass

def _touch(vid):
    """Mark a voice as most recently used (dict insertion order = LRU order)."""
    registry[vid] = registry.pop(vid)

def _evict_lru_if_needed():
    """Keep registry size bounded; evict the least-recently-used voice (batch use touches it)."""
    while len(registry) > MAX_REGISTERED_VOICES:
        oldest_id = next(iter(registry))
        old = registry.pop(oldest_id)
        _safe_unlink(old['path'])

# ── 3. FastAPI ──
api_app = FastAPI()

# Only extension origins may call the API. CORS alone is not enough: a website's simple
# multipart POST would still reach the server, so the middleware below also rejects them.
ORIGIN_PATTERN = r"^(chrome-extension|moz-extension|safari-web-extension)://.*$"
_ORIGIN_RE = re.compile(ORIGIN_PATTERN)
api_app.add_middleware(
    CORSMiddleware,
    allow_origin_regex=ORIGIN_PATTERN,
    allow_methods=["*"],
    allow_headers=["*"],
    allow_private_network=True,  # the pinned Starlette rejects these preflights with 400 otherwise
)


@api_app.middleware("http")
async def _extension_origin_only(request: Request, call_next):
    origin = request.headers.get("origin")
    if origin is not None and not _ORIGIN_RE.match(origin):
        return JSONResponse({'ok': False, 'error': 'origin not allowed'}, status_code=403)
    _seen()  # the panel shows when Neko Dub last called, health pings included
    return await call_next(request)


@api_app.middleware("http")
async def _private_network_access(request: Request, call_next):
    # Chrome Private Network Access preflight: allow public pages/extension to reach LAN/loopback.
    response = await call_next(request)
    response.headers["Access-Control-Allow-Private-Network"] = "true"
    return response

@api_app.get("/api/health")
async def health():
    # Drop dead entries (temp file gone) before reporting.
    for vid in list(registry.keys()):
        if not os.path.exists(registry[vid]['path']):
            registry.pop(vid, None)
    default_ref = registry.get(DEFAULT_VOICE_ID)
    return {
        "ok": True,
        "voices": list(registry.keys()),
        "ref_audio_set": DEFAULT_VOICE_ID in registry,                # back-compat
        "ref_text": default_ref['ref_text'] if default_ref else None,  # back-compat
        "voice_design": True,
        "max_voices": MAX_REGISTERED_VOICES,
        "max_batch_items": MAX_BATCH_ITEMS,
        "max_batch_concurrency": MAX_BATCH_CONCURRENCY,
    }

@api_app.post("/api/set-ref-audio")
async def set_ref_audio(
    file: UploadFile = File(...),
    ref_text: str = Form(default=''),
    voice_id: str = Form(default=DEFAULT_VOICE_ID),   # NEW — back-compat default
):
    content = await file.read()
    ext = os.path.splitext(file.filename or '.wav')[1] or '.wav'

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
    tmp.write(content); tmp.close()

    # Transcribe before touching the previous voice, so a failed transcription keeps the old one.
    try:
        if ref_text.strip():
            rt = ref_text.strip()
        else:
            await acquire_async()  # hold the model slot while Whisper runs
            try:
                rt = await asyncio.to_thread(_transcribe_with_whisper, tmp.name)
            finally:
                release()
    except Exception as e:
        _safe_unlink(tmp.name)
        _error(f'Could not transcribe the reference audio: {e}')
        return JSONResponse(
            {'ok': False, 'error': f'Could not transcribe the reference audio: {e}'},
            status_code=500,
        )

    # Remove the previous temp file for THIS voice only (don't touch others).
    old = registry.pop(voice_id, None)  # pop so re-insert becomes LRU-newest
    if old:
        _safe_unlink(old['path'])

    registry[voice_id] = {"path": tmp.name, "ref_text": rt}
    _evict_lru_if_needed()

    return {"ok": True, "voice_id": voice_id, "ref_text": rt}

@api_app.post("/api/design-voice")
async def design_voice(request: Request):
    body = await request.json()
    voice_id = (body.get('voice_id') or '').strip()
    instruct = (body.get('instruct') or '').strip()
    if not voice_id or not instruct:
        return JSONResponse({'ok': False, 'error': 'voice_id + instruct required'}, status_code=400)

    seed_text = (body.get('seed_text') or '').strip() or 'Xin chào, đây là đoạn văn bản mẫu.'
    async with generate_semaphore:
        m = await acquire_async()
        try:
            audio = await asyncio.to_thread(m.generate, text=seed_text, instruct=instruct, num_step=16)
        finally:
            release()
    buf = io.BytesIO()
    sf.write(buf, audio[0], 24000, format='OGG')
    raw = buf.getvalue()

    old = registry.pop(voice_id, None)
    if old:
        _safe_unlink(old['path'])

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.ogg')
    tmp.write(raw); tmp.close()

    registry[voice_id] = {"path": tmp.name, "ref_text": seed_text}
    _evict_lru_if_needed()

    return {
        "ok": True,
        "voice_id": voice_id,
        "ref_text": seed_text,
        "audioBase64": base64.b64encode(raw).decode(),
    }

@api_app.post("/api/tts/batch")
async def tts_batch(request: Request):
    body = await request.json()
    items = body.get('items', [])
    if not isinstance(items, list):
        return JSONResponse({'ok': False, 'error': 'items must be a list'}, status_code=400)
    if len(items) > MAX_BATCH_ITEMS:
        return JSONResponse(
            {'ok': False, 'error': f'Too many items. Max {MAX_BATCH_ITEMS}.'},
            status_code=413,
        )

    # Per-request default voice (back-compat: omits → __default__).
    default_voice = body.get('voice_id') or DEFAULT_VOICE_ID

    # Drop dead registry entries before resolving voices.
    for vid in list(registry.keys()):
        if not os.path.exists(registry[vid]['path']):
            registry.pop(vid, None)

    needed = {(it.get('voice_id') or default_voice) for it in items}
    missing = sorted(needed - registry.keys())
    if missing:
        return JSONResponse(
            {'ok': False, 'error': f'Missing ref audio for: {missing}', 'missingVoices': missing},
            status_code=400,
        )

    for vid in needed:
        _touch(vid)  # true LRU: batch use keeps voices alive

    speed = float(body.get('speed', 1.0))
    num_step = int(body.get('num_step', 16))
    request_default_ref_text = body.get('ref_text')  # legacy override for default voice
    started_at = time.perf_counter()

    async def generate_one(index: int, item: dict):
        iid = item.get('id', index)
        text = str(item.get('text', '')).strip()
        if not text:
            return index, {'id': iid, 'audioBase64': None, 'byteLength': 0, 'error': 'Empty text'}

        vid = item.get('voice_id') or default_voice
        ref = registry[vid]
        # Item-level ref_text wins, else request-level (legacy), else registry ref_text.
        item_ref_text = item.get('ref_text') or request_default_ref_text or ref['ref_text']

        async with generate_semaphore:
            item_started_at = time.perf_counter()
            try:
                audio = await asyncio.to_thread(m.generate, text=text, ref_audio=ref['path'], ref_text=item_ref_text, speed=speed, num_step=num_step)
                buf = io.BytesIO()
                sf.write(buf, audio[0], 24000, format='OGG')
                wav_bytes = buf.getvalue()
                b64 = base64.b64encode(wav_bytes).decode()
                elapsed_ms = int((time.perf_counter() - item_started_at) * 1000)
                _count(True)
                return index, {
                    'id': iid,
                    'audioBase64': b64,
                    'byteLength': len(wav_bytes),
                    'error': None,
                    'elapsedMs': elapsed_ms,
                }
            except Exception as e:
                elapsed_ms = int((time.perf_counter() - item_started_at) * 1000)
                _count(False)
                _error(f'Sentence #{iid} failed: {e}')
                return index, {
                    'id': iid,
                    'audioBase64': None,
                    'byteLength': 0,
                    'error': str(e),
                    'elapsedMs': elapsed_ms,
                }

    m = await acquire_async()  # one hold for the whole batch
    try:
        indexed_results = await asyncio.gather(*(generate_one(index, item) for index, item in enumerate(items)))
    finally:
        release()
    results = [result for _, result in sorted(indexed_results, key=lambda pair: pair[0])]
    elapsed_ms = int((time.perf_counter() - started_at) * 1000)
    ok_count = sum(1 for result in results if result.get('audioBase64'))

    return {
        'ok': True,
        'data': results,
        'meta': {
            'requested': len(items),
            'returned': ok_count,
            'elapsedMs': elapsed_ms,
            'maxBatchItems': MAX_BATCH_ITEMS,
            'maxBatchConcurrency': MAX_BATCH_CONCURRENCY,
            'voices': sorted(needed),
        },
    }

# ── 4. Status panel + start server ──
def _status_panel():
    """The one panel shown after startup. Rebuilt on every refresh from the live state."""
    with _ui_lock:
        ui = dict(_ui)
    if _loading:
        status = "[yellow]Reloading model...[/]"
    elif _busy > 0:
        status = "[cyan]Generating...[/]"
    elif _model is None:
        status = "[blue]Sleeping[/] [dim](VRAM freed, wakes on the next dub)[/]"
    elif args.idle_minutes > 0:
        status = f"[green]Ready[/] [dim](model unloads after {args.idle_minutes} idle min)[/]"
    else:
        status = "[green]Ready[/]"
    if ui["last_seen"]:
        client = f"[green]Connected[/] [dim](last request {ui['last_seen']})[/]"
    else:
        client = "[dim]Waiting for Neko Dub...[/]"
    sentences = f"{ui['ok']} dubbed"
    if ui["failed"]:
        sentences += f" [red]· {ui['failed']} failed[/]"

    shown_host = "127.0.0.1" if args.host in ("0.0.0.0", "") else args.host
    grid = Table.grid(padding=(0, 2))
    grid.add_column(style="dim", no_wrap=True)
    grid.add_column()
    grid.add_row("Status", status)
    grid.add_row("Neko Dub", client)
    grid.add_row("Sentences", sentences)
    grid.add_row("Address", f"http://{shown_host}:{args.port}")
    if args.host == "0.0.0.0":
        grid.add_row("Network", f"[yellow]Open to your network, no password: http://<this-machine-LAN-IP>:{args.port}[/]")
    if ui["last_error"]:
        # Error text comes from exceptions; escape it so "[...]" is not read as markup.
        from rich.markup import escape
        grid.add_row("Last error", f"[red]{escape(ui['last_error'])}[/]")
    # The hint is a body line, not a panel subtitle: a subtitle is cut off when the rows are short.
    hint = "[dim]Neko Dub uses this server automatically. Close this window to stop.[/]"
    return Panel(
        Group(grid, "", hint), title="[bold green]READY[/]", title_align="left",
        border_style="green", expand=False, padding=(1, 2),
    )

if __name__ == "__main__":
    if args.idle_minutes > 0:
        threading.Thread(target=_idle_watcher, daemon=True).start()
    serve = lambda: uvicorn.run(api_app, host=args.host, port=args.port, log_level="warning", access_log=False)
    if not IS_TTY:
        rich_print(_status_panel())
        serve()
    else:
        if UI_TOP:
            # Clear the install and download log below the launcher's banner.
            sys.stdout.write(f"\x1b[{UI_TOP};1H\x1b[J")
            sys.stdout.flush()
        # Live also routes stray stdout/stderr lines (warnings, uvicorn errors) above the panel.
        with Live(get_renderable=_status_panel, refresh_per_second=2):
            _live_on = True
            try:
                serve()
            finally:
                _live_on = False
NEKO_OMNIVOICE_PY
printf '\n\033[36m[4/4] Starting the OmniVoice server. First run downloads the model (~3.6 GB).\033[0m\n'
printf '\033[2m%s\033[0m\n' 'Do not close this terminal.'
printf '%s\n' 'When the READY panel appears, Neko Dub uses this server automatically.'
# No exec: the EXIT trap above must still run when the server stops.
venv/bin/python omnivoice_local.py "$@"
exit $?