siftlode/backend/app/worker.py

"""Download worker process entry point.

Runs in a DEDICATED container (`command: python -m app.worker`, `WORKER_ENABLED=1`), separate
from the API so long yt-dlp downloads + ffmpeg postprocessing never block web requests.

Design:
  * N loop threads (download_worker_concurrency) each claim a `queued` DownloadJob with
    FOR UPDATE SKIP LOCKED and process it.
  * The shared MediaAsset is the dedup unit. A job whose asset is already `ready` finishes
    instantly (no re-download); a fresh asset is claimed pending→downloading by exactly one
    worker (a DB compare-and-set), so two jobs for the same video+format never double-download.
  * Live progress is written to the job row (short-lived sessions, throttled ~1/s) — the
    frontend polls it; the worker can't reach the API's in-process WebSocket registry.
  * Pause/cancel are cooperative: the progress hook re-reads the job status and aborts the
    yt-dlp download if it's no longer `running`.
  * Crash-safe: on startup, orphaned `running` jobs are requeued and `downloading` assets reset.
"""
import logging
import shutil
import signal
import threading
import time
from datetime import date, datetime, timedelta, timezone
from pathlib import Path

import yt_dlp
from sqlalchemy import text

from app.config import settings
from app.db import SessionLocal
from app.downloads import edit as editmod
from app.downloads import formats, quota, service, storage
from app.models import DownloadJob, MediaAsset
from app.titles import normalize_title

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

_stop = threading.Event()


class _Aborted(Exception):
    """Raised inside the progress hook when the job was paused/canceled mid-download."""


def _handle_signal(signum, frame):
    _stop.set()


# --- small DB helpers (short-lived sessions; the download itself holds no transaction) -------

def _set_job(job_id: int, **fields) -> None:
    with SessionLocal() as db:
        job = db.get(DownloadJob, job_id)
        if job is None:
            return
        for k, v in fields.items():
            setattr(job, k, v)
        db.commit()


def _job_status(job_id: int) -> str | None:
    with SessionLocal() as db:
        return db.execute(
            text("SELECT status FROM download_jobs WHERE id = :id"), {"id": job_id}
        ).scalar()


def _parse_upload_date(raw) -> date | None:
    if not raw:
        return None
    try:
        return datetime.strptime(str(raw), "%Y%m%d").date()
    except ValueError:
        return None


# --- claim + recovery -----------------------------------------------------------------------

def _wait_for_schema(timeout: float = 300.0) -> None:
    """Block until the download tables exist before doing any DB work.

    The API applies migrations on its OWN startup, and the worker is a separate container that may
    boot first (they start in parallel). Without this, the worker would hit a missing `download_jobs`
    table and crash-loop until the API's migrations land. Poll patiently instead."""
    deadline = time.monotonic() + timeout
    while not _stop.is_set():
        try:
            with SessionLocal() as db:
                db.execute(text("SELECT 1 FROM download_jobs LIMIT 1"))
            return
        except Exception as exc:  # noqa: BLE001 — table not migrated yet / DB not up yet
            if time.monotonic() > deadline:
                log.warning("download schema still absent after %ss; continuing: %s", timeout, str(exc)[:120])
                return
            log.info("waiting for the database schema (API migrations)…")
            _stop.wait(3.0)


def _recover_orphans() -> None:
    """Reset state left behind by a previous crash so nothing is stuck."""
    with SessionLocal() as db:
        db.execute(text("UPDATE download_jobs SET status='queued' WHERE status='running'"))
        db.execute(text("UPDATE media_assets SET status='pending' WHERE status='downloading'"))
        db.commit()


def _claim_job() -> int | None:
    """Claim the oldest queued job whose user is under their per-user concurrency limit.

    Fetches a window of locked candidates (FOR UPDATE SKIP LOCKED) and picks the first eligible
    one, so one user flooding the queue can't starve others beyond their max_concurrent."""
    with SessionLocal() as db:
        rows = db.execute(
            text(
                """
                SELECT id, user_id FROM download_jobs
                WHERE status='queued'
                ORDER BY queue_pos, id
                FOR UPDATE SKIP LOCKED
                LIMIT 50
                """
            )
        ).fetchall()
        chosen = next(
            (jid for jid, uid in rows if not quota.at_concurrency_limit(db, uid)), None
        )
        if chosen is None:
            db.commit()  # release the row locks; nothing eligible right now
            return None
        db.execute(
            text(
                "UPDATE download_jobs SET status='running', phase='starting', updated_at=now() "
                "WHERE id = :id"
            ),
            {"id": chosen},
        )
        db.commit()
        return chosen


