siftlode/backend/app/plex/stream.py

"""On-the-fly HLS remux for Plex playback (seek-restart session model).

Playback is from the LOCAL physical file. Browser-compatible files (playable="direct") are served
raw with HTTP range requests. Everything else that is h264 video (playable="remux") is remuxed on
the fly to HLS with **video stream-copy** (cheap, I/O-bound — no video re-encode) and audio to AAC
when needed. HEVC/VP9 (playable="transcode") needs a full re-encode and is deferred to P3.

Seek model (Jellyfin-style): one ffmpeg session per item, started at a given offset via `-ss`
(fast keyframe seek). A seek beyond the generated region restarts the session at the new offset,
so seeking is responsive even on long movies without pre-generating the whole file. ffmpeg's own
HLS muxer does the segmentation (the only reliable way to cut a stream-copy at keyframes).

Sessions live in this (single) API process; a periodic reaper kills idle ones and frees the temp
segments. The remux is CPU-light (video copy + a tiny audio transcode), so it runs fine on the
CPU-only prod host; only P3's full transcode is CPU-heavy.
"""
import logging
import shutil
import subprocess
import threading
import time
from pathlib import Path

from sqlalchemy.orm import Session as DbSession

from app.config import settings
from app.models import PlexItem
from app.plex import paths

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

_HLS_ROOT = Path(settings.plex_hls_dir) if settings.plex_hls_dir else Path(settings.download_root) / ".plex-hls"
_SEG_SECONDS = 6
_MAX_SESSIONS = 4  # concurrent remux sessions (safety valve for the CPU-only host)
_SESSION_IDLE_S = 600  # reap a session with no access for this long

_lock = threading.Lock()
_sessions: dict[str, "HlsSession"] = {}


class HlsSession:
    def __init__(self, key: str, directory: Path, proc: subprocess.Popen, start_s: float, entry: str):
        self.key = key
        self.dir = directory
        self.proc = proc
        self.start_s = start_s  # the REQUESTED seek offset
        # The REAL media start: with `-ss X -c:v copy` ffmpeg can only cut at the nearest earlier
        # keyframe K<=X, and hls.js zero-bases the timeline to K (so video.currentTime=0 is K, not X).
        # We measure K from the first segment's absolute PTS (kept absolute by `-copyts`) and report
        # THAT as the session start, so the client's absolute clock (start + currentTime) and the
        # subtitle cue-shift both key off the true content position — no drift. Provisional = X until
        # the first segment is probed.
        self.media_start_s = start_s
        self.entry = entry  # the playlist filename hls.js should load (master.m3u8 with subs, else index.m3u8)
        self.last_access = time.time()


def _seek_opts(start_s: float) -> list[str]:
    # -noaccurate_seek: with `-ss` before -i, video (stream-copy) backs up to the keyframe K<=X,
    # but the re-encoded audio is normally trimmed to the exact X — so audio starts (X-K) AFTER
    # video → seconds of silence at the start of every seek/audio-switch. noaccurate_seek makes
    # audio ALSO start at K, aligned with the video (K is our real clock anchor anyway).
    return ["-noaccurate_seek", "-ss", f"{start_s:.3f}"] if start_s > 0 else []


def _ffmpeg_cmd(src: Path, start_s: float, out_dir: Path, audio_ord: int | None, aoff: float = 0.0) -> list[str]:
    ao = audio_ord if audio_ord is not None else 0
    args = ["ffmpeg", "-nostdin", "-loglevel", "error"]
    # Keep ORIGINAL timestamps on the output so the first segment carries the true absolute PTS of the
    # keyframe ffmpeg actually seeked to. hls.js still zero-bases playback to that PTS; we read it back
    # (see _probe_first_pts) to learn the real content offset. Without -copyts the segment PTS is
    # rewritten toward 0 and the true offset is unknowable, which caused the ~seconds subtitle lead.
    args += _seek_opts(start_s) + ["-copyts", "-i", str(src)]
    if aoff:
        # A/V-sync offset: read the file a SECOND time for audio only, with `-itsoffset` shifting the
        # audio PTS by ±aoff relative to the (input-0) video. Only when the user set a non-zero offset
        # (default single-input path is unchanged, no double read). Maps audio from input 1 then.
        args += ["-itsoffset", f"{aoff:.3f}"] + _seek_opts(start_s) + ["-copyts", "-i", str(src)]
        amap = f"1:a:{ao}?"
    else:
        amap = f"0:a:{ao}?"
    # Video always stream-copied; the selected audio track (default = first) is transcoded to AAC.
    # Subtitles are NOT muxed here — they're served as standalone WebVTT tracks (GET /subtitle) that
    # the browser overlays, so choosing a subtitle no longer restarts this session (and external
    # sidecar subs, which aren't in the file, work — the old `-map 0:s` broke on those).
    args += ["-map", "0:v:0", "-map", amap, "-sn"]
    args += ["-c:v", "copy", "-c:a", "aac", "-ac", "2", "-b:a", "192k"]
    args += [
        "-f", "hls",
        "-hls_time", str(_SEG_SECONDS),
        "-hls_list_size", "0",
        # EVENT (not VOD): ffmpeg writes/appends the playlist AS segments complete, so playback can
        # start immediately from this session's offset. VOD only writes the playlist at the end. The
        # full seekbar comes from our known duration + the seek-restart model, not from the playlist.
        "-hls_playlist_type", "event",
        "-hls_segment_type", "mpegts",
        "-hls_flags", "independent_segments+temp_file",
    ]
    args += ["-hls_segment_filename", str(out_dir / "seg_%d.ts"), str(out_dir / "index.m3u8")]
    return args


