"""Celery beat that remembers when each job last ran — from task_run, not a file.

Celery's default PersistentScheduler keeps its memory in a shelve file in the
working directory. Nothing mounts that directory, so every container recreate
forgets it, and a scheduler that remembers nothing seeds every entry with
`last_run_at = now`: each job waits a FULL interval after startup. A daily job
therefore needs 24 hours without a redeploy to fire. Since the one-container
image (172e33d) every redeploy restarts beat, and on 2026-09-24 no daily or
weekly job had run since the 21st (#4408).

task_run already records every task that starts (celery_signals), indexed on
(task_name, started_at DESC), and prune_task_runs keeps the newest row of each
task however old it is. So on startup each entry takes its last_run_at from
there:
  - a job that is overdue runs at once;
  - a job that is not due waits only the remainder of its interval;
  - a job that has never run is due now.

Beat keeps last_run_at in memory from then on, as the default scheduler does;
only the startup seed changes. If the database cannot be read, the entries keep
Celery's own default rather than beat failing to start.
"""

from __future__ import annotations

import logging
from datetime import UTC, datetime

from celery.beat import Scheduler
from sqlalchemy import func, select

log = logging.getLogger(__name__)

# Seed for a job with no recorded run: far enough back that any interval or
# crontab reads as due.
NEVER = datetime(2000, 1, 1, tzinfo=UTC)


def last_runs(session, task_names: list[str]) -> dict[str, datetime]:
    """The latest recorded start of each task, by task name."""
    from .models import TaskRun

    if not task_names:
        return {}
    rows = session.execute(
        select(TaskRun.task_name, func.max(TaskRun.started_at))
        .where(TaskRun.task_name.in_(sorted(set(task_names))))
        .group_by(TaskRun.task_name)
    ).all()
    return dict(rows)


def seed(entries, last: dict[str, datetime]) -> None:
    """Set each entry's last_run_at from `last`; a task with none is due now."""
    for entry in entries:
        entry.last_run_at = last.get(entry.task, NEVER)


class TaskRunScheduler(Scheduler):
    """An in-memory beat seeded from task_run history at startup."""

    def setup_schedule(self):
        super().setup_schedule()
        try:
            from .tasks._sync_engine import sync_session_factory

            with sync_session_factory()() as session:
                last = last_runs(session, [e.task for e in self.schedule.values()])
        except Exception:
            log.exception("beat: could not read task_run; every job waits a full interval")
            return
        seed(self.schedule.values(), last)
        due = sum(1 for e in self.schedule.values() if e.is_due()[0])
        log.info(
            "beat: seeded %d job(s) from task_run, %d due now",
            len(self.schedule), due,
        )