"""tsayt: yet another asyncio cron scheduling micro-library for Python."""

from __future__ import annotations

import asyncio
import enum
from collections.abc import Callable, Coroutine
from datetime import datetime, timezone
from datetime import datetime as _datetime
from typing import Any

from cronsim import CronSim
from cronsim import CronSimError as _CronSimError

__all__ = ["Scheduler", "InvalidExpressionError", "OverlapPolicy"]


class InvalidExpressionError(ValueError):
    """Raised when an invalid cron expression is passed to
    :meth:`~tsayt.Scheduler.register`.

    This is a :exc:`ValueError` subclass.  The original parsing error from
    the underlying cron library is available as the ``__cause__``.
    """


class OverlapPolicy(enum.Enum):
    """Controls what happens when a scheduled task is still running at the
    time its next firing is due.

    .. attribute:: SKIP

        If the previous invocation is still in flight, the new firing is
        silently skipped.  The task will next be considered at its following
        scheduled time.  This is the default and is the safer choice for
        non-reentrant tasks (database operations, file I/O, anything that
        holds a lock or maintains state).

    .. attribute:: ALLOW

        A new invocation is always started on schedule, regardless of whether
        a previous one is still running.  Choose this for fully independent,
        stateless work where overlapping runs are harmless.
    """

    SKIP = "skip"
    ALLOW = "allow"


# A zero-argument async callable that returns nothing of interest.
Task = Callable[[], Coroutine[Any, Any, Any]]

# Called with the exception when a scheduled task raises.
ErrorHandler = Callable[[BaseException], Coroutine[Any, Any, Any]]


def _now() -> datetime:
    """Return the current UTC time. Exists as a named module-level function
    so that tests can patch ``tsayt._now`` with a clean, typed seam instead
    of replacing the ``datetime`` class itself on the module."""
    return datetime.now(tz=timezone.utc)


class _Registration:
    """Internal record of a registered cron task."""

    def __init__(self, expr: str, func: Task, overlap: OverlapPolicy) -> None:
        self.expr = expr
        self.func = func
        self.overlap = overlap
        self.in_flight: asyncio.Task[None] | None = None

    def next_after(self, dt: datetime) -> datetime:
        """Return the next scheduled datetime strictly after *dt*.

        CronSim always ticks forward by at least one second before resolving
        the next match, so the returned value is guaranteed to be > *dt*.
        """
        return next(iter(CronSim(self.expr, dt)))


class Scheduler:
    """A simple asyncio cron scheduler.

    Usage::

        scheduler = tsayt.Scheduler()
        scheduler.register("0 0 * * *", my_async_func)
        await scheduler.run()

    Args:
        error_handler: An optional async callable that receives any exception
            raised by a scheduled task.  If not provided, exceptions are
            re-raised into the abandoned ``asyncio.Task`` and will produce a
            ``"Task exception was never retrieved"`` warning from the event
            loop.

    Clock-skew / suspension note:
        The scheduler sleeps until the next scheduled time using
        ``asyncio.sleep``.  If the process is suspended and resumed (e.g. a
        laptop waking from sleep) or the system clock is stepped backward by
        NTP, the scheduler will behave as follows:

        * **Forward jump** (clock ahead, or process wakes late): every task
          whose window was passed during the gap fires once on the next loop
          iteration.  Missed intermediate firings are not back-filled.
        * **Backward jump** (clock stepped back): the scheduler will
          re-sleep for the remaining interval and fire once when the original
          target time is reached again.

        There is currently no configurable catch-up or skip policy.
    """

    def __init__(self, error_handler: ErrorHandler | None = None) -> None:
        self._registrations: list[_Registration] = []
        self._error_handler = error_handler

    def register(
        self,
        expr: str,
        func: Task,
        overlap: OverlapPolicy = OverlapPolicy.SKIP,
    ) -> None:
        """Register *func* to be called according to the cron expression *expr*.

        Args:
            expr: A standard 5-field cron expression (e.g. ``"0 0 * * *"``).
            func: A zero-argument async callable.
            overlap: What to do when *func* is still running at the time its next firing is due. Defaults to :attr:`~tsayt.OverlapPolicy.SKIP`, which silently skips the new firing. Pass :attr:`~tsayt.OverlapPolicy.ALLOW` to always start a new invocation on schedule.

        Raises:
            InvalidExpressionError: If *expr* is not a valid cron expression.
        """
        # Validate the expression eagerly by constructing a CronSim against an
        # arbitrary epoch and immediately discarding it.  Any parse error is
        # caught and re-raised as InvalidExpressionError so callers never need
        # to import from cronsim directly.
        try:
            CronSim(expr, _datetime(2000, 1, 1, tzinfo=timezone.utc))
        except _CronSimError as e:
            raise InvalidExpressionError(str(e)) from e
        self._registrations.append(_Registration(expr, func, overlap))

    async def _run_task(self, func: Task) -> None:
        """Invoke *func*, routing any exception to the error handler."""
        try:
            await func()
        except Exception as exc:
            if self._error_handler is not None:
                await self._error_handler(exc)
            else:
                raise

    async def run(self) -> None:
        """Run the scheduler indefinitely, firing registered tasks on schedule.

        Each task is launched as an independent ``asyncio.Task`` so that
        long-running coroutines do not delay the firing of other scheduled
        jobs.

        Tasks registered *after* ``run()`` has started will not be picked up;
        call ``register()`` before starting the scheduler.
        """
        if not self._registrations:
            return

        now = _now()

        # Map each registration to its next scheduled fire time.
        schedule: list[tuple[datetime, _Registration]] = [
            (reg.next_after(now), reg) for reg in self._registrations
        ]

        while True:
            # Find the earliest next fire time across all registrations.
            next_dt, _ = min(schedule, key=lambda item: item[0])

            now = _now()
            delay = (next_dt - now).total_seconds()
            if delay > 0:
                await asyncio.sleep(delay)

            # Fire every registration whose scheduled time has arrived.
            fired_now = _now()
            new_schedule: list[tuple[datetime, _Registration]] = []
            for fire_dt, reg in schedule:
                if fire_dt <= fired_now:
                    should_fire = (
                        reg.overlap is OverlapPolicy.ALLOW
                        or reg.in_flight is None
                        or reg.in_flight.done()
                    )
                    if should_fire:
                        reg.in_flight = asyncio.ensure_future(self._run_task(reg.func))
                    new_schedule.append((reg.next_after(fired_now), reg))
                else:
                    new_schedule.append((fire_dt, reg))
            schedule = new_schedule