def _probe_first_pts(seg: Path) -> float | None:
    """Read the first video packet's absolute PTS from a finished HLS segment (`-copyts` kept it
    absolute). This is the true content position of the segment's first frame — the value hls.js
    zero-bases the timeline to. Returns None on any failure (caller falls back to the requested offset)."""
    try:
        out = subprocess.run(
            ["ffprobe", "-v", "error", "-select_streams", "v:0", "-read_intervals", "%+#1",
             "-show_entries", "packet=pts_time", "-of", "csv=p=0", str(seg)],
            capture_output=True, text=True, timeout=10,
        )
        for line in out.stdout.splitlines():
            tok = line.strip().strip(",")
            if tok:
                return float(tok)
    except (subprocess.SubprocessError, ValueError, OSError):
        pass
    return None


def _kill(s: HlsSession) -> None:
    try:
        s.proc.terminate()
        try:
            s.proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            s.proc.kill()
    except Exception:
        pass
    shutil.rmtree(s.dir, ignore_errors=True)


def _enforce_cap() -> None:
    # Caller holds _lock. Drop the least-recently-accessed sessions over the cap.
    if len(_sessions) <= _MAX_SESSIONS:
        return
    for key, s in sorted(_sessions.items(), key=lambda kv: kv[1].last_access)[: len(_sessions) - _MAX_SESSIONS]:
        _kill(s)
        _sessions.pop(key, None)


def start_session(
    db: DbSession,
    item: PlexItem,
    start_s: float,
    audio_ord: int | None = None,
    aoff: float = 0.0,
) -> HlsSession | None:
    """(Re)start the HLS remux for an item at the given offset, with an optional selected audio track
    (ordinal among audio streams). Subtitles are separate WebVTT tracks (not muxed here). Returns
    None if the local file can't be read."""
    src = paths.local_media_path(db, item.file_path)
    if src is None:
        return None
    key = item.rating_key
    start_s = max(0.0, float(start_s))
    with _lock:
        old = _sessions.pop(key, None)
        if old is not None:
            _kill(old)
        directory = _HLS_ROOT / f"{key}_{int(start_s)}_{audio_ord}_{aoff:+.2f}"
        shutil.rmtree(directory, ignore_errors=True)
        directory.mkdir(parents=True, exist_ok=True)
        proc = subprocess.Popen(
            _ffmpeg_cmd(src, start_s, directory, audio_ord, aoff),
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        s = HlsSession(key, directory, proc, start_s, "index.m3u8")
        _sessions[key] = s
        _enforce_cap()
    # Learn the REAL start offset K from the first segment (outside the lock — this blocks on ffmpeg
    # producing seg_0, and we must not hold up other sessions). The client reads `media_start_s`, so
    # its absolute clock + subtitle shift key off the true keyframe, not the requested offset. On
    # timeout/probe failure we keep the provisional requested offset (old behaviour, small drift).
    if start_s > 0:
        seg0 = directory / "seg_0.ts"
        if wait_for(seg0, timeout=25.0):
            k = _probe_first_pts(seg0)
            if k is not None and k >= 0:
                s.media_start_s = k
    log.info(
        "plex hls session start key=%s req=%.1f real=%.3f audio=%s",
        key, start_s, s.media_start_s, audio_ord,
    )
    return s


def current_session(key: str) -> HlsSession | None:
    with _lock:
        s = _sessions.get(key)
        if s is not None:
            s.last_access = time.time()
        return s


def wait_for(path: Path, timeout: float = 20.0) -> bool:
    """Wait until a session file (playlist / segment) exists and is non-empty. Segments are
    produced ~faster than realtime, so a segment just ahead of playback appears quickly; a segment
    far beyond the generated region won't (the frontend restarts the session at a seek instead)."""
    end = time.time() + timeout
    while time.time() < end:
        try:
            if path.exists() and path.stat().st_size > 0:
                return True
        except OSError:
            pass
        time.sleep(0.15)
    return path.exists()


def reap_idle() -> int:
    now = time.time()
    dropped = 0
    with _lock:
        for key, s in list(_sessions.items()):
            done = s.proc.poll() is not None
            if now - s.last_access > _SESSION_IDLE_S or (done and now - s.last_access > 30):
                _kill(s)
                _sessions.pop(key, None)
                dropped += 1
    return dropped