"""The worker's scheduler: APScheduler 3 in site time, settings read on every fire, "run now" over NOTIFY.

Each job has a SettingsTrigger whose next fire time comes from the job's cadence and the current feed settings,
so a changed interval applies on the next fire. A save NOTIFYs `ps_settings` and the listener calls
`reschedule()` so it applies at once. A "run now" NOTIFYs `ps_job_request` and the listener moves the job to now.
After a run the job's next fire is measured from the end of the run.
"""

from __future__ import annotations

import contextlib
import json
import logging
import os
import socket
import threading
from collections.abc import Callable, Iterable
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo

from apscheduler.executors.pool import ThreadPoolExecutor
from apscheduler.schedulers.background import BackgroundScheduler
from apscheduler.triggers.base import BaseTrigger

from .. import db
from ..feed import config
from ..feed.config import FeedSettings
from ..settings import Settings
from . import heartbeat, runs
from .jobs import JobSpec, Stack, load_jobs
from .redaction import redact

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

CHANNELS = ("ps_settings", "ps_job_request")
IDLE_RECHECK_SECONDS = 3600  # a job with no schedule wakes this often and does nothing
BACKOFF_MAX = 30.0
THREADS = 4
MISFIRE_GRACE = 300


class SettingsTrigger(BaseTrigger):
    """Fires `seconds(now)` after now. `seconds` reads the current settings, so a change applies on the next fire.
    None means "no schedule": the trigger wakes again after an hour and the fire does nothing."""

    __slots__ = ("seconds",)

    def __init__(self, seconds: Callable[[datetime], int | None]):
        self.seconds = seconds

    def get_next_fire_time(self, previous_fire_time: datetime | None, now: datetime) -> datetime:
        secs = self.seconds(now)
        return now + timedelta(seconds=max(1, IDLE_RECHECK_SECONDS if secs is None else int(secs)))

    def __str__(self) -> str:
        return "settings"


def next_due(
    spec: JobSpec, feed: FeedSettings, now_site: datetime, last_end_site: datetime | None = None
) -> datetime | None:
    """When the job is next due: its cadence measured from the end of its last run (or from now when it never ran),
    never before now. None when the cadence says "no schedule"."""
    base = last_end_site if last_end_site is not None and last_end_site <= now_site else now_site
    secs = spec.cadence(feed, base)
    if secs is None:
        return None
    due = base + timedelta(seconds=max(1, int(secs)))
    return due if due > now_site else now_site