def _claim_asset(asset_id: int) -> bool:
    """Compare-and-set pending→downloading. True if THIS worker won the right to download."""
    with SessionLocal() as db:
        row = db.execute(
            text(
                "UPDATE media_assets SET status='downloading' "
                "WHERE id = :id AND status='pending' RETURNING id"
            ),
            {"id": asset_id},
        ).fetchone()
        db.commit()
        return row is not None


# --- processing -----------------------------------------------------------------------------

def _process_job(job_id: int) -> None:
    with SessionLocal() as db:
        job = db.get(DownloadJob, job_id)
        if job is None:
            return
        asset = db.get(MediaAsset, job.asset_id) if job.asset_id else None
        spec = dict(job.profile_snapshot or {})
        source_kind, source_ref = job.source_kind, job.source_ref
        job_kind = job.job_kind
        asset_id = asset.id if asset else None
        asset_status = asset.status if asset else None

    if asset is None:
        _set_job(job_id, status="error", error="No media asset", phase=None)
        return

    if asset_status == "ready":
        _finish_from_cache(job_id, asset_id)
        return
    if asset_status == "error":
        _set_job(job_id, status="error", error=asset.error or "Download failed", phase=None)
        return
    if asset_status == "downloading":
        # Another worker owns this asset; wait and let the job be reclaimed later.
        _set_job(job_id, status="queued", phase="waiting")
        time.sleep(2)
        return

    # asset_status == "pending": try to win the download / edit.
    if not _claim_asset(asset_id):
        _set_job(job_id, status="queued", phase="waiting")
        time.sleep(2)
        return

    if job_kind == "edit":
        _process_edit(job_id, asset_id)
    else:
        _download(job_id, asset_id, source_kind, source_ref, spec)


def _finish_from_cache(job_id: int, asset_id: int) -> None:
    with SessionLocal() as db:
        job = db.get(DownloadJob, job_id)
        asset = db.get(MediaAsset, asset_id)
        if job and asset:
            job.status = "done"
            job.progress = 100
            job.phase = None
            if not job.display_name:
                job.display_name = asset.title
            asset.last_access_at = datetime.now(timezone.utc)
        db.commit()


def _fill_asset_meta_early(asset_id: int, info: dict) -> None:
    """Populate an asset's display metadata from yt-dlp's info_dict the first time it's seen
    (only if not already filled at enqueue from our catalog) — so an ad-hoc URL's row shows a
    real title + thumbnail while it downloads, not just the URL. Zero extra network cost."""
    with SessionLocal() as db:
        asset = db.get(MediaAsset, asset_id)
        if asset is None or asset.title:
            return
        asset.title = normalize_title(info.get("title"))
        asset.uploader = info.get("uploader") or info.get("channel")
        asset.thumbnail_url = info.get("thumbnail")
        asset.duration_s = int(info["duration"]) if info.get("duration") else None
        asset.upload_date = _parse_upload_date(info.get("upload_date"))
        db.commit()


def _make_progress_hook(job_id: int, asset_id: int):
    state = {"last_db": 0.0, "last_check": 0.0, "meta_done": False}

    def hook(d: dict) -> None:
        now = time.monotonic()
        # Cooperative pause/cancel: if the job is no longer 'running', abort the download.
        if now - state["last_check"] >= 2.0:
            state["last_check"] = now
            if _job_status(job_id) not in ("running", None):
                raise _Aborted
        if not state["meta_done"] and d.get("info_dict"):
            state["meta_done"] = True
            try:
                _fill_asset_meta_early(asset_id, d["info_dict"])
            except Exception:  # noqa: BLE001 — metadata is best-effort, never fail a download
                pass
        if d.get("status") != "downloading":
            return
        if now - state["last_db"] < 1.0:
            return
        state["last_db"] = now
        # yt-dlp downloads the video and audio streams SEPARATELY (each 0→100%), then merges.
        # Label the phase so the user understands the bar resetting between streams instead of
        # seeing a mysterious 95%→5% jump.
        info = d.get("info_dict") or {}
        vcodec = info.get("vcodec") or "none"
        acodec = info.get("acodec") or "none"
        if vcodec != "none" and acodec == "none":
            phase = "video"
        elif acodec != "none" and vcodec == "none":
            phase = "audio"
        else:
            phase = "media"
        total = d.get("total_bytes") or d.get("total_bytes_estimate") or 0
        done = d.get("downloaded_bytes") or 0
        pct = int(done * 100 / total) if total else 0
        _set_job(
            job_id,
            progress=min(pct, 99),
            speed_bps=int(d["speed"]) if d.get("speed") else None,
            eta_s=int(d["eta"]) if d.get("eta") else None,
            phase=phase,
        )

    return hook


