"""Reference solution — In-Memory Job Scheduler.

DO NOT READ BEFORE YOU HAVE RUN THE PROBLEM UNDER THE CLOCK.
Teaching text: ../../../WARMUP.md § Chapter 5.

Three decisions that carry the whole problem:

1. THE TIE-BREAKER. Pushing (fire_at, job) into a heap means heapq compares
   the payloads when two jobs share an instant -- TypeError if they are not
   comparable, or worse, silent ordering by job CONTENTS if they are. Push
   (fire_at, priority, sequence, job_id) so ties resolve deterministically and
   the payload is never compared.

2. LAZY CANCELLATION. heapq has no remove, and finding an element is O(n). So
   mark the id cancelled and skip it when it surfaces -- O(1). The failure mode
   is tombstone accumulation when a workload schedules far ahead and cancels
   most of it, so rebuild once cancellations exceed half the heap. That bounds
   the waste at 2x and is amortised O(1). It is the same pattern asyncio's loop
   uses for cancelled timers.

3. CATCH-UP IS A PRODUCT DECISION. A fixed-rate job that has fallen behind must
   not silently decide whether to fire every missed occurrence. Surface it.
"""

from __future__ import annotations

import heapq
import itertools
import random
import time


class Scheduler:
    """Delayed and recurring execution with O(1) cancel, retries, fairness."""

    def __init__(self, clock=time.monotonic, rng=None, starvation_passes=3):
        self._heap = []                 # (fire_at, priority, seq, job_id)
        self._jobs = {}                 # job_id -> record
        self._cancelled = set()
        self._seq = itertools.count()
        self._ids = itertools.count(1)
        self._clock = clock
        self._rng = rng if rng is not None else random.Random(0)
        self._starvation_passes = starvation_passes
        self._pass = 0
        self._draining = False
        self.dead_letters = []

    # ---- scheduling ------------------------------------------------------
    def schedule(self, fn, delay=0.0, *, period=None, mode="fixed_delay",
                 catch_up="run_latest_only", max_attempts=5,
                 base_backoff=0.2, cap_backoff=30.0,
                 priority=0, tenant=None, max_concurrency=None):
        if self._draining:
            raise RuntimeError("scheduler is draining; not accepting new jobs")
        if period is not None:
            if mode not in ("fixed_rate", "fixed_delay"):
                raise ValueError("mode must be fixed_rate or fixed_delay")
            if catch_up not in ("run_all", "run_latest_only", "skip"):
                raise ValueError("unknown catch_up policy")
            if period <= 0:
                raise ValueError("period must be positive")
        if max_attempts < 1:
            raise ValueError("max_attempts must be >= 1")

        job_id = next(self._ids)
        self._jobs[job_id] = {
            "fn": fn, "period": period, "mode": mode, "catch_up": catch_up,
            "max_attempts": max_attempts, "attempt": 0,
            "base": base_backoff, "cap": cap_backoff,
            "priority": priority, "base_priority": priority,
            "tenant": tenant, "max_concurrency": max_concurrency,
            "waiting_since": None,
        }
        self._push(self._clock() + delay, job_id)
        return job_id

    def _push(self, fire_at, job_id):
        job = self._jobs[job_id]
        heapq.heappush(
            self._heap, (fire_at, job["priority"], next(self._seq), job_id))

    def cancel(self, job_id):
        if job_id not in self._jobs:
            return False
        del self._jobs[job_id]
        self._cancelled.add(job_id)
        self._maybe_rebuild()
        return True

    def _maybe_rebuild(self):
        if len(self._cancelled) > max(32, len(self._heap) // 2):
            self._heap = [e for e in self._heap if e[3] not in self._cancelled]
            heapq.heapify(self._heap)
            self._cancelled.clear()

    def drain(self):
        """Refuse new work; already-due work still runs."""
        self._draining = True

    @property
    def draining(self):
        return self._draining

    @property
    def pending(self):
        return len(self._jobs)

    # ---- running ---------------------------------------------------------
    def _backoff(self, job):
        # Full jitter: uniform over [0, min(cap, base * 2**attempt)]. Fixed
        # backoff synchronises every client that failed at the same instant.
        ceiling = min(job["cap"], job["base"] * (2 ** job["attempt"]))
        return self._rng.uniform(0, ceiling)

    def next_fire_time(self):
        while self._heap and self._heap[0][3] in self._cancelled:
            self._cancelled.discard(heapq.heappop(self._heap)[3])
        return self._heap[0][0] if self._heap else None

    def run_due(self, limit=None):
        """Run everything due at the current clock. Returns (ran, errors)."""
        now = self._clock()
        self._pass += 1
        ran, errors = 0, []
        deferred = []                   # (fire_at, job_id) held by a tenant cap
        per_tenant = {}

        while self._heap and (limit is None or ran < limit):
            fire_at, _prio, _seq, job_id = self._heap[0]
            if fire_at > now:
                break
            heapq.heappop(self._heap)

            if job_id in self._cancelled:
                self._cancelled.discard(job_id)
                continue
            job = self._jobs.get(job_id)
            if job is None:
                continue

            # Per-tenant concurrency cap: one tenant may not consume the whole
            # pass. Defer the rest — they stay due for the next call.
            tenant, cap = job["tenant"], job["max_concurrency"]
            if tenant is not None and cap is not None:
                if per_tenant.get(tenant, 0) >= cap:
                    deferred.append((fire_at, job_id))
                    continue
                per_tenant[tenant] = per_tenant.get(tenant, 0) + 1

            job["priority"] = job["base_priority"]     # reset any promotion
            job["waiting_since"] = None

            try:
                job["fn"]()
            except Exception as exc:
                errors.append((job_id, exc))
                job["attempt"] += 1
                if job["attempt"] < job["max_attempts"]:
                    self._push(self._clock() + self._backoff(job), job_id)
                else:
                    self.dead_letters.append((job_id, job["attempt"], exc))
                    del self._jobs[job_id]
                continue

            ran += 1
            job["attempt"] = 0
            period = job["period"]
            if period is None:
                del self._jobs[job_id]
                continue

            if job["mode"] == "fixed_delay":
                self._push(self._clock() + period, job_id)
            else:                                       # fixed_rate
                nxt = fire_at + period
                if nxt <= now:                          # we overran
                    if job["catch_up"] == "run_latest_only":
                        missed = int((now - nxt) // period) + 1
                        nxt += missed * period
                    elif job["catch_up"] == "skip":
                        nxt = now + period
                    # "run_all" leaves nxt in the past: it fires again
                    # immediately, once per missed occurrence, until caught up.
                self._push(nxt, job_id)

        # Re-arm the tenant-capped jobs, promoting anything that has waited too
        # long so low priority cannot starve indefinitely.
        for fire_at, job_id in deferred:
            job = self._jobs.get(job_id)
            if job is None:
                continue
            if job["waiting_since"] is None:
                job["waiting_since"] = self._pass
            elif self._pass - job["waiting_since"] >= self._starvation_passes:
                job["priority"] = job["base_priority"] - 1     # promote
                job["waiting_since"] = self._pass
            self._push(fire_at, job_id)

        return ran, errors