class Scheduler:
    def __init__(
        self,
        settings: Settings,
        *,
        build: str | None = None,
        jobs: dict[str, JobSpec] | None = None,
        connect: Callable[[], Any] | None = None,
    ):
        self.settings = settings
        self.build = build or settings.ps_build
        self.jobs = dict(jobs) if jobs is not None else dict(load_jobs())
        self.tz = ZoneInfo(settings.site_tz)
        self._connect = connect or (lambda: db.connect_as("worker", settings, autocommit=True))
        self.stack = Stack()
        self.host = socket.gethostname()
        self.pid = os.getpid()
        self.feed = FeedSettings()
        self._lock = threading.Lock()
        self._running: set[str] = set()
        self._next_trigger: dict[str, str] = {}
        self._again: dict[str, str] = {}
        self._last_end: dict[str, datetime] = {}
        self._db_ready = False
        self._stop = threading.Event()
        self._listener: threading.Thread | None = None
        self.listening = False
        self.aps = BackgroundScheduler(
            executors={"default": ThreadPoolExecutor(THREADS)},
            job_defaults={"coalesce": True, "max_instances": 1, "misfire_grace_time": MISFIRE_GRACE},
            timezone=self.tz,
        )
        self.started = False

    # ------------------------------------------------ small helpers used by runs.run_job

    def connect(self) -> Any:
        return self._connect()

    def now(self) -> datetime:
        return datetime.now(self.tz)

    def set_feed(self, feed: FeedSettings) -> None:
        self.feed = feed

    def running(self) -> list[str]:
        with self._lock:
            return sorted(self._running)

    def running_add(self, name: str) -> None:
        with self._lock:
            self._running.add(name)

    def running_discard(self, name: str) -> None:
        with self._lock:
            self._running.discard(name)

    # ------------------------------------------------ lifecycle

    def start(self) -> None:
        """Set up the database side when it is reachable, add the jobs and start the scheduler and listener."""
        self._setup_db()
        now = self.now()
        for name, spec in self.jobs.items():
            self.aps.add_job(
                self._fire,
                trigger=SettingsTrigger(self._seconds_for(spec)),
                args=(name,),
                id=name,
                name=name,
                next_run_time=self._next_time(spec, now),
                replace_existing=True,
            )
        self.aps.start()
        self.started = True
        self._write_schedule(self.jobs)
        self._stop.clear()
        self._listener = threading.Thread(target=self._listen, name="ps-worker-listen", daemon=True)
        self._listener.start()
        log.info("scheduler started", extra={"ps": {"jobs": sorted(self.jobs), "host": self.host}})

    def shutdown(self, wait: bool = True) -> None:
        self._stop.set()
        if self.started:
            self.aps.shutdown(wait=wait)
            self.started = False
        if self._listener is not None:
            self._listener.join(5.0)
            self._listener = None

    def _setup_db(self) -> bool:
        """Seed the feed settings, mark this host's leftover runs abandoned, load the settings and the last run
        ends. Safe to call again; returns False when the database is not reachable."""
        if self._db_ready:
            return True
        try:
            with self.connect() as conn:
                with conn.transaction():
                    if config.seed(conn):
                        log.info("feed settings seeded with the defaults")
                self._abandon(conn)
                self.feed = config.load(conn)
                for job, end in conn.execute(
                    "SELECT job, max(finished_at) FROM job_runs WHERE finished_at IS NOT NULL"
                    " AND status <> 'skipped' GROUP BY job"
                ).fetchall():
                    if end is not None:
                        self._last_end[job] = end.astimezone(self.tz)
        except Exception as e:
            log.warning("worker database setup deferred: %s", redact(f"{type(e).__name__}: {e}"))
            return False
        self._db_ready = True
        return True

    def _abandon(self, conn: Any) -> None:
        """Runs left `running` by a dead worker become `abandoned`: this host's always, another host's when no
        session holds that job's advisory lock (a container recreated with a new host name)."""
        rows = conn.execute("SELECT id, job, host FROM job_runs WHERE status = 'running'").fetchall()
        for rid, job, host in rows:
            if host != self.host:
                got = conn.execute("SELECT pg_try_advisory_lock(%s, hashtext(%s))", (runs.LOCK_SPACE, job)).fetchone()[
                    0
                ]
                if not got:
                    continue
                conn.execute("SELECT pg_advisory_unlock(%s, hashtext(%s))", (runs.LOCK_SPACE, job))
            conn.execute(
                "UPDATE job_runs SET status = 'abandoned', finished_at = now(),"
                " errors = errors || '[\"The worker stopped during this run.\"]'::jsonb"
                " WHERE id = %s AND status = 'running'",
                (rid,),
            )
            log.warning("job run %s (%s) marked abandoned", rid, job)

    # ------------------------------------------------ timing

    def _seconds_for(self, spec: JobSpec) -> Callable[[datetime], int | None]:
        def seconds(now: datetime) -> int | None:
            return spec.cadence(self.feed, now.astimezone(self.tz))

        return seconds

    def _next_time(self, spec: JobSpec, now: datetime) -> datetime:
        due = next_due(spec, self.feed, now, self._last_end.get(spec.name))
        return due if due is not None else now + timedelta(seconds=IDLE_RECHECK_SECONDS)

    def reschedule(self, names: Iterable[str] | None = None) -> None:
        """Recompute the next run of every job (or of `names`) from the current settings."""
        if not self.started:
            return
        now = self.now()
        todo = list(names) if names is not None else list(self.jobs)
        for name in todo:
            spec = self.jobs.get(name)
            if spec is None:
                continue
            with contextlib.suppress(Exception):  # JobLookupError during shutdown
                self.aps.modify_job(name, next_run_time=self._next_time(spec, now))
        self._write_schedule(todo)

    def run_soon(self, name: str, trigger: str = "request") -> None:
        """Move a job to now. A job that is running now runs again right after."""
        if name not in self.jobs:
            log.warning("ignoring a run request for unknown job %r", name[:60])
            return
        with self._lock:
            if name in self._running:
                self._again[name] = trigger
                return
            self._next_trigger[name] = trigger
        if self.started:
            with contextlib.suppress(Exception):
                self.aps.modify_job(name, next_run_time=self.now())

    def after_run(self, name: str, status: str, pending: bool) -> None:
        """Called by run_job when a run ends: chain, run again for new requests, else time the next run from now."""
        self._last_end[name] = self.now()
        spec = self.jobs[name]
        with self._lock:
            again = self._again.pop(name, None)
        if status in ("ok", "degraded"):
            for c in spec.chain:
                self.run_soon(c, "chain")
        if pending or again:
            self.run_soon(name, again or "request")
            self._write_schedule([name])
        else:
            self.reschedule([name])

    def _fire(self, name: str) -> None:
        with self._lock:
            trigger = self._next_trigger.pop(name, "schedule")
        spec = self.jobs[name]
        if not self._db_ready and not self._setup_db():
            log.warning("job %s not run: database not reachable", name)
            return
        if trigger == "schedule" and spec.cadence(self.feed, self.now()) is None:
            self._write_schedule([name])
            return
        try:
            runs.run_job(self, name, trigger)
        except Exception as e:
            log.error("job %s could not run: %s", name, redact(f"{type(e).__name__}: {e}"))

    def _write_schedule(self, names: Iterable[str]) -> None:
        """job_schedule rows for the Admin and Clerical screens. Best effort."""
        if not self.started:
            return
        rows = []
        now = self.now()
        for name in names:
            spec = self.jobs.get(name)
            if spec is None:
                continue
            job = self.aps.get_job(name)
            nxt = job.next_run_time if job is not None else None
            secs = spec.cadence(self.feed, now)
            if secs is None:
                nxt = None
            rows.append((name, nxt, secs, bool(spec.pausable and self.feed.paused)))
        if not rows:
            return
        try:
            with self.connect() as conn:
                for r in rows:
                    conn.execute(
                        "INSERT INTO job_schedule (job, next_run_at, interval_seconds, paused, updated_at)"
                        " VALUES (%s, %s, %s, %s, now()) ON CONFLICT (job) DO UPDATE SET"
                        " next_run_at = EXCLUDED.next_run_at, interval_seconds = EXCLUDED.interval_seconds,"
                        " paused = EXCLUDED.paused, updated_at = now()",
                        r,
                    )
        except Exception as e:
            log.warning("job_schedule not written: %s", type(e).__name__)

    # ------------------------------------------------ settings

    def reload_settings(self, conn: Any | None = None) -> None:
        try:
            if conn is not None:
                self.feed = config.load(conn)
            else:
                with self.connect() as c:
                    self.feed = config.load(c)
        except Exception as e:
            log.warning("feed settings not reloaded: %s", type(e).__name__)
            return
        self.reschedule()

    # ------------------------------------------------ listener

    def _on_connect(self, conn: Any) -> None:
        """After every (re)connect: catch up on what was missed while not listening."""
        self._setup_db()
        self.reload_settings(conn)
        for (job,) in conn.execute("SELECT DISTINCT job FROM job_requests WHERE claimed_at IS NULL").fetchall():
            self.run_soon(job, "request")

    def _notified(self, channel: str, payload: str) -> None:
        if channel == "ps_settings":
            self.reload_settings()
            return
        try:
            data = json.loads(payload)
        except ValueError:
            log.warning("ignoring a non-JSON ps_job_request payload")
            return
        job = data.get("job") if isinstance(data, dict) else None
        if isinstance(job, str):
            self.run_soon(job, "request")

    def _listen(self) -> None:
        delay = 1.0
        while not self._stop.is_set():
            conn = None
            try:
                conn = self.connect()
                for ch in CHANNELS:
                    conn.execute(f"LISTEN {ch}")
                self.listening = True
                delay = 1.0
                log.info("worker listening on %s", ", ".join(CHANNELS))
                self._on_connect(conn)
                while not self._stop.is_set():
                    for n in conn.notifies(timeout=1.0):
                        self._notified(n.channel, n.payload)
            except Exception as e:
                if not self._stop.is_set():
                    log.warning("worker listener: %s; retrying in %.0f s", type(e).__name__, delay)
            finally:
                self.listening = False
                if conn is not None:
                    with contextlib.suppress(Exception):
                        conn.close()
            if self._stop.wait(delay):
                break
            delay = min(delay * 2, BACKOFF_MAX)

    # ------------------------------------------------ heartbeat loop (main thread)

    def serve(self, beat_file: Path, stop: threading.Event, interval: float) -> None:
        """Beat until `stop`: the beat file and the worker_heartbeat row."""
        conn = None
        while not stop.is_set():
            try:
                heartbeat.touch(beat_file)
            except OSError as e:
                log.error("can't touch beat file %s: %s", beat_file, e.strerror)
            try:
                if conn is None or conn.closed:
                    conn = self.connect()
                heartbeat.write_row(conn, build=self.build, host=self.host, pid=self.pid, running=self.running())
            except Exception as e:
                log.warning("heartbeat row not written: %s", type(e).__name__)
                if conn is not None:
                    with contextlib.suppress(Exception):
                        conn.close()
                conn = None
            stop.wait(interval)
        if conn is not None:
            with contextlib.suppress(Exception):
                conn.close()