pigway-frame-flow/pigway_frame_flow/scheduler/engine.py

import json
import logging
import random
import threading
from datetime import UTC, datetime, timedelta
from uuid import uuid4
from zoneinfo import ZoneInfo

from croniter import croniter

from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.models import (
    DisplayRequest,
    ImageContent,
    JobInput,
    Options,
    PlaylistInput,
    RestoreContent,
)
from pigway_frame_flow.storage.database import encode, now_iso


def next_run(schedule, after=None):
    after = after or datetime.now(UTC)
    if schedule.type == "interval":
        if schedule.interval_minutes is None:
            raise FlowError(E.INVALID_REQUEST, "Interval requires interval_minutes")
        return (after + timedelta(minutes=schedule.interval_minutes)).isoformat()
    if schedule.type in ("date", "temporary"):
        if not schedule.at or schedule.at.tzinfo is None:
            raise FlowError(E.INVALID_REQUEST, "One-time date must include timezone")
        return schedule.at.astimezone(UTC).isoformat() if schedule.at > after else None
    if schedule.type == "cron":
        expression = schedule.cron or ""
        if len(expression.split()) != 5 or not croniter.is_valid(expression):
            raise FlowError(E.INVALID_REQUEST, "Cron requires five valid fields")
    else:
        if not schedule.time:
            raise FlowError(E.INVALID_REQUEST, "Daily schedule requires HH:MM")
        hour, minute = schedule.time.split(":")
        day = (
            "1-5"
            if schedule.type == "weekdays"
            else ("6,0" if schedule.type == "weekends" else "*")
        )
        expression = f"{minute} {hour} * * {day}"
    return (
        croniter(expression, after.astimezone(ZoneInfo("Asia/Shanghai")))
        .get_next(datetime)
        .astimezone(UTC)
        .isoformat()
    )


