"""Schedule evaluation — computes next run times."""
from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone

from croniter import croniter

from crondag.sdk.schedule import Cron, Interval, Schedule

logger = logging.getLogger(__name__)


def next_run_time(schedule: Schedule, after: datetime | None = None) -> datetime | None:
    """Compute the next run time for a schedule.

    Args:
        schedule: Cron or Interval schedule.
        after: Compute next time after this datetime. Defaults to now(UTC).

    Returns:
        Next datetime in UTC, or None if no schedule.
    """
    if schedule is None:
        return None

    if after is None:
        after = datetime.now(timezone.utc)

    # Ensure after is timezone-aware
    if after.tzinfo is None:
        after = after.replace(tzinfo=timezone.utc)

    if isinstance(schedule, Cron):
        cron = croniter(schedule.expression, after)
        next_dt = cron.get_next(datetime)
        if next_dt.tzinfo is None:
            next_dt = next_dt.replace(tzinfo=timezone.utc)
        return next_dt

    if isinstance(schedule, Interval):
        return after + schedule.delta

    return None


def is_due(schedule: Schedule, next_run_at: datetime | None, now: datetime | None = None) -> bool:
    """Check if a DAG is due to run.

    Args:
        schedule: The DAG's schedule.
        next_run_at: The precomputed next run time.
        now: Current time (defaults to now UTC).
    """
    if schedule is None or next_run_at is None:
        return False

    if now is None:
        now = datetime.now(timezone.utc)

    if now.tzinfo is None:
        now = now.replace(tzinfo=timezone.utc)
    if next_run_at.tzinfo is None:
        next_run_at = next_run_at.replace(tzinfo=timezone.utc)

    return now >= next_run_at
