caption-local/server.py

"""Multiple-stream, loopback-only transcription service. No audio persistence."""
from contextlib import asynccontextmanager
import argparse
import ipaddress
import logging
import os
import asyncio
from pathlib import Path
import threading
import time
import hashlib
import re
import sysconfig
import secrets
import json
from collections import OrderedDict

import numpy as np
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, JSONResponse
from starlette.middleware.trustedhost import TrustedHostMiddleware
from starlette.middleware.cors import CORSMiddleware
from urllib.parse import urlsplit
from starlette.staticfiles import StaticFiles

from rooms import MAX_CAST, MAX_CAST_NAME

ROOT = Path(__file__).resolve().parent
RATE = 16000
MAX_BYTES = RATE * 4 * 12
VERSION = "1.1.0"
DEFAULT_MODEL_REVISION = "536b0662742c02347bc0e980a01041f333bce120"
MODEL_REVISIONS = {"small": DEFAULT_MODEL_REVISION, "base": "ebe41f70d5b6dfa9166e2c581c45c9c0cfc57b66",
                   "medium": "08e178d48790749d25932bbc082711ddcfdfbc4f",
                   "large-v3": "edaa852ec7e145841d8ffdb056a99866b5f0a478"}
_CUDA_DLL_HANDLES = []


def configure_windows_cuda():
    """Make optional NVIDIA wheels in this environment visible to native loaders.

    Only this process is changed; no driver, registry or system PATH edits.
    Keep add_dll_directory handles alive for delayed cuDNN dependency loading.
    """
    if os.name != 'nt' or _CUDA_DLL_HANDLES:
        return
    base = Path(sysconfig.get_path('purelib')) / 'nvidia'
    directories = [base / name / 'bin' for name in ('cublas', 'cuda_nvrtc', 'cudnn')]
    directories = [directory for directory in directories if directory.is_dir()]
    for directory in directories:
        _CUDA_DLL_HANDLES.append(os.add_dll_directory(str(directory)))
    if directories:
        # CTranslate2/cuDNN also use LoadLibrary, which consults process PATH.
        os.environ['PATH'] = os.pathsep.join(map(str, directories)) + os.pathsep + os.environ.get('PATH', '')


class Engine:
    def __init__(self, model, threads=4, offline=False, device="cpu", compute_type="auto", workers=1, beam_size=5):
        if device in ('cuda', 'auto'):
            configure_windows_cuda()
        import ctranslate2
        from faster_whisper import WhisperModel
        requested_device = device
        if device == "auto":
            try:
                device = "cuda" if ctranslate2.get_cuda_device_count() > 0 else "cpu"
            except RuntimeError:
                device = "cpu"
        def load(selected):
            precision = compute_type if compute_type != "auto" else ("float16" if selected == "cuda" else "int8")
            instance = WhisperModel(model, device=selected, compute_type=precision,
                                    cpu_threads=threads, num_workers=workers, local_files_only=offline,
                                    revision=MODEL_REVISIONS.get(model))
            return instance, precision
        try:
            self.model, self.compute_type = load(device)
        except RuntimeError:
            if requested_device != "auto" or device != "cuda" or compute_type != "auto":
                raise
            logging.warning("CUDA model loading failed; auto mode is falling back to CPU int8")
            device = "cpu"
            self.model, self.compute_type = load(device)
        self.beam_size = beam_size
        self.workers = workers
        self.threads = threads
        self.device = device
        self.model_name = Path(model).name
        self.multilingual = self.model.model.is_multilingual
        self.actual_compute_type = getattr(self.model.model, 'compute_type', self.compute_type)
        logging.getLogger("uvicorn.error").info("Inference device: %s; compute type: %s", device, self.compute_type)

    def transcribe(self, audio, language, task="transcribe"):
        segments, info = self.model.transcribe(
            audio, language=language, task=task, beam_size=getattr(self, "beam_size", 5), temperature=0, vad_filter=True,
            condition_on_previous_text=False,
            vad_parameters={"min_silence_duration_ms": 500},
        )
        return " ".join(s.text.strip() for s in segments).strip(), info.language


    def window(self, audio, language, context_seconds=0, final=True, words=False):
        """Commit complete words; retain the last second for the next window."""
        duration = len(audio) / RATE
        segments, info = self.model.transcribe(
            audio, language=language, beam_size=getattr(self, "beam_size", 5), temperature=0, vad_filter=True,
            word_timestamps=True, condition_on_previous_text=False,
            hallucination_silence_threshold=1.0,
            vad_parameters={"min_silence_duration_ms": 500},
        )
        cutoff = duration if final else max(context_seconds, duration - 1.0)
        selected = []
        for segment in segments:
            for word in segment.words or []:
                # Midpoint excludes previously committed context, even if a
                # subsequent decode adjusts its timestamps slightly.
                repeated_boundary = (context_seconds > 0 and word.end - word.start <= 0.001
                                     and word.end <= context_seconds + 0.08)
                if not repeated_boundary and (word.start + word.end) / 2 > context_seconds and word.end <= cutoff + 0.001:
                    selected.append(word)
        text = "".join(word.word for word in selected).strip()
        if not any(char.isalnum() for char in text):
            text = ""
        committed = duration if final else (selected[-1].end if selected else cutoff)
        committed = max(context_seconds, min(committed, duration))
        if words:
            timed = [(w.start, w.end, w.word) for w in selected] if text else []
            return text, info.language, committed, timed
        return text, info.language, committed

    def warmup(self):
        # Exercise the actual decoder before readiness, including CUDA libraries.
        segments, _ = self.model.transcribe(np.zeros(RATE, dtype=np.float32),
                                           language="en", beam_size=1, temperature=0,
                                           vad_filter=False, max_new_tokens=1)
        list(segments)


