siftlode/backend/app/plex/subtitles.py

"""WebVTT conversion, HLS-clock time-shift, and the on-disk cache + background extractor for Plex
subtitle streams.

Text subtitles (embedded or external) are turned into WebVTT so the browser <track> can render them;
bitmap subtitles (Blu-ray/DVD) can't be (would need burn-in). Converted VTTs are cached on the shared
Plex on-disk cache root (``cache.CACHE_DIR``).

WHY THERE IS A BACKGROUND EXTRACTOR HERE — measured on prod, 2026-09-10, against a 17.06 GB 1080p
MKV with two embedded SubRip tracks on the 81 TB array:

===========================================================  ========
  full pass, ONE track (what this module used to do)           154 s
  full pass, BOTH text tracks in a SINGLE ffmpeg run           111 s
  `-c:s copy` instead of converting to WebVTT                  143 s
  from 1:00:00 to the end of the file                           66 s
  a 630 s WINDOW at 1:00:00, both tracks                          8 s
===========================================================  ========

The conversion is not the cost (dropping it saved 7%): ffmpeg has to stream every byte of the
container between its seek point and its stop point, and the array delivers ~140 MB/s — so the time
is (bytes read) / 140 MB/s, and the ONLY lever is reading less. Hence:

* extract a WINDOW around the playhead first — the viewer sees subtitles in seconds — then complete
  the whole track in the background so every later seek is instant and the next session is free;
* extract EVERY embedded text track in ONE pass, because the read is shared: picking Hungarian used
  to cost a full read and then switching to English cost another one (the "hektikus" report);
* never do any of it inside a request (it froze the <track> with no feedback AND held an API worker
  thread for the whole read).

`-copyts` CANNOT be used for a window: it makes ffmpeg ignore `-t`/`-to` (measured — a [0, 600 s]
window still took the full 145 s and emitted the whole track), so windows are extracted on a
zero-based clock and shifted back to absolute time here with `shift_vtt`.
"""

import json
import logging
import re
import subprocess
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from sqlalchemy.orm import Session

from app.db import SessionLocal
from app.models import PlexItem
from app.plex import paths as plex_paths
from app.plex.cache import CACHE_DIR
from app.plex.client import PlexClient, PlexError
from app.plex.lookups import by_rating_key

log = logging.getLogger("siftlode.plex")

# Image-based subtitle codecs (Blu-ray/DVD) — can't be turned into text WebVTT (would need burn-in).
BITMAP_SUBS = {
    "pgs",
    "hdmv_pgs_subtitle",
    "dvd_subtitle",
    "dvdsub",
    "vobsub",
    "dvbsub",
    "dvb_subtitle",
}


def srt_to_vtt(text: str) -> str:
    """Minimal SubRip→WebVTT: add the header and switch the ms separator (`,`→`.`). SRT cue-index
    lines are valid in WebVTT, so they're left as-is."""
    body = re.sub(r"(\d\d:\d\d:\d\d),(\d\d\d)", r"\1.\2", text.replace("\r\n", "\n").lstrip(""))
    return "WEBVTT\n\n" + body.strip() + "\n"


def _vtt_ts_to_s(ts: str) -> float:
    parts = ts.split(":")
    h = "0" if len(parts) == 2 else parts[0]
    m, rest = parts[-2], parts[-1]
    sec, _, ms = rest.partition(".")
    return int(h) * 3600 + int(m) * 60 + int(sec) + int(ms or 0) / 1000.0


