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 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 _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
        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.
    if not _claim_asset(asset_id):
        _set_job(job_id, status="queued", phase="waiting")
        time.sleep(2)
        return

    _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 _make_pp_hook(job_id: int):
    """Postprocessor hook: surface the ffmpeg merge/convert/extract step so the bar shows a
    'Merging/Processing' label instead of freezing at 100% after the streams finish."""

    def hook(d: dict) -> None:
        if d.get("status") in ("started", "processing"):
            pp = d.get("postprocessor") or ""
            phase = "merging" if "Merg" in pp else "processing"
            _set_job(job_id, phase=phase, 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


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)

        outtmpl = str(staging / "%(id)s.%(ext)s")
        opts = formats.build_ydl_opts(spec, outtmpl, _make_progress_hook(job_id, asset_id))
        opts["postprocessor_hooks"] = [_make_pp_hook(job_id)]
        url = service.source_url(source_kind, source_ref)

        _set_job(job_id, phase="downloading", progress=0)
        with yt_dlp.YoutubeDL(opts) as ydl:
            info = ydl.extract_info(url, download=True)

        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()


# --- 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

    _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()