def _loopback_host(host):
    value = (host or '').strip('[]').lower()
    if value in ('localhost', '127.0.0.1', '::1'):
        return True
    try:
        return ipaddress.ip_address(value).is_loopback
    except ValueError:
        return False


def _private_lan_host(host):
    value = (host or '').strip('[]').lower()
    if not value:
        return False
    try:
        address = ipaddress.ip_address(value)
        if address.version == 4:
            return any(address in ipaddress.ip_network(network) for network in
                       ('10.0.0.0/8', '172.16.0.0/12', '192.168.0.0/16'))
        return address in ipaddress.ip_network('fc00::/7')
    except ValueError:
        return value.endswith(('.lan', '.local', '.internal'))


def relay_target_from_env(env):
    """Validated relay address and caption site, or (None, '').

    Requires wss:// except for loopback and explicitly allowed private LAN use.
    """
    allow_insecure_lan = env.get('CAPTION_ALLOW_INSECURE_LAN', '').strip().lower() in ('1', 'true', 'yes', 'on')
    address = env.get('RELAY_ADDRESS', '').strip()
    parsed = urlsplit(address)
    local_ws = parsed.scheme == 'ws' and _loopback_host(parsed.hostname)
    private_ws = parsed.scheme == 'ws' and allow_insecure_lan and _private_lan_host(parsed.hostname)
    if not (parsed.scheme == 'wss' or local_ws or private_ws) or not parsed.hostname:
        if address:
            print('RELAY_ADDRESS must be wss://, ws://localhost, or an allowed private LAN address; '
                  'relay auto-configuration disabled.', flush=True)
        return None, ''
    site = env.get('RELAY_SITE', '').strip()
    if site:
        parsed_site = urlsplit(site)
        local_http = parsed_site.scheme == 'http' and _loopback_host(parsed_site.hostname)
        private_http = parsed_site.scheme == 'http' and allow_insecure_lan and _private_lan_host(parsed_site.hostname)
        if (parsed_site.scheme != 'https' and not local_http and not private_http) or not parsed_site.hostname \
                or parsed_site.username or parsed_site.password or parsed_site.query or parsed_site.fragment:
            print('RELAY_SITE must be https://, http://localhost, or an allowed private LAN address; ignoring it.', flush=True)
            site = ''
        elif not site.endswith('/'):
            site += '/'
    return address, site