def _vtt_s_to_ts(x: float) -> str:
    x = max(0.0, x)
    h = int(x // 3600)
    x -= h * 3600
    m = int(x // 60)
    x -= m * 60
    s = int(x)
    ms = int(round((x - s) * 1000))
    if ms >= 1000:
        s += 1
        ms -= 1000
    return f"{h:02d}:{m:02d}:{s:02d}.{ms:03d}"


def shift_vtt(vtt: str, offset: float) -> str:
    """Subtract `offset` seconds from every cue's start/end. WebVTT cues carry ABSOLUTE content times,
    but an HLS remux session's clock is zero-based at the real keyframe (see stream.py media_start_s),
    so a subtitle overlaid on that session must be shifted by that offset to line up with speech.

    Processes cue BLOCKS so a cue that shifts fully into the past (end <= 0) can be DROPPED entirely.
    (Collapsing them to 00:00:00.000-->00:00:00.000 made EVERY past cue count as active at currentTime 0
    — where a resume/seek lands — piling all of them on screen until playback advanced past 0.) A cue
    straddling 0 is clamped to start at 0."""
    if not offset:
        return vtt
    out_blocks: list[str] = []
    for block in re.split(r"\n[ \t]*\n", vtt):
        if "-->" not in block:
            out_blocks.append(block)  # header / NOTE / STYLE — keep verbatim
            continue
        lines = block.split("\n")
        keep = True
        for i, ln in enumerate(lines):
            if "-->" not in ln:
                continue
            left, _, right = ln.partition("-->")
            rparts = right.strip().split(None, 1)
            settings = (" " + rparts[1]) if len(rparts) > 1 else ""
            try:
                st = _vtt_ts_to_s(left.strip()) - offset
                en = _vtt_ts_to_s(rparts[0]) - offset
            except (ValueError, IndexError):
                break  # unparseable timing → leave the block as-is
            if en <= 0:
                keep = False  # whole cue is before the session start → drop the block
            else:
                lines[i] = (
                    f"{_vtt_s_to_ts(st)} --> {_vtt_s_to_ts(en)}{settings}"  # _vtt_s_to_ts clamps st>=0
                )
            break
        if keep:
            out_blocks.append("\n".join(lines))
    return "\n\n".join(out_blocks)


def ffmpeg_to_vtt(src: str, sub_ord: int | None) -> str:
    """Convert one subtitle stream of a file to WebVTT via ffmpeg (for embedded subs, and non-SRT
    external subs written to a temp file). `sub_ord` = ordinal among the file's subtitle streams."""
    cmd = [
        "ffmpeg",
        "-nostdin",
        "-loglevel",
        "error",
        "-i",
        src,
        "-map",
        f"0:s:{sub_ord if sub_ord is not None else 0}",
        "-f",
        "webvtt",
        "pipe:1",
    ]
    r = subprocess.run(cmd, capture_output=True, timeout=EXTRACT_TIMEOUT_S)
    if r.returncode != 0 or not r.stdout:
        raise PlexError(
            (r.stderr or b"").decode("utf-8", "replace")[:200] or "ffmpeg subtitle convert failed"
        )
    return r.stdout.decode("utf-8", "replace")


def cache_read(key: str) -> str | None:
    try:
        p = CACHE_DIR / f"{key}.vtt"
        return p.read_text(encoding="utf-8") if p.exists() else None
    except OSError:
        return None


def cache_write(key: str, vtt: str) -> None:
    try:
        CACHE_DIR.mkdir(parents=True, exist_ok=True)
        tmp = CACHE_DIR / f"{key}.vtt.tmp"
        tmp.write_text(vtt, encoding="utf-8")
        tmp.replace(CACHE_DIR / f"{key}.vtt")  # atomic publish
    except OSError:
        pass


# --- cache identity -----------------------------------------------------------------------------

# The cached VTT holds ABSOLUTE cue times (per rating_key + subtitle ordinal); the per-session shift
# is applied on the way out, so a resume/seek to any offset reuses the same file. Beside it, a small
# JSON sidecar records WHICH absolute-time ranges the file actually covers — a windowed extraction
# publishes a usable but incomplete track, and "incomplete" must never be mistaken for "all there is".
_FULL_END_S = 1_000_000_000.0  # stand-in for "to the end of the file" (JSON can't hold infinity)


def cache_key(rating_key: str, sub_ord: int) -> str:
    return f"sub_{rating_key}_{sub_ord}"


def _vtt_path(key: str) -> Path:
    return CACHE_DIR / f"{key}.vtt"


def _meta_path(key: str) -> Path:
    return CACHE_DIR / f"{key}.json"


def read_meta(key: str) -> tuple[list[list[float]], bool]:
    """Covered absolute-time ranges + whether the track is fully extracted. A VTT with no sidecar
    (written by an older build, which only ever cached complete tracks) counts as complete."""
    try:
        raw = json.loads(_meta_path(key).read_text(encoding="utf-8"))
        return [[float(a), float(b)] for a, b in raw.get("ranges", [])], bool(raw.get("complete"))
    except (OSError, ValueError, TypeError):
        return ([[0.0, _FULL_END_S]], True) if _vtt_path(key).exists() else ([], False)


def _write_meta(key: str, ranges: list[list[float]], complete: bool) -> None:
    try:
        CACHE_DIR.mkdir(parents=True, exist_ok=True)
        tmp = CACHE_DIR / f"{key}.json.tmp"
        tmp.write_text(json.dumps({"ranges": ranges, "complete": complete}), encoding="utf-8")
        tmp.replace(_meta_path(key))
    except OSError:
        pass


def generation(key: str) -> int:
    """Monotonic id of the cached file's current content. It goes into the <track> URL so that when a
    window is superseded by a wider one the browser refetches instead of keeping the short VTT."""
    try:
        return _vtt_path(key).stat().st_mtime_ns
    except OSError:
        return 0


def merge_ranges(ranges: list[list[float]]) -> list[list[float]]:
    """Sort + coalesce touching/overlapping covered ranges."""
    out: list[list[float]] = []
    for a, b in sorted((r for r in ranges if r[1] > r[0]), key=lambda r: r[0]):
        if out and a <= out[-1][1]:
            out[-1][1] = max(out[-1][1], b)
        else:
            out.append([a, b])
    return out


def covers(ranges: list[list[float]], at: float) -> bool:
    return any(a <= at <= b for a, b in ranges)


def merge_vtt(*sources: str) -> str:
    """Union of the cue blocks of several absolute-time VTTs, ordered by start time.

    Two extraction passes over the same file emit identical cues for the region they share, so
    de-duplication keys off the TIMING line — the one part of a cue that identifies it uniquely and
    that neither pass can renumber (ffmpeg's webvtt muxer writes no cue identifiers)."""
    seen: set[str] = set()
    cues: list[tuple[float, str]] = []
    for src in sources:
        for block in re.split(r"\n[ \t]*\n", src or ""):
            b = block.strip()
            if not b or "-->" not in b:
                continue  # header / NOTE / STYLE — the merged file gets its own single header
            timing = next(ln for ln in b.split("\n") if "-->" in ln)
            if timing in seen:
                continue
            seen.add(timing)
            try:
                start = _vtt_ts_to_s(timing.partition("-->")[0].strip())
            except (ValueError, IndexError):
                continue
            cues.append((start, b))
    cues.sort(key=lambda c: c[0])
    return "WEBVTT\n\n" + "\n\n".join(b for _, b in cues) + "\n"


def publish(key: str, vtt: str, covered: tuple[float, float], complete: bool) -> None:
    """Merge a freshly extracted region into the cached track and widen its covered ranges."""
    old = cache_read(key)
    old_ranges, was_complete = read_meta(key)
    cache_write(key, merge_vtt(old or "", vtt))
    _write_meta(key, merge_ranges(old_ranges + [list(covered)]), complete or was_complete)


# --- background extraction ----------------------------------------------------------------------

# How much of the track to extract before the viewer sees anything. WINDOW_PRE_S starts a little
# before the playhead so a cue that is already on screen is not missed; WINDOW_LEN_S is the span
# published as the first, provisional VTT (~4 s of work on the measured file). It only has to hold
# until the full pass lands behind it — 111 s for that file — so it is deliberately much longer than
# any realistic full pass.
WINDOW_PRE_S = 30.0
WINDOW_LEN_S = 300.0
# Ceiling on one ffmpeg run so a wedged extraction can't hold a pool thread forever (the old
# in-request extraction had no timeout at all).
EXTRACT_TIMEOUT_S = 900.0
# A failed track is remembered briefly so a 1 s status poll doesn't respawn ffmpeg every second, and
# forgotten again soon after so fixing the file (or Plex coming back) doesn't need a restart.
FAIL_TTL_S = 60.0
# Two items' full passes plus a window for a third: enough that a seek into an un-extracted region is
# served while a full pass runs, bounded so a burst can't thrash the array.
_POOL = ThreadPoolExecutor(max_workers=3, thread_name_prefix="plex-subs")

_LOCK = threading.Lock()
_INFLIGHT: set[tuple] = set()
_FAILURES: dict[str, tuple[str, float]] = {}


class SubtitleUnavailable(Exception):
    """A track that cannot be served as text, carrying the stable code the route maps to a status."""

    def __init__(self, code: str):
        super().__init__(code)
        self.code = code


# code -> HTTP status. An extraction/upstream failure is 422 with a structured detail (502 is
# reserved for the raw byte proxies); the rest are genuine "no such thing" / "not text".
UNAVAILABLE_STATUS = {
    "no_item": 404,
    "no_track": 404,
    "no_file": 404,
    "not_fetchable": 404,
    "bitmap": 415,
    "extract_failed": 422,
    "plex_error": 422,
}


def reset_state() -> None:
    """Drop the in-flight/failure registries (tests; nothing in the app needs it)."""
    with _LOCK:
        _INFLIGHT.clear()
        _FAILURES.clear()


def _fail(key: str, code: str) -> None:
    with _LOCK:
        _FAILURES[key] = (code, time.monotonic())


def _recent_failure(key: str) -> str | None:
    with _LOCK:
        hit = _FAILURES.get(key)
        if hit and time.monotonic() - hit[1] < FAIL_TTL_S:
            return hit[0]
        _FAILURES.pop(key, None)
        return None


def ensure(rating_key: str, sub_ord: int, at: float = 0.0) -> dict:
    """State of one subtitle track, scheduling whatever extraction is still missing. NEVER blocks.

    ``ready`` = fully extracted · ``partial`` = usable, covers ``at``, still growing · ``preparing``
    = nothing to show for ``at`` yet · ``failed`` = ``code`` says why. ``gen`` changes whenever the
    cached file grows, so the client can refetch a window that has been superseded."""
    key = cache_key(rating_key, sub_ord)
    ranges, complete = read_meta(key)
    have = cache_read(key) is not None
    covered = have and covers(ranges, at)
    if complete and have:
        return {"state": "ready", "code": None, "gen": generation(key)}
    code = _recent_failure(key)
    if code and not covered:
        return {"state": "failed", "code": code, "gen": generation(key)}
    if not code:
        _schedule(rating_key, sub_ord, at, want_window=not covered)
    return {"state": "partial" if covered else "preparing", "code": None, "gen": generation(key)}


def _schedule(rating_key: str, sub_ord: int, at: float, want_window: bool) -> None:
    """Kick off the jobs this track still needs: the window that makes it watchable now, and the one
    full pass that makes it complete. Each is deduplicated by identity, so a 1 s status poll (or a
    second viewer) joins the running job instead of starting another ffmpeg."""
    jobs: list[tuple[tuple, float | None]] = []
    if want_window:
        w_start = max(0.0, at - WINDOW_PRE_S)
        jobs.append((("win", rating_key, round(w_start)), w_start))
    jobs.append((("full", rating_key), None))
    for ident, window_start in jobs:
        with _LOCK:
            if ident in _INFLIGHT:
                continue
            _INFLIGHT.add(ident)
        _POOL.submit(_run, ident, rating_key, sub_ord, window_start)


def _run(ident: tuple, rating_key: str, sub_ord: int, window_start: float | None) -> None:
    try:
        _extract_and_publish(rating_key, sub_ord, window_start)
    except SubtitleUnavailable as e:
        _fail(cache_key(rating_key, sub_ord), e.code)
    except PlexError as e:
        log.warning("subtitle extraction failed for %s/%s: %s", rating_key, sub_ord, e)
        _fail(cache_key(rating_key, sub_ord), "plex_error")
    except (
        Exception
    ):  # a worker thread has nobody to raise to — record it, let the next poll see it
        log.exception("subtitle extraction crashed for %s/%s", rating_key, sub_ord)
        _fail(cache_key(rating_key, sub_ord), "extract_failed")
    finally:
        with _LOCK:
            _INFLIGHT.discard(ident)


def _text_streams(db: Session, rating_key: str) -> tuple[PlexItem, list[dict]]:
    """The item plus its subtitle streams, in the ordinal order the API and the UI both use."""
    it = by_rating_key(db, PlexItem, rating_key)
    if it is None:
        raise SubtitleUnavailable("no_item")
    with PlexClient(db) as plex:
        part = ((plex.metadata(it.rating_key) or {}).get("Media") or [{}])[0].get("Part") or [{}]
        return it, [s for s in (part[0].get("Stream") or []) if s.get("streamType") == 3]


def _is_bitmap(stream: dict) -> bool:
    return (stream.get("codec") or stream.get("format") or "").lower() in BITMAP_SUBS


def _extract_and_publish(rating_key: str, sub_ord: int, window_start: float | None) -> None:
    """One extraction job. A sidecar track is fetched whole from Plex (cheap, no ffmpeg); an embedded
    one is extracted together with EVERY other embedded text track in the file, because they all ride
    on the same single read — the whole point of the measurement at the top of this module."""
    with SessionLocal() as db:
        it, subs = _text_streams(db, rating_key)
        if sub_ord < 0 or sub_ord >= len(subs):
            raise SubtitleUnavailable("no_track")
        want = subs[sub_ord]
        if _is_bitmap(want):
            raise SubtitleUnavailable("bitmap")
        if want.get("key"):  # external sidecar → one fetch, always the complete track
            if window_start is not None:
                return  # nothing to window — the full job serves it in a single fetch
            vtt = _sidecar_vtt(db, want)
            publish(cache_key(rating_key, sub_ord), vtt, (0.0, _FULL_END_S), True)
            return
        if want.get("index") is None:
            raise SubtitleUnavailable("not_fetchable")
        src = plex_paths.local_media_path(db, it.file_path)
        if src is None:
            raise SubtitleUnavailable("no_file")
        # Ordinal among the file's EMBEDDED subtitle streams — what `-map 0:s:N` counts. Sidecars sit
        # in the same list but not in the file, so they must not shift the mapping.
        targets = {
            sum(1 for x in subs[:i] if x.get("index") is not None): i
            for i, s in enumerate(subs)
            if s.get("index") is not None and not _is_bitmap(s)
        }
    start = window_start or 0.0
    length = (WINDOW_LEN_S + WINDOW_PRE_S) if window_start is not None else None
    end = _FULL_END_S if length is None else start + length
    for emb_ord, vtt in _ffmpeg_extract(str(src), sorted(targets), window_start, length).items():
        publish(
            cache_key(rating_key, targets[emb_ord]),
            shift_vtt(vtt, -start),  # zero-based window clock → absolute cue times
            (start, end),
            length is None,
        )


def _sidecar_vtt(db: Session, stream: dict) -> str:
    """An external subtitle file, fetched through Plex and converted to WebVTT."""
    with PlexClient(db) as plex:
        raw = plex.raw_get(stream["key"])
    codec = (stream.get("codec") or stream.get("format") or "").lower()
    if codec in ("srt", "subrip", ""):
        return srt_to_vtt(raw.decode("utf-8", "replace"))
    with tempfile.NamedTemporaryFile(suffix="." + (codec or "sub"), delete=False) as fh:
        fh.write(raw)
        tmp = fh.name
    try:
        return ffmpeg_to_vtt(tmp, None)
    finally:
        Path(tmp).unlink(missing_ok=True)


def _ffmpeg_extract(
    src: str, emb_ords: list[int], start: float | None, length: float | None
) -> dict[int, str]:
    """Extract several embedded subtitle streams to WebVTT in ONE ffmpeg run (= one read of the file).

    `start`/`length` bound the read to a window. Without `-copyts` the cues come out on a zero-based
    clock (the caller shifts them back) — which is exactly why `-copyts` is absent: with it ffmpeg
    ignores `-t` and reads the whole container anyway (measured)."""
    if not emb_ords:
        return {}
    with tempfile.TemporaryDirectory(prefix="plexsub-") as tmpdir:
        cmd = ["ffmpeg", "-nostdin", "-loglevel", "error"]
        if start:
            cmd += ["-ss", f"{start:.3f}"]
        cmd += ["-i", src]
        outs = {o: Path(tmpdir) / f"{o}.vtt" for o in emb_ords}
        for emb_ord, out in outs.items():
            cmd += ["-map", f"0:s:{emb_ord}"]
            if length is not None:
                cmd += ["-t", f"{length:.3f}"]
            cmd += ["-f", "webvtt", "-y", str(out)]
        r = subprocess.run(cmd, capture_output=True, timeout=EXTRACT_TIMEOUT_S)
        got = {
            o: p.read_text(encoding="utf-8", errors="replace")
            for o, p in outs.items()
            if p.exists()
        }
    if not got:
        raise SubtitleUnavailable("extract_failed")
    if r.returncode != 0:
        # Partial output with a non-zero exit still beats nothing (a later pass widens it) — say so.
        log.warning("ffmpeg subtitle pass exited %s: %s", r.returncode, (r.stderr or b"")[:200])
    return got