class Scheduler:
    def __init__(self, service):
        self.service, self.db = service, service.db
        self.stop_event = threading.Event()
        self.thread = None
        self.webhook_thread = None
        self.tick_lock = threading.Lock()

    def jobs(self):
        rows = [
            dict(
                json.loads(x["payload"]),
                id=x["id"],
                created_at=x["created_at"],
                updated_at=x["updated_at"],
                last_run=x["last_run"],
                last_result=x["last_result"],
                next_run=x["next_run"],
            )
            for x in self.db.all("SELECT * FROM jobs ORDER BY created_at DESC")
        ]

        for row in rows:
            if row["last_result"]:
                operation = self.db.one(
                    "SELECT status FROM operations WHERE id=?", (row["last_result"],)
                )
                if operation:
                    row["last_result"] = operation["status"]
        return rows

    def save_job(self, job, id=None):
        self.service.registry.get(job.device_id)
        if job.source.type != "restore":
            self.service.validate(
                DisplayRequest(device_id=job.device_id, content=job.source),
                writing=False,
            )
        upcoming = next_run(job.schedule)
        if job.enabled and upcoming is None:
            raise FlowError(E.INVALID_REQUEST, "One-time task must be in the future")
        id = id or uuid4().hex
        now = now_iso()
        self.db.execute(
            "INSERT INTO jobs(id,payload,enabled,next_run,created_at,updated_at) VALUES(?,?,?,?,?,?) ON CONFLICT(id) DO UPDATE SET payload=excluded.payload,enabled=excluded.enabled,next_run=excluded.next_run,updated_at=excluded.updated_at",
            (
                id,
                encode(self.service.settings.sanitize(job.model_dump(mode="json"))),
                int(job.enabled),
                upcoming,
                now,
                now,
            ),
        )
        return next(x for x in self.jobs() if x["id"] == id)

    def playlists(self):
        return [
            dict(
                json.loads(x["payload"]),
                id=x["id"],
                position=x["position"],
                next_run=x["next_run"],
            )
            for x in self.db.all("SELECT * FROM playlists")
        ]

    def save_playlist(self, playlist, id=None):
        self.service.registry.get(playlist.device_id)
        for existing in self.playlists():
            if (
                existing["enabled"]
                and playlist.enabled
                and existing["device_id"] == playlist.device_id
                and existing["id"] != id
            ):
                raise FlowError(
                    E.CONFLICT, "Only one Frame Flow playlist may run per device", 409
                )
        id = id or uuid4().hex
        self.db.execute(
            "INSERT INTO playlists VALUES(?,?,?,?) ON CONFLICT(id) DO UPDATE SET payload=excluded.payload,next_run=excluded.next_run",
            (
                id,
                encode(playlist.model_dump()),
                -1,
                now_iso() if playlist.enabled else None,
            ),
        )
        return next(x for x in self.playlists() if x["id"] == id)

    def tick(self, current=None):
        current = current or datetime.now(UTC)
        with self.tick_lock:
            for row in self.db.all(
                "SELECT * FROM jobs WHERE enabled=1 AND next_run<=?",
                (current.isoformat(),),
            ):
                try:
                    job = JobInput.model_validate_json(row["payload"])
                    slot = row["next_run"]
                    if self.db.one(
                        "SELECT * FROM job_runs WHERE job_id=? AND scheduled_at=?",
                        (row["id"], slot),
                    ):
                        continue
                    # Skip overdue tasks (>5 minutes): avoid an unexpected artwork
                    # overwrite after downtime. One-time tasks are retired too.
                    late = (
                        current - datetime.fromisoformat(slot)
                    ).total_seconds() > 300
                    operation_id = None
                    if not late:
                        action = (
                            "restore" if job.source.type == "restore" else "display"
                        )
                        payload = (
                            {}
                            if isinstance(job.source, RestoreContent)
                            else DisplayRequest(
                                device_id=job.device_id,
                                content=job.source,
                                options=Options(
                                    duration=job.display_duration,
                                    restore_previous=job.restore_previous,
                                    priority=job.priority,
                                    cleanup_unstarred=job.cleanup_unstarred,
                                ),
                            ).model_dump(mode="json")
                        )
                        operation_id = self.service.submit(
                            action,
                            payload,
                            job.device_id,
                            "scheduler",
                            "schedule-" + uuid4().hex,
                            "job-" + row["id"] + "-" + slot,
                        )
                    following = next_run(job.schedule, current)
                    enabled = bool(following)
                    updated = job.model_copy(update={"enabled": enabled}).model_dump(
                        mode="json"
                    )
                    with self.db.transaction() as tx:
                        tx.execute(
                            "INSERT OR IGNORE INTO job_runs VALUES(?,?,?)",
                            (row["id"], slot, operation_id),
                        )
                        tx.execute(
                            "UPDATE jobs SET payload=?,enabled=?,next_run=?,last_run=?,last_result=? WHERE id=?",
                            (
                                encode(updated),
                                int(enabled),
                                following,
                                now_iso(),
                                "SKIPPED_OVERDUE" if late else operation_id,
                                row["id"],
                            ),
                        )
                except Exception:
                    # Keep this slot pending for retry or the existing overdue policy.
                    # A single device must not prevent other jobs or queue consumption.
                    logging.getLogger("frame-flow").error(
                        "Scheduled job failed; details suppressed"
                    )
            for row in self.db.all(
                "SELECT * FROM playlists WHERE next_run<=?", (current.isoformat(),)
            ):
                p = PlaylistInput.model_validate_json(row["payload"])
                if (
                    not p.enabled
                    or self.service.top(p.device_id)
                    or self.service.blocked(p.device_id)
                ):
                    continue
                if (
                    not self.service.settings.mock
                    and not self.service.settings.write_enabled
                ):
                    continue
                pos = (
                    random.randrange(len(p.image_ids))
                    if p.mode == "random"
                    else (row["position"] + 1) % len(p.image_ids)
                )
                request = DisplayRequest(
                    device_id=p.device_id,
                    content=ImageContent(type="image", image_id=p.image_ids[pos]),
                    options=Options(priority=0),
                )
                self.service.submit(
                    "display",
                    request.model_dump(mode="json"),
                    p.device_id,
                    "scheduler",
                    "playlist-" + uuid4().hex,
                    "playlist-" + row["id"] + "-" + row["next_run"],
                )
                following = (
                    current + timedelta(minutes=p.interval_minutes)
                ).isoformat()
                self.db.execute(
                    "UPDATE playlists SET position=?,next_run=? WHERE id=?",
                    (pos, following, row["id"]),
                )

    def start(self):
        self.thread = threading.Thread(
            target=self.run, name="frame-flow-worker", daemon=True
        )
        self.thread.start()
        if hasattr(self.service, "webhooks"):
            self.webhook_thread = threading.Thread(
                target=self.run_webhooks, name="frame-flow-webhooks", daemon=True
            )
            self.webhook_thread.start()

    def run_webhooks(self):
        while not self.stop_event.is_set():
            try:
                self.service.device_observer.tick()
                self.service.webhooks.tick()
            except Exception:
                logging.getLogger("frame-flow").error(
                    "Webhook cycle failed; details suppressed"
                )
            self.stop_event.wait(60)

    def run(self):
        while not self.stop_event.is_set():
            try:
                # Resolve late sync before submitting due jobs: a write lock
                # otherwise makes tick fail before reconciliation can run.
                self.service.sync()
                self.tick()
                self.service.expire()
                self.service.cloud_library.sweep()
                if hasattr(self.service, "local_retention"):
                    self.service.local_retention.sweep()
                processed = self.service.process_one()
                self.service.sync()
            except Exception:
                logging.getLogger("frame-flow").error(
                    "Background cycle failed; details suppressed"
                )
                processed = False
            self.stop_event.wait(1 if processed else 5)

    def stop(self):
        self.stop_event.set()
        if self.thread:
            self.thread.join(timeout=45)
        if self.webhook_thread:
            self.webhook_thread.join(timeout=25)

    def healthy(self):
        return bool(
            self.thread and self.thread.is_alive() and not self.stop_event.is_set()
        )