def create_app(engine, max_streams=12, queue_timeout=20, registry=None, relay_address=None, relay_site='',
               log_requests=False, allowed_origins=(), admin_key=None, allow_creation=True,
               cookie_secure=False, create_limit=5, create_window=3600, clock=time.time, speakers=None):
    registry = registry or None
    admin_key = admin_key or None
    if admin_key is not None and (len(admin_key) < 24 or not admin_key.isascii() or any(c.isspace() for c in admin_key)):
        raise ValueError('CAPTION_ADMIN_KEY must contain at least 24 ASCII characters without whitespace')
    allowed_origins = tuple(allowed_origins)
    for origin in allowed_origins:
        parsed = urlsplit(origin)
        if (parsed.scheme not in ('https', 'http') or not parsed.hostname or
                parsed.username or parsed.password or parsed.path or parsed.query or parsed.fragment or
                origin != f'{parsed.scheme}://{parsed.netloc}' or '*' in origin or
                any(c.isspace() or ord(c) < 32 for c in origin) or
                (parsed.scheme == 'http' and parsed.hostname not in ('localhost', '127.0.0.1', '::1'))):
            raise ValueError('Allowed origins must be exact HTTPS origins (HTTP only for localhost)')
        parsed.port  # Validate malformed ports before starting.
    pending_tasks = set()
    @asynccontextmanager
    async def lifespan(app):
        yield
        # Finish shielded requests before releasing model resources at shutdown.
        if pending_tasks:
            await asyncio.gather(*list(pending_tasks), return_exceptions=True)
    app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None, lifespan=lifespan)
    app.state.allowed_origins = allowed_origins
    app.state.registry = registry
    app.state.relay_address = relay_address
    app.state.relay_site = relay_site
    app.add_middleware(CORSMiddleware, allow_origins=list(allowed_origins),
                       allow_methods=['GET', 'POST', 'PATCH', 'DELETE'],
                       allow_headers=['Authorization', 'Content-Type', 'X-Caption-Local', 'X-Stream-ID', 'X-Request-ID'],
                       expose_headers=['Retry-After'], max_age=600)
    app.add_middleware(TrustedHostMiddleware, allowed_hosts=["localhost", "127.0.0.1", "[::1]"])
    # One outstanding request per stream bounds both body memory and queue size.
    # asyncio's semaphore queues waiters FIFO; the model owns matching CT2 workers.
    workers = getattr(engine, "workers", 1)
    slots = asyncio.Semaphore(workers)
    streams = OrderedDict()
    running = {}
    def release(stream_id):
        streams[stream_id]["busy"] = False
        streams[stream_id]["seen"] = time.monotonic()
        if streams[stream_id].get("close_when_idle"):
            del streams[stream_id]
        running.pop(stream_id, None)
        state["busy_since"] = min(running.values(), default=None)

    def admit(stream_id):
        now = time.monotonic()
        for key in list(streams):
            if not streams[key]["busy"] and now - streams[key]["seen"] > 120:
                del streams[key]
        if stream_id not in streams:
            if len(streams) >= max_streams:
                state["rejected"] += 1
                raise HTTPException(429, "Stream capacity reached; idle sessions expire after 120 seconds", headers={"Retry-After": "2"})
            streams[stream_id] = {"busy": False, "seen": now, "cache": OrderedDict()}
        session = streams[stream_id]
        if session["busy"]:
            state["rejected"] += 1
            raise HTTPException(429, "Stream busy; retry shortly", headers={"Retry-After": "1"})
        session["busy"] = True
        session["seen"] = now
        return session["cache"]
    state = {"busy_since": None, "completed": 0, "failed": 0, "rejected": 0, "queue_timeouts": 0}
    app.state.inference_status = state

    public_pages = ('/', '/index.html', '/capture-local.html', '/admin.html', '/health')
    session_path = re.compile(r'^/api/rooms/[A-Za-z0-9_-]+/session$')
    cookie_name = 'caption_room'

    def bearer(request):
        value = request.headers.get('authorization', '')
        return value[7:].strip() if value[:7].lower() == 'bearer ' else ''

    def client_ip(request):
        return (request.headers.get('x-real-ip', '').strip() or
                request.headers.get('x-forwarded-for', '').split(',')[0].strip() or
                (request.client.host if request.client else 'unknown'))

    def origin_allowed(request):
        origin = request.headers.get('origin')
        if not origin:
            return True
        expected = (relay_site or str(request.base_url)).rstrip('/')
        return origin == expected or origin in allowed_origins

    def is_public(request):
        path = request.url.path
        if path in public_pages or path.startswith('/static/') or path.startswith('/api/public/'):
            return True
        if request.method == 'POST' and path == '/api/rooms':
            return True
        return request.method == 'POST' and bool(session_path.match(path))

    @app.middleware("http")
    async def response_headers(request, call_next):
        started = time.perf_counter()
        path = request.url.path
        preflight = request.method == 'OPTIONS' and request.headers.get('origin') and request.headers.get('access-control-request-method')
        authorized = True
        protected = (path.startswith(('/api/', '/transcribe', '/streams/', '/v1/', '/relay-config')) or
                     path in ('/api', '/docs', '/redoc', '/openapi.json'))
        if path.startswith('/api/admin'):
            token = bearer(request)
            authorized = bool(admin_key) and secrets.compare_digest(token.encode(), admin_key.encode())
        elif registry is not None and protected and not is_public(request) and not preflight:
            room = registry.room_from_session(request.cookies.get(cookie_name), now=clock())
            if room:
                request.state.room = room
                request.state.room_public_id = room['public_id']
                registry.touch(room['public_id'])
            else:
                authorized = False
        if not authorized:
            response = JSONResponse({'detail': 'Open your room link or provide a valid service key'}, status_code=401,
                                    headers={'WWW-Authenticate': 'Bearer'})
        else:
            response = await call_next(request)
        if log_requests and not is_public(request):
            # Fixed route labels only: no URL queries, tokens, IDs, IPs, audio or text.
            route = request.scope.get('route')
            print(json.dumps({'event':'http_request','method':request.method,
                  'route':getattr(route, 'path', 'unmatched'), 'status':response.status_code,
                  'duration_seconds':round(time.perf_counter()-started,4)}),flush=True)
        if request.headers.get('origin') in allowed_origins:
            response.headers['Access-Control-Allow-Origin'] = request.headers['origin']
            response.headers.add_vary_header('Origin')
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["Referrer-Policy"] = "no-referrer"
        response.headers["Cache-Control"] = "no-store"
        response.headers["Content-Security-Policy"] = (
            "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; "
            "connect-src 'self' wss: ws://localhost:* ws://127.0.0.1:*; object-src 'none'; "
            "base-uri 'none'; frame-ancestors 'self'"
        )
        return response

    def require_registry():
        if registry is None:
            raise HTTPException(404, 'Not found')

    def require_master():
        if registry is None or not registry.master or not relay_address:
            raise HTTPException(404, 'Room links are not configured')

    async def json_body(request, limit=1024):
        raw = await request.body()
        if len(raw) > limit:
            raise HTTPException(413, 'Request body is too large')
        if not raw:
            return {}
        try:
            payload = json.loads(raw)
        except ValueError:
            raise HTTPException(400, 'Invalid JSON body')
        if not isinstance(payload, dict):
            raise HTTPException(400, 'Expected a JSON object')
        return payload

    def portal_base(request):
        base = relay_site or str(request.base_url)
        return base if base.endswith('/') else base + '/'

    def room_links(request, public_id, secret=None):
        base = portal_base(request)
        admin = base + 'room/' + public_id
        if secret:
            admin += '#key=' + secret
        return {'adminUrl': admin, 'audienceUrl': base + 'watch/' + public_id}

    def room_json(room):
        return {'id': room['public_id'], 'label': room['label'], 'created': room['created'],
                'revoked': bool(room['revoked']), 'lastSeen': room['last_seen']}

    def set_session_cookie(response, room):
        token, _ = registry.issue_session(room, now=clock())
        if not token:
            raise HTTPException(404, 'Room links are not configured')
        response.set_cookie(cookie_name, token, max_age=registry.session_ttl, httponly=True,
                            secure=bool(cookie_secure), samesite='lax', path='/')
        return response

    def room_for_path(request, public_id):
        room = getattr(request.state, 'room', None)
        if not room or room['public_id'] != public_id:
            raise HTTPException(401, 'Open this room with its admin link')
        return room

    create_attempts = {}

    def allow_create(ip):
        if create_limit <= 0:
            return True
        now = time.monotonic()
        recent = [t for t in create_attempts.get(ip, ()) if now - t < create_window]
        if len(recent) >= create_limit:
            create_attempts[ip] = recent
            return False
        recent.append(now)
        create_attempts[ip] = recent
        return True

    public_attempts = {}

    def allow_public(ip):
        now = time.monotonic()
        recent = [t for t in public_attempts.get(ip, ()) if now - t < 60]
        if len(recent) >= 240:
            public_attempts[ip] = recent
            return False
        recent.append(now)
        public_attempts[ip] = recent
        return True

    @app.get("/")
    async def index():
        return FileResponse(ROOT / "static/index.html")

    @app.get('/capture-local.html')
    async def capture_local():
        return FileResponse(ROOT / 'static/capture-local.html')

    @app.get('/admin.html')
    async def admin_page():
        return FileResponse(ROOT / 'static/admin.html')

    @app.post('/api/rooms')
    async def create_room(request: Request):
        require_registry()
        require_master()
        if not allow_creation:
            raise HTTPException(403, 'Room creation is disabled on this deployment')
        if not allow_create(client_ip(request)):
            raise HTTPException(429, 'Too many rooms created; try again later', headers={'Retry-After': '3600'})
        payload = await json_body(request)
        label = payload.get('label', '')
        if not isinstance(label, str) or len(label) > 128:
            raise HTTPException(400, 'Label must be at most 128 characters')
        try:
            _, public_id, admin_secret = registry.create_room(label)
        except RuntimeError:
            raise HTTPException(429, 'Room capacity reached')
        room = registry.get_room(public_id)
        response = JSONResponse({'room': room_json(room), **room_links(request, public_id, admin_secret)}, status_code=201)
        set_session_cookie(response, room)
        return response

    @app.post('/api/rooms/{public_id}/session')
    async def room_session(public_id: str, request: Request):
        require_registry()
        require_master()
        secret = bearer(request)
        if not secret:
            payload = await json_body(request)
            secret = payload.get('key') if isinstance(payload.get('key'), str) else ''
        room = registry.verify_admin(public_id, secret)
        if not room:
            raise HTTPException(401, 'This room link is invalid or has been closed')
        response = JSONResponse({'room': room_json(room)})
        return set_session_cookie(response, room)

    @app.get('/api/me')
    async def current_room(request: Request):
        require_registry()
        room = getattr(request.state, 'room', None)
        if not room:
            raise HTTPException(401, 'Open your room link first')
        return {'room': room_json(room), 'audienceUrl': room_links(request, room['public_id'])['audienceUrl']}

    @app.get('/api/rooms/{public_id}')
    async def room_details(public_id: str, request: Request):
        require_registry()
        room = room_for_path(request, public_id)
        return {'room': room_json(room),
                'audienceUrl': room_links(request, public_id)['audienceUrl']}

    @app.patch('/api/rooms/{public_id}')
    async def room_update(public_id: str, request: Request):
        require_registry()
        room_for_path(request, public_id)
        payload = await json_body(request)
        label = payload.get('label')
        if not isinstance(label, str) or len(label) > 128:
            raise HTTPException(400, 'Label must be at most 128 characters')
        if not registry.set_label(public_id, label):
            raise HTTPException(404, 'Room not found')
        return {'room': room_json(registry.get_room(public_id))}

    @app.post('/api/rooms/{public_id}/rotate')
    async def room_rotate(public_id: str, request: Request):
        require_registry()
        room_for_path(request, public_id)
        secret = registry.rotate_admin(public_id)
        if not secret:
            raise HTTPException(404, 'Room not found')
        room = registry.get_room(public_id)
        response = JSONResponse({'room': room_json(room), **room_links(request, public_id, secret)})
        return set_session_cookie(response, room)

    @app.delete('/api/rooms/{public_id}')
    async def room_close(public_id: str, request: Request):
        require_registry()
        room_for_path(request, public_id)
        if not registry.revoke_room(public_id):
            raise HTTPException(404, 'Room not found')
        if speakers:
            speakers.forget(public_id)
        response = JSONResponse({'closed': public_id})
        response.delete_cookie(cookie_name, path='/')
        return response

    def speaker_state(public_id):
        live = speakers.snapshot(public_id) if speakers else {'assignments': {}, 'detected': []}
        return {'enabled': speakers is not None, 'cast': registry.get_cast(public_id), **live}

    def clean_name(value, what):
        if not isinstance(value, str) or len(value.strip()) > MAX_CAST_NAME or any(ord(c) < 32 for c in value):
            raise HTTPException(400, f'{what} must be text of at most {MAX_CAST_NAME} characters')
        return value.strip()

    @app.get('/api/rooms/{public_id}/speakers')
    async def room_speakers(public_id: str, request: Request):
        require_registry()
        room_for_path(request, public_id)
        return speaker_state(public_id)

    @app.patch('/api/rooms/{public_id}/speakers')
    async def room_speakers_update(public_id: str, request: Request):
        require_registry()
        room_for_path(request, public_id)
        payload = await json_body(request, limit=8192)
        if 'cast' in payload:
            cast = payload['cast']
            if not isinstance(cast, list) or len(cast) > MAX_CAST:
                raise HTTPException(400, f'Cast must be a list of at most {MAX_CAST} names')
            names = []
            for name in (clean_name(n, 'Cast names') for n in cast):
                if name and name not in names:
                    names.append(name)
            registry.set_cast(public_id, names)
        if 'assignments' in payload:
            assignments = payload['assignments']
            if not isinstance(assignments, dict) or len(assignments) > 64 or \
                    not all(re.fullmatch(r'S[1-9][0-9]?', key) for key in assignments):
                raise HTTPException(400, 'Assignments map speaker IDs such as S1 to names')
            if not speakers:
                raise HTTPException(409, 'Speaker labels are not enabled on this service')
            speakers.assign(public_id, {k: clean_name(v, 'Speaker names') for k, v in assignments.items()})
        return speaker_state(public_id)

    @app.post('/api/rooms/{public_id}/speakers/reset')
    async def room_speakers_reset(public_id: str, request: Request):
        require_registry()
        room_for_path(request, public_id)
        if speakers:
            speakers.reset(public_id)
        return speaker_state(public_id)

    @app.get('/api/rooms/{public_id}/config')
    async def room_config(public_id: str, request: Request):
        require_registry()
        require_master()
        room_for_path(request, public_id)
        config = registry.relay_config(public_id, relay_address, relay_site, now=clock())
        if config is None:
            raise HTTPException(404, 'Room links are not configured')
        return config

    @app.get('/api/public/rooms/{public_id}')
    async def public_room_config(public_id: str, request: Request):
        require_registry()
        require_master()
        if not allow_public(client_ip(request)):
            raise HTTPException(429, 'Too many requests', headers={'Retry-After': '60'})
        config = registry.public_config(public_id, relay_address, relay_site, now=clock())
        if config is None:
            raise HTTPException(404, 'Room not found')
        return config

    @app.get('/api/admin/rooms')
    async def admin_list_rooms():
        require_registry()
        return {'rooms': registry.list_rooms()}

    @app.delete('/api/admin/rooms/{public_id}')
    async def admin_close_room(public_id: str):
        require_registry()
        if not registry.revoke_room(public_id):
            raise HTTPException(404, 'Room not found or already closed')
        return {'closed': public_id}

    @app.post('/api/admin/rooms/{public_id}/rotate')
    async def admin_rotate_room(public_id: str, request: Request):
        require_registry()
        secret = registry.rotate_admin(public_id)
        if not secret:
            raise HTTPException(404, 'Room not found or already closed')
        return {'room': room_json(registry.get_room(public_id)),
                **room_links(request, public_id, secret)}

    @app.get("/relay-config")
    async def relay_config_route(request: Request):
        require_registry()
        require_master()
        room = getattr(request.state, 'room', None)
        if not room:
            raise HTTPException(401, 'Open your room link first')
        config = registry.relay_config(room['public_id'], relay_address, relay_site, now=clock())
        if config is None:
            raise HTTPException(404, "Relay not configured")
        return config

    @app.get("/health")
    async def health():
        from faster_whisper.tokenizer import _LANGUAGE_CODES
        age = time.monotonic() - state["busy_since"] if state["busy_since"] is not None else 0
        return {"ready": age < 90, "version": VERSION, "sample_rate": RATE, "max_seconds": 12,
                "rooms_enabled": registry is not None, "speakers": speakers is not None,
                "max_streams": max_streams, "workers": workers, "beam_size": getattr(engine, "beam_size", 5),
                "threads_per_worker": getattr(engine, "threads", None),
                "sessions": len(streams), "running": len(running),
                "pending": sum(v["busy"] for v in streams.values()) - len(running),
                "queue_timeout": queue_timeout, "rejected": state["rejected"],
                "queue_timeouts": state["queue_timeouts"],
                "multilingual": engine.multilingual,
                "model": getattr(engine, "model_name", "unknown"),
                "languages": sorted(_LANGUAGE_CODES) if engine.multilingual else ["en"],
                "device": engine.device, "compute_type": engine.compute_type,
                "actual_compute_type": getattr(engine, 'actual_compute_type', engine.compute_type),
                "busy": state["busy_since"] is not None, "completed": state["completed"], "failed": state["failed"],
                "modes": ["transcribe", "translate", "both"] if engine.multilingual else ["transcribe"],
                "translation_target": "en" if engine.multilingual else None}

    @app.delete("/streams/{stream_id}")
    async def close_stream(stream_id: str, request: Request):
        if request.headers.get("x-caption-local") != "1" or not origin_allowed(request):
            raise HTTPException(403, "Use the local capture page")
        session = streams.get(stream_id)
        if session and session["busy"]:
            raise HTTPException(409, "Finish pending audio before closing a stream")
        streams.pop(stream_id, None)
        return {"closed": stream_id}

    @app.post("/transcribe")
    async def transcribe(request: Request):
        if request.headers.get("x-caption-local") != "1" or not origin_allowed(request):
            raise HTTPException(403, "Use the local capture page")
        if request.headers.get("content-type") != "application/octet-stream":
            raise HTTPException(415, "Expected little-endian float32 mono PCM at 16000 Hz")
        language = request.query_params.get("language", "en")
        from faster_whisper.tokenizer import _LANGUAGE_CODES
        if language != "auto" and language not in _LANGUAGE_CODES:
            raise HTTPException(400, "Unknown language code")
        mode = request.query_params.get("mode", "transcribe")
        if mode not in ("transcribe", "translate", "both"):
            raise HTTPException(400, "Mode must be transcribe, translate, or both")
        if not engine.multilingual and (mode != "transcribe" or language not in ("auto", "en")):
            raise HTTPException(400, "Restart with a multilingual model such as --model small for this language or translation")
        window = request.query_params.get("window", "0") == "1"
        label_speakers = request.query_params.get("speakers", "0") == "1"
        if label_speakers and (not speakers or not window or mode == "translate"):
            raise HTTPException(400, "Speaker labels need a window transcription on a service started with CAPTION_DIARIZATION=1")
        speaker_room = getattr(request.state, "room_public_id", None) or "local"
        final = request.query_params.get("final", "1") == "1"
        try:
            context = float(request.query_params.get("context_seconds", "0"))
            if not np.isfinite(context) or not 0 <= context <= 1:
                raise ValueError()
        except ValueError:
            raise HTTPException(400, "Context must be between 0 and 1 seconds")
        request_id = request.headers.get("x-request-id", "")
        if request_id and not re.fullmatch(r"[A-Za-z0-9_-]{8,80}", request_id):
            raise HTTPException(400, "Invalid request ID")
        stream_id = request.headers.get("x-stream-id", "legacy")
        if stream_id != "legacy" and not re.fullmatch(r"[A-Za-z0-9_-]{8,80}", stream_id):
            raise HTTPException(400, "Invalid stream ID")
        cache = admit(stream_id)
        handed_to_worker = False
        try:
            body = bytearray()
            try:
                async with asyncio.timeout(10):
                    async for part in request.stream():
                        if len(body) + len(part) > MAX_BYTES:
                            raise HTTPException(413, "Audio chunk exceeds 12 seconds")
                        body.extend(part)
            except TimeoutError:
                raise HTTPException(408, "Audio upload timed out")
            if len(body) < 4 * 160 or len(body) % 4:
                raise HTTPException(400, "Invalid PCM length")
            audio = np.frombuffer(body, dtype="<f4").copy()
            if not np.isfinite(audio).all() or np.max(np.abs(audio)) > 1.01:
                raise HTTPException(400, "Invalid PCM samples")
            duration = len(audio) / RATE
            if context >= duration or (window and not final and duration - context < 2):
                raise HTTPException(400, "Window must contain new audio; unfinished windows need at least two seconds")
            signature = hashlib.sha256(body + str((language, mode, window, final, context, label_speakers)).encode()).hexdigest()
            now = time.monotonic()
            while cache and (len(cache) > 64 or now - next(iter(cache.values()))[0] > 120):
                cache.popitem(last=False)
            if request_id and request_id in cache:
                _, previous_signature, result = cache[request_id]
                if signature != previous_signature:
                    raise HTTPException(409, "Request ID was already used for different audio or options")
                return result

            def infer():
                try:
                    start = time.perf_counter()
                    transcript = translation = ""
                    detected = None if language == "auto" else language
                    committed = duration if final else duration - 1
                    segments, diarization_seconds = [], 0
                    if np.max(np.abs(audio)) >= 0.0001:
                        if window and label_speakers:
                            transcript, detected, committed, words = engine.window(audio, detected, context, final, words=True)
                            began = time.perf_counter()
                            segments = speakers.label(speaker_room, audio, words)
                            diarization_seconds = round(time.perf_counter() - began, 3)
                        elif window:
                            transcript, detected, committed = engine.window(audio, detected, context, final)
                        elif mode in ("transcribe", "both"):
                            transcript, detected = engine.transcribe(audio, detected)
                        if mode in ("translate", "both"):
                            if detected == "en" and (window or mode == "both"):
                                translation = transcript
                            elif (not window and mode == "translate") or transcript:
                                selected_audio = audio[round(context * RATE):round(committed * RATE)] if window else audio
                                translation, detected = engine.transcribe(selected_audio, detected, task="translate")
                    if mode == "translate":
                        transcript = ""
                    result = {"text": translation if mode == "translate" else transcript,
                              "transcript": transcript, "translation": translation,
                              "language": detected, "output_language": "en" if mode == "translate" else detected,
                              "translation_language": "en" if mode in ("translate", "both") else None,
                              "mode": mode, "audio_seconds": duration, "committed_seconds": committed,
                              "request_id": request_id, "stream_id": stream_id,
                              **({"segments": segments, "diarization_seconds": diarization_seconds} if label_speakers else {}),
                              "inference_seconds": round(time.perf_counter() - start, 3)}
                    return result
                except Exception:
                    logging.exception("Transcription worker failed")
                    raise

            async def scheduled():
                queued = time.monotonic()
                try:
                    try:
                        await asyncio.wait_for(slots.acquire(), timeout=queue_timeout)
                    except TimeoutError:
                        state["queue_timeouts"] += 1
                        raise HTTPException(429, "Inference queue deadline reached; retry pending audio", headers={"Retry-After": "1"})
                    try:
                        running[stream_id] = time.monotonic()
                        state["busy_since"] = min(running.values())
                        wait = running[stream_id] - queued
                        try:
                            result = await asyncio.to_thread(infer)
                        except Exception:
                            state["failed"] += 1
                            raise
                        result["queue_seconds"] = round(wait, 3)
                        if request_id:
                            cache[request_id] = (time.monotonic(), signature, result)
                            while len(cache) > 64:
                                cache.popitem(last=False)
                        state["completed"] += 1
                        return result
                    finally:
                        slots.release()
                finally:
                    release(stream_id)

            task = asyncio.create_task(scheduled())
            pending_tasks.add(task)
            task.add_done_callback(pending_tasks.discard)
            # Consume late errors if the HTTP client disconnects while inference
            # finishes. The result remains cached for its request ID.
            task.add_done_callback(lambda done: done.exception() if not done.cancelled() else None)
            handed_to_worker = True
            try:
                return await asyncio.shield(task)
            except HTTPException:
                raise
            except Exception as exc:
                raise HTTPException(500, "Transcription failed; pending audio can be retried") from exc
        finally:
            if not handed_to_worker:
                release(stream_id)

    from api_compat import register_audio_api
    def close_ephemeral(stream_id):
        session = streams.get(stream_id)
        if session:
            if session['busy']:
                # A disconnected WAV client has no stream ID to reuse. Keep
                # admission locked until its shielded worker finishes, then
                # release the otherwise abandoned session immediately.
                session['close_when_idle'] = True
            else:
                streams.pop(stream_id, None)
    register_audio_api(app, transcribe, close_ephemeral, max_streams, getattr(engine, 'model_name', 'local'))
    app.mount("/static", StaticFiles(directory=ROOT / "static"), name="static")
    return app