def _pp_phase(pp: str) -> str:
    """Map a yt-dlp postprocessor class name to a user-facing phase label. ffmpeg post-steps
    aren't byte-progress operations (no %), so the UI shows the step name + an indeterminate
    pulse — this makes the long merge/finalize legible instead of a silent 'Processing'."""
    if "Merg" in pp:
        return "merging"
    if "ExtractAudio" in pp:
        return "audio_extract"
    if "Thumbnail" in pp:
        return "thumbnail"
    if "SponsorBlock" in pp or "ModifyChapters" in pp:
        return "sponsorblock"
    if "Metadata" in pp:
        return "metadata"
    return "processing"


def _make_pp_hook(job_id: int):
    """Postprocessor hook: surface the current ffmpeg step so the bar shows a concrete label
    (Merging / Embedding thumbnail / …) instead of freezing at 100% after the streams finish."""

    def hook(d: dict) -> None:
        if d.get("status") in ("started", "processing"):
            _set_job(job_id, phase=_pp_phase(d.get("postprocessor") or ""), speed_bps=None, eta_s=None)

    return hook


def _staging_dir(asset_id: int) -> Path:
    return Path(settings.download_root) / ".staging" / f"asset-{asset_id}"


def _find_thumbnail(staging: Path, media: Path) -> Path | None:
    for p in staging.iterdir():
        if p != media and p.suffix.lower() in (".jpg", ".jpeg", ".png", ".webp"):
            return p
    return None


# Errors worth retrying with a different player-client set (YouTube per-client flakiness).
_RETRIABLE_MARKERS = (
    "not a bot",
    "DRM protected",
    "Requested format is not available",
    "Only images are available",
    "Sign in to confirm",
)


def _extract_with_fallback(job_id: int, asset_id: int, url: str, spec: dict, staging: Path):
    """Download via yt-dlp, trying the primary player clients then the POT-backed fallback set
    on a bot/DRM/format failure. Returns the info dict of the successful attempt."""
    client_sets = [
        c for c in (settings.player_client_list, settings.player_client_fallback_list) if c
    ] or [None]
    outtmpl = str(staging / "%(id)s.%(ext)s")
    last_exc: Exception | None = None
    for i, clients in enumerate(client_sets):
        opts = formats.build_ydl_opts(
            spec, outtmpl, _make_progress_hook(job_id, asset_id),
            settings.download_pot_base_url, clients,
        )
        opts["postprocessor_hooks"] = [_make_pp_hook(job_id)]
        try:
            with yt_dlp.YoutubeDL(opts) as ydl:
                return ydl.extract_info(url, download=True)
        except yt_dlp.utils.DownloadError as exc:
            last_exc = exc
            retriable = any(m in str(exc) for m in _RETRIABLE_MARKERS)
            if i < len(client_sets) - 1 and retriable:
                log.info("job %s: clients %s failed (%s); retrying with fallback", job_id, clients, str(exc)[:70])
                for p in list(staging.iterdir()):  # drop partials before the retry
                    try:
                        p.unlink()
                    except OSError:
                        pass
                _set_job(job_id, progress=0, phase="downloading")
                continue
            raise
    raise last_exc  # unreachable (loop either returns or raises)


