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
        self.entry = entry  # the playlist filename hls.js should load (master.m3u8 with subs, else index.m3u8)
        self.last_access = time.time()


def _ffmpeg_cmd(
    src: Path, start_s: float, out_dir: Path, audio_ord: int | None, sub_ord: int | None
) -> list[str]:
    args = ["ffmpeg", "-nostdin", "-loglevel", "error"]
    if start_s > 0:
        args += ["-ss", f"{start_s:.3f}"]  # fast keyframe seek before -i
    args += ["-i", str(src)]
    # Video always stream-copied. The selected audio track (default = first) is transcoded to AAC.
    # A selected subtitle track is muxed in as a WebVTT rendition (ffmpeg aligns it to this session's
    # timeline, so it stays in sync even after a seek-restart); no selection → drop subtitles.
    args += ["-map", "0:v:0", "-map", f"0:a:{audio_ord if audio_ord is not None else 0}?"]
    if sub_ord is not None:
        args += ["-map", f"0:s:{sub_ord}?", "-c:s", "webvtt"]
    else:
        args += ["-sn"]
    args += ["-c:v", "copy", "-c:a", "aac", "-ac", "2", "-b:a", "192k"]
    if sub_ord is not None:
        args += ["-c:s", "webvtt"]
    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",
    ]
    if sub_ord is not None:
        # A subtitle rendition needs a MASTER playlist tying the video variant to the WebVTT group
        # (ffmpeg generates the subtitle segments per-session, so they stay in sync after a seek).
        args += ["-master_pl_name", "master.m3u8", "-var_stream_map", "v:0,a:0,s:0,sgroup:subs"]
    args += ["-hls_segment_filename", str(out_dir / "seg_%d.ts"), str(out_dir / "index.m3u8")]
    return args


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,
    sub_ord: int | None = None,
) -> HlsSession | None:
    """(Re)start the HLS remux for an item at the given offset, with an optional selected audio /
    subtitle track (ordinals among their stream type). 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}_{sub_ord}"
        shutil.rmtree(directory, ignore_errors=True)
        directory.mkdir(parents=True, exist_ok=True)
        proc = subprocess.Popen(
            _ffmpeg_cmd(src, start_s, directory, audio_ord, sub_ord),
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        entry = "master.m3u8" if sub_ord is not None else "index.m3u8"
        s = HlsSession(key, directory, proc, start_s, entry)
        _sessions[key] = s
        _enforce_cap()
        log.info(
            "plex hls session start key=%s start=%.1f audio=%s sub=%s", key, start_s, audio_ord, sub_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