def start_watchdog(state, timeout=90, interval=5):
    def watch():
        while True:
            time.sleep(interval)
            since = state["busy_since"]
            if since is not None and time.monotonic() - since > timeout:
                logging.critical("Inference exceeded its deadline; exiting for the service supervisor to restart")
                os._exit(70)
    thread = threading.Thread(target=watch, name="inference-watchdog", daemon=True)
    thread.start()
    return thread


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default=os.environ.get("CAPTION_MODEL", "small"), help="Model name or downloaded directory")
    parser.add_argument("--beam-size", type=int, choices=range(1,6), default=int(os.environ.get("CAPTION_BEAM_SIZE", "5")))
    parser.add_argument("--workers", type=int, default=int(os.environ.get("CAPTION_WORKERS", "1")))
    parser.add_argument("--max-streams", type=int, default=int(os.environ.get("CAPTION_MAX_STREAMS", "12")))
    parser.add_argument("--threads", type=int, default=int(os.environ.get("CAPTION_THREADS", "4")))
    parser.add_argument("--port", type=int, default=8765)
    parser.add_argument("--offline", action="store_true", default=os.environ.get("CAPTION_OFFLINE") == "1")
    parser.add_argument("--log-requests", action="store_true", default=os.environ.get("CAPTION_LOG_REQUESTS") == "1",
                        help="Write request status/duration JSON to stdout; exclude captions and credentials")
    parser.add_argument("--host", choices=["127.0.0.1", "0.0.0.0"], default="127.0.0.1",
                        help="Use 0.0.0.0 only inside a container with a loopback host port mapping")
    parser.add_argument("--device", choices=["cpu", "cuda", "auto"], default=os.environ.get("CAPTION_DEVICE", "cpu"))
    parser.add_argument("--compute-type", choices=["auto", "int8", "float32", "float16", "int8_float16"],
                        default=os.environ.get("CAPTION_COMPUTE_TYPE", "auto"))
    args = parser.parse_args()
    if args.threads < 1 or not 1 <= args.beam_size <= 5 or not 1 <= args.workers <= 8 or not 1 <= args.max_streams <= 64 or not 1 <= args.port <= 65535:
        parser.error("Threads must be positive; beam size 1–5; workers 1–8; streams 1–64; port 1–65535")
    print("Loading speech model; first run downloads weights. Audio stays on this host.", flush=True)
    engine = Engine(args.model, args.threads, args.offline, args.device, args.compute_type, args.workers, args.beam_size)
    engine.warmup()
    print(f"Inference ready: {engine.device} / {engine.compute_type}", flush=True)
    speakers = None
    if os.environ.get("CAPTION_DIARIZATION") == "1":
        from speakers import OnnxBackend, SpeakerService
        speaker_models = os.environ.get("HF_HOME") or str(ROOT / "models")
        speakers = SpeakerService(lambda: OnnxBackend(speaker_models, min(args.threads, 2), args.offline),
                                  threshold=float(os.environ.get("CAPTION_SPEAKER_THRESHOLD", "0.5")),
                                  max_speakers=int(os.environ.get("CAPTION_MAX_SPEAKERS", "12")))
        speakers.load()
    import uvicorn
    from rooms import RoomRegistry, decode_master
    relay_address, relay_site = relay_target_from_env(os.environ)
    master = decode_master(os.environ.get('CAPTION_RELAY_MASTER', ''))
    admin_key = os.environ.get('CAPTION_ADMIN_KEY', '').strip() or None
    rooms_db = os.environ.get('CAPTION_ROOMS_DB', '').strip() or os.environ.get('CAPTION_KEYS_DB', '').strip()
    registry = None
    if rooms_db or master or admin_key:
        registry = RoomRegistry(
            rooms_db or str(ROOT / 'rooms.db'), master=master,
            token_ttl=int(os.environ.get('CAPTION_TOKEN_TTL', str(60 * 60))),
            session_ttl=int(os.environ.get('CAPTION_SESSION_TTL', str(30 * 24 * 3600))),
            max_rooms=int(os.environ.get('CAPTION_MAX_ROOMS', '400')))
    allow_creation = os.environ.get('CAPTION_ROOM_CREATION', 'open').strip().lower() not in ('off', '0', 'false', 'no')
    cookie_setting = os.environ.get('CAPTION_COOKIE_SECURE', 'auto').strip().lower()
    cookie_secure = cookie_setting == '1' if cookie_setting in ('0', '1') else relay_site.startswith('https://')
    app = create_app(engine, args.max_streams, registry=registry, relay_address=relay_address, relay_site=relay_site,
                     log_requests=args.log_requests, allow_creation=allow_creation, cookie_secure=cookie_secure,
                     create_limit=int(os.environ.get('CAPTION_CREATE_LIMIT', '5')),
                     create_window=int(os.environ.get('CAPTION_CREATE_WINDOW', str(3600))),
                     allowed_origins=tuple(s.strip() for s in os.environ.get('CAPTION_ALLOWED_ORIGINS', '').split(',') if s.strip()),
                     admin_key=admin_key, speakers=speakers)
    start_watchdog(app.state.inference_status)
    # Browsers can retain six HTTP/1.1 connections per origin while loading assets.
    # This transport bound includes idle sockets; inference admission stays separate.
    uvicorn.run(app, host=args.host, port=args.port, access_log=False, limit_concurrency=max(64, args.max_streams * 6 + 16), timeout_keep_alive=5, timeout_graceful_shutdown=90)