def _download(job_id: int, asset_id: int, source_kind: str, source_ref: str, spec: dict) -> None:
    staging = _staging_dir(asset_id)
    try:
        if staging.exists():
            shutil.rmtree(staging, ignore_errors=True)
        staging.mkdir(parents=True, exist_ok=True)

        _set_job(job_id, phase="downloading", progress=0)
        url = service.source_url(source_kind, source_ref)
        info = _extract_with_fallback(job_id, asset_id, url, spec, staging)

        req = (info.get("requested_downloads") or [{}])[0]
        produced = Path(req.get("filepath") or "")
        if not produced.exists():
            # Fallback: the single media file left in staging (ignore the thumbnail).
            cands = [p for p in staging.iterdir() if p.suffix.lower() not in (".jpg", ".jpeg", ".png", ".webp")]
            if not cands:
                raise RuntimeError("yt-dlp produced no output file")
            produced = cands[0]

        ext = produced.suffix.lstrip(".").lower()
        meta = storage.MediaMeta(
            video_id=info.get("id") or source_ref,
            title=normalize_title(info.get("title")) or source_ref,
            uploader=info.get("uploader") or info.get("channel") or "Unknown Channel",
            upload_date=_parse_upload_date(info.get("upload_date")),
            duration_s=int(info["duration"]) if info.get("duration") else None,
            description=info.get("description"),
            thumbnail_url=info.get("thumbnail"),
        )

        rel = storage.rel_path(meta, ext, settings.download_layout)
        thumb = _find_thumbnail(staging, produced)
        storage.place_file(produced, settings.download_root, rel)
        nfo_ok = storage.write_sidecars(settings.download_root, rel, meta, thumb)

        final = storage.abs_path(settings.download_root, rel)
        size = final.stat().st_size if final.exists() else None
        expires = datetime.now(timezone.utc) + timedelta(days=settings.download_retention_days)

        with SessionLocal() as db:
            asset = db.get(MediaAsset, asset_id)
            if asset:
                asset.status = "ready"
                asset.rel_path = rel
                asset.size_bytes = size
                asset.duration_s = meta.duration_s
                asset.width = int(info["width"]) if info.get("width") else None
                asset.height = int(info["height"]) if info.get("height") else None
                asset.container = ext
                asset.vcodec = info.get("vcodec") if info.get("vcodec") not in (None, "none") else None
                asset.acodec = info.get("acodec") if info.get("acodec") not in (None, "none") else None
                asset.title = meta.title
                asset.uploader = meta.uploader
                asset.upload_date = meta.upload_date
                asset.thumbnail_url = meta.thumbnail_url
                asset.nfo_written = nfo_ok
                asset.error = None
                asset.last_access_at = datetime.now(timezone.utc)
                asset.expires_at = expires
            job = db.get(DownloadJob, job_id)
            if job:
                job.status = "done"
                job.progress = 100
                job.phase = None
                job.speed_bps = None
                job.eta_s = None
                if not job.display_name:
                    job.display_name = meta.title
            db.commit()
        log.info("job %s done: %s", job_id, rel)

    except _Aborted:
        # Paused/canceled mid-download: free the asset for a later retry; keep the job's status
        # (the route set it to paused/canceled). ref_count is adjusted by the route.
        _reset_asset_pending(asset_id)
        log.info("job %s aborted (pause/cancel)", job_id)
    except Exception as exc:  # noqa: BLE001 — surface any failure to the user, keep worker alive
        msg = str(exc)[:500]
        with SessionLocal() as db:
            asset = db.get(MediaAsset, asset_id)
            if asset:
                asset.status = "error"
                asset.error = msg
            job = db.get(DownloadJob, job_id)
            if job:
                job.status = "error"
                job.error = msg
                job.phase = None
            db.commit()
        log.warning("job %s failed: %s", job_id, msg)
    finally:
        shutil.rmtree(staging, ignore_errors=True)


def _reset_asset_pending(asset_id: int) -> None:
    with SessionLocal() as db:
        asset = db.get(MediaAsset, asset_id)
        if asset and asset.status == "downloading":
            asset.status = "pending"
            db.commit()


# --- editor derivatives (phase 2: trim/crop via ffmpeg) -------------------------------------

def _edit_staging_dir(asset_id: int) -> Path:
    return Path(settings.download_root) / ".staging" / f"edit-{asset_id}"


def _fail_edit(job_id: int, asset_id: int, msg: str) -> None:
    with SessionLocal() as db:
        asset = db.get(MediaAsset, asset_id)
        if asset:
            asset.status = "error"
            asset.error = msg
        job = db.get(DownloadJob, job_id)
        if job:
            job.status = "error"
            job.error = msg
            job.phase = None
        db.commit()
    log.warning("edit job %s failed: %s", job_id, msg)


