Lyra/worker/lyra_worker/qobuz_gate.py

"""Qobuz pacing ("budget gate").

Keep using Qobuz for hi-res, but at a human-looking volume and cadence so the account isn't
flagged for bulk downloading (which got it 403 USER_BLOCKED). One gate decides, per album,
whether Qobuz is an eligible source right now: within active hours, under today's cap (a warm-up
ramp of a steady cap), and past a randomized spacing since the last Qobuz download. When the gate
is closed the pipeline drops Qobuz candidates and falls through to the next source; monitored
albums upgrade to Qobuz later once budget frees up.

Pure decision logic (`today_cap`, `is_allowed`, `next_gap_seconds`) is separated from the two DB
helpers (`gate_open`, `record_download`) so the policy is unit-tested without a clock or DB.
"""
from dataclasses import dataclass
from datetime import date

_TRUE = {"1", "true", "yes", "on"}
_JITTER = 0.3   # ± fraction around the spread interval
_GAP_FLOOR_S = 60.0


@dataclass(frozen=True)
class QobuzPacing:
    enabled: bool = True
    active_start: int = 8
    active_end: int = 23
    daily_cap: int = 40
    warmup_start_date: str = ""  # ISO "YYYY-MM-DD", or "" for no warm-up ramp

    @classmethod
    def from_config(cls, config: dict) -> "QobuzPacing":
        def _int(key, default):
            try:
                return int(config[key])
            except (KeyError, TypeError, ValueError):
                return default

        return cls(
            enabled=str(config.get("qobuz.pacing.enabled", "true")).strip().lower() in _TRUE,
            active_start=_int("qobuz.pacing.activeStartHour", 8),
            active_end=_int("qobuz.pacing.activeEndHour", 23),
            daily_cap=_int("qobuz.pacing.dailyCap", 40),
            warmup_start_date=str(config.get("qobuz.pacing.warmupStartDate", "") or "").strip(),
        )


def today_cap(pacing: QobuzPacing, today: date) -> int:
    """Today's Qobuz download cap: the steady cap, ramped up by account age (days since
    ``warmup_start_date``). No warm-up date -> steady cap. A small steady cap is never exceeded."""
    steady = max(0, pacing.daily_cap)
    if not pacing.warmup_start_date:
        return steady
    try:
        start = date.fromisoformat(pacing.warmup_start_date)
    except ValueError:
        return steady
    days = max(0, (today - start).days)
    if days < 3:
        return min(5, steady)
    if days < 7:
        return min(10, steady)
    if days < 14:
        return min(20, steady)
    return steady


def is_allowed(pacing: QobuzPacing, now_hour: int, count_today: int, cap: int,
               now_epoch: float, next_allowed_epoch: float) -> bool:
    """Pure gate decision. Qobuz is eligible only within active hours, under today's cap, and past
    the randomized spacing since the last download. Pacing disabled -> always allowed."""
    if not pacing.enabled:
        return True
    if not (pacing.active_start <= now_hour < pacing.active_end):
        return False
    if count_today >= cap:
        return False
    if now_epoch < next_allowed_epoch:
        return False
    return True


def next_gap_seconds(pacing: QobuzPacing, cap: int, rand01: float) -> float:
    """Seconds until the next Qobuz download is allowed: the day's budget spread across the active
    window (``activeWindow / cap``), ± _JITTER, floored at _GAP_FLOOR_S. ``rand01`` in [0,1] is the
    injected randomness (so tests are deterministic; callers pass random.random())."""
    window_s = max(1, pacing.active_end - pacing.active_start) * 3600
    base = window_s / max(1, cap)
    lo, hi = base * (1 - _JITTER), base * (1 + _JITTER)
    gap = lo + (hi - lo) * rand01
    return max(_GAP_FLOOR_S, gap)


def _now_parts(conn):
    """(hour:int, today:date, epoch:float) from the DB clock — avoids container-TZ ambiguity."""
    with conn.cursor() as cur:
        cur.execute("SELECT extract(hour from now())::int, current_date, extract(epoch from now())")
        hour, today, epoch = cur.fetchone()
    return int(hour), today, float(epoch)


def _read(conn, keys):
    with conn.cursor() as cur:
        cur.execute('SELECT key, value FROM "Config" WHERE key = ANY(%s)', (list(keys),))
        return dict(cur.fetchall())


def gate_open(conn) -> bool:
    """Whether Qobuz may be used for a download right now (reads pacing config + counter/spacing
    state from Config and the DB clock). A counter from a previous day is treated as 0."""
    from lyra_worker.config import get_config

    pacing = QobuzPacing.from_config(get_config(conn))
    hour, today, epoch = _now_parts(conn)
    state = _read(conn, ("qobuz.pacing.countDay", "qobuz.pacing.countToday", "qobuz.pacing.nextAllowedAt"))
    count = 0
    if state.get("qobuz.pacing.countDay") == today.isoformat():
        try:
            count = int(state.get("qobuz.pacing.countToday") or 0)
        except (TypeError, ValueError):
            count = 0
    try:
        next_allowed = float(state.get("qobuz.pacing.nextAllowedAt") or 0)
    except (TypeError, ValueError):
        next_allowed = 0.0
    return is_allowed(pacing, hour, count, today_cap(pacing, today), epoch, next_allowed)


def record_download(conn, rand01: float) -> None:
    """Record a successful Qobuz download: bump today's counter (resetting it if the stored day is
    stale) and set the randomized next-allowed time. ``rand01`` is injected randomness."""
    from lyra_worker.config import get_config

    pacing = QobuzPacing.from_config(get_config(conn))
    _hour, today, epoch = _now_parts(conn)
    state = _read(conn, ("qobuz.pacing.countDay", "qobuz.pacing.countToday"))
    if state.get("qobuz.pacing.countDay") == today.isoformat():
        try:
            count = int(state.get("qobuz.pacing.countToday") or 0)
        except (TypeError, ValueError):
            count = 0
    else:
        count = 0
    count += 1
    next_allowed = epoch + next_gap_seconds(pacing, today_cap(pacing, today), rand01)
    _write(conn, {
        "qobuz.pacing.countDay": today.isoformat(),
        "qobuz.pacing.countToday": str(count),
        "qobuz.pacing.nextAllowedAt": repr(next_allowed),
    })


def _write(conn, kv: dict) -> None:
    with conn.cursor() as cur:
        for key, value in kv.items():
            cur.execute(
                'INSERT INTO "Config"(key,value,secret,"updatedAt") VALUES (%s,%s,false,now()) '
                'ON CONFLICT (key) DO UPDATE SET value=EXCLUDED.value, "updatedAt"=now()',
                (key, value),
            )
    conn.commit()