def _process_edit(job_id: int, asset_id: int) -> None:
    """Produce a per-user trim/crop clip from the source asset's downloaded file via ffmpeg."""
    with SessionLocal() as db:
        job = db.get(DownloadJob, job_id)
        src = db.get(MediaAsset, job.source_asset_id) if job and job.source_asset_id else None
        edit_spec = dict(job.edit_spec or {}) if job else {}
        user_id = job.user_id if job else None
        src_ready = src is not None and src.status == "ready" and bool(src.rel_path)
        src_rel = src.rel_path if src else None
        src_w, src_h, src_dur = (src.width, src.height, src.duration_s) if src else (None, None, None)

    if not src_ready or not src_rel:
        _fail_edit(job_id, asset_id, "The source download is no longer available.")
        return
    src_path = storage.abs_path(settings.download_root, src_rel)
    if not src_path.exists():
        _fail_edit(job_id, asset_id, "The source file is missing.")
        return

    src_ext = src_path.suffix.lstrip(".").lower() or "mp4"
    crop = editmod.has_crop(edit_spec)
    out_ext = "mp4" if editmod.needs_reencode(edit_spec) else src_ext
    total_s = editmod.clip_duration(edit_spec, src_dur) or src_dur or 0
    staging = _edit_staging_dir(asset_id)

    try:
        if staging.exists():
            shutil.rmtree(staging, ignore_errors=True)
        staging.mkdir(parents=True, exist_ok=True)
        dest_staging = staging / f"out.{out_ext}"

        _set_job(job_id, phase="editing", progress=0, speed_bps=None, eta_s=None)
        if edit_spec.get("segments"):
            cmd, prep = editmod.build_concat_plan(src_path, dest_staging, edit_spec, out_ext, staging)
            for p, content in prep.items():
                p.write_text(content, encoding="utf-8")
        else:
            cmd = editmod.build_edit_cmd(src_path, dest_staging, edit_spec, out_ext)
        editmod.run_ffmpeg(
            cmd,
            total_s,
            on_progress=lambda pct: _set_job(job_id, progress=int(pct), phase="editing"),
            should_cancel=lambda: _job_status(job_id) not in ("running", None),
        )
        if not dest_staging.exists():
            raise RuntimeError("ffmpeg produced no output file")

        rel = storage.edit_rel_path(user_id, asset_id, out_ext)
        storage.place_file(dest_staging, settings.download_root, rel)
        final = storage.abs_path(settings.download_root, rel)
        size = final.stat().st_size if final.exists() else None
        expires = datetime.now(timezone.utc) + timedelta(days=settings.download_retention_days)

        with SessionLocal() as db:
            asset = db.get(MediaAsset, asset_id)
            if asset:
                asset.status = "ready"
                asset.rel_path = rel
                asset.size_bytes = size
                asset.container = out_ext
                asset.duration_s = total_s or asset.duration_s
                if crop:
                    asset.width = edit_spec["crop"]["w"]
                    asset.height = edit_spec["crop"]["h"]
                else:
                    asset.width, asset.height = src_w, src_h
                asset.error = None
                asset.last_access_at = datetime.now(timezone.utc)
                asset.expires_at = expires
            job = db.get(DownloadJob, job_id)
            if job:
                job.status = "done"
                job.progress = 100
                job.phase = None
            db.commit()
        log.info("edit job %s done: %s", job_id, rel)

    except editmod.EditAborted:
        # Paused/canceled mid-encode: free the asset for a later retry (ref_count is the route's job).
        _reset_asset_pending(asset_id)
        log.info("edit job %s aborted (pause/cancel)", job_id)
    except Exception as exc:  # noqa: BLE001 — surface the failure, keep the worker alive
        _fail_edit(job_id, asset_id, str(exc)[:500])
    finally:
        shutil.rmtree(staging, ignore_errors=True)


# --- loop -----------------------------------------------------------------------------------

def _worker_loop(idx: int) -> None:
    while not _stop.is_set():
        try:
            job_id = _claim_job()
        except Exception:  # noqa: BLE001
            log.exception("claim failed")
            _stop.wait(settings.download_poll_seconds)
            continue
        if job_id is None:
            _stop.wait(settings.download_poll_seconds)
            continue
        try:
            _process_job(job_id)
        except Exception:  # noqa: BLE001
            log.exception("processing job %s crashed", job_id)
            _set_job(job_id, status="error", error="Worker error", phase=None)


def main() -> None:
    logging.basicConfig(level=logging.INFO)
    signal.signal(signal.SIGTERM, _handle_signal)
    signal.signal(signal.SIGINT, _handle_signal)

    if not settings.worker_enabled:
        log.info("WORKER_ENABLED is off; worker exiting (nothing to do).")
        return

    _wait_for_schema()  # the API may still be applying migrations when we boot
    if _stop.is_set():
        return
    _recover_orphans()
    n = max(1, settings.download_worker_concurrency)
    log.info("Download worker started (concurrency=%d, root=%s).", n, settings.download_root)
    threads = [threading.Thread(target=_worker_loop, args=(i,), daemon=True) for i in range(n)]
    for t in threads:
        t.start()
    while not _stop.is_set():
        _stop.wait(1.0)
    log.info("Download worker stopping…")
    for t in threads:
        t.join(timeout=5)


if __name__ == "__main__":
    main()