cron_system.py

"""Cron-style job scheduler for the Wcode agent.

Supports both one-shot and recurring jobs with optional disk persistence.
A background daemon thread evaluates cron expressions every 5 seconds and
enqueues jobs that are due.  The agent loop consumes the queue each turn.
"""

import json
import time
import random
import threading
from pathlib import Path
from datetime import datetime
from dataclasses import dataclass, asdict

from config import get_logger
from wcode_errors import CronError

logger = get_logger(__name__)


@dataclass
class CronJob:
    id: str
    cron: str        # "0 9 * * *"
    prompt: str      # message to inject when fired
    recurring: bool  # True = recurring, False = one-shot
    durable: bool    # True = persist to disk


class CronScheduler:
    """Manages scheduled cron jobs with optional disk persistence.

    A background daemon thread evaluates cron expressions every 5 seconds.
    The agent loop consumes fired jobs each turn via ``consume_queue()``.

    Usage::

        cron = CronScheduler(Path(".cron/scheduled_tasks.json"))
        cron.start()
        cron.schedule("0 9 * * *", "daily standup")
        ...
        for job in cron.consume_queue():
            messages.append({"role": "user", "content": job.prompt})
    """

    def __init__(self, durable_path: Path | None = None) -> None:
        if durable_path is None:
            durable_path = Path.cwd() / ".cron" / "scheduled_tasks.json"
        self.durable_path = durable_path

        self._jobs: dict[str, CronJob] = {}
        self._queue: list[CronJob] = []
        self._lock = threading.Lock()
        self.agent_lock = threading.Lock()
        self._last_fired: dict[str, str] = {}

    # ── Lifecycle ──────────────────────────────────────────────────────

    def start(self) -> None:
        """Load durable jobs from disk and start the scheduler daemon thread."""
        self._load_durable()
        threading.Thread(target=self._scheduler_loop, daemon=True).start()
        logger.info("[cron] scheduler thread started")

    # ── Job management ─────────────────────────────────────────────────

    def schedule(
        self, cron: str, prompt: str, recurring: bool = True, durable: bool = True
    ) -> CronJob:
        """Schedule a new cron job.

        Args:
            cron: 5-field cron expression.
            prompt: Message text to inject when the job fires.
            recurring: True for recurring, False for one-shot.
            durable: True to persist across restarts.

        Returns:
            ``CronJob`` on success.

        Raises:
            CronError: If the cron expression is invalid.
        """
        validate_cron(cron)
        job = CronJob(
            id=f"cron_{random.randint(0, 999999):06d}",
            cron=cron,
            prompt=prompt,
            recurring=recurring,
            durable=durable,
        )
        with self._lock:
            self._jobs[job.id] = job
        if job.durable:
            self._save_durable()
        logger.info(f"[cron register] {job.id} '{cron}' → {prompt[:40]}")
        return job

    def cancel(self, job_id: str) -> str:
        """Cancel a scheduled job by ID."""
        with self._lock:
            job = self._jobs.pop(job_id, None)
        if not job:
            return f"Job {job_id} not found"
        if job.durable:
            self._save_durable()
        logger.info(f"[cron cancel] {job_id}")
        return f"Cancelled {job_id}"

    def list_jobs(self) -> list[CronJob]:
        """Return all currently registered jobs."""
        with self._lock:
            return list(self._jobs.values())

    # ── Queue interface (called by agent loop) ─────────────────────────

    def consume_queue(self) -> list[CronJob]:
        """Atomically drain and return the pending job queue."""
        with self._lock:
            fired = list(self._queue)
            self._queue.clear()
        return fired

    def has_queue(self) -> bool:
        """Return True if there are pending jobs to deliver."""
        with self._lock:
            return bool(self._queue)

    # ── Persistence ────────────────────────────────────────────────────

    def _save_durable(self) -> None:
        durable = [asdict(j) for j in self._jobs.values() if j.durable]
        self.durable_path.parent.mkdir(parents=True, exist_ok=True)
        self.durable_path.write_text(json.dumps(durable, indent=2))

    def _load_durable(self) -> None:
        if not self.durable_path.exists():
            return
        jobs = json.loads(self.durable_path.read_text())
        for j in jobs:
            job = CronJob(**j)
            try:
                validate_cron(job.cron)
            except CronError as e:
                logger.error(f"[cron] skipping invalid job {job.id}: {e}")
                continue
            self._jobs[job.id] = job
        valid = [j for j in jobs if j["id"] in self._jobs]
        if valid:
            logger.info(f"[cron] loaded {len(valid)} durable job(s)")

    # ── Scheduler loop ─────────────────────────────────────────────────

    def _scheduler_loop(self) -> None:
        while True:
            time.sleep(5)
            now = datetime.now()
            minute_marker = now.strftime("%Y-%m-%d %H:%M")
            with self._lock:
                for job in list(self._jobs.values()):
                    if (
                        not cron_matches(job.cron, now)
                        or self._last_fired.get(job.id) == minute_marker
                    ):
                        continue
                    self._queue.append(job)
                    self._last_fired[job.id] = minute_marker
                    logger.info(f"[cron fire] {job.id} → {job.prompt[:40]}")
                    if not job.recurring:
                        self._jobs.pop(job.id)
                        if job.durable:
                            self._save_durable()


# ── Cron expression helpers (module-level, stateless) ───────────────────


def _cron_field_matches(field: str, value: int) -> bool:
    if field == "*":
        return True
    if field.startswith("*/"):
        step = int(field[2:])
        return step > 0 and value % step == 0
    if "," in field:
        return any(_cron_field_matches(f.strip(), value) for f in field.split(","))
    if "-" in field:
        lo, hi = field.split("-", 1)
        return int(lo) <= value <= int(hi)
    return value == int(field)


def cron_matches(cron_expr: str, dt: datetime) -> bool:
    fields = cron_expr.strip().split()
    if len(fields) != 5:
        return False
    minute, hour, dom, month, dow = fields
    dow_val = (dt.weekday() + 1) % 7  # Python Monday=0 → cron Sunday=0

    m = _cron_field_matches(minute, dt.minute)
    h = _cron_field_matches(hour, dt.hour)
    dom_ok = _cron_field_matches(dom, dt.day)
    month_ok = _cron_field_matches(month, dt.month)
    dow_ok = _cron_field_matches(dow, dow_val)
    if not (m and h and month_ok):
        return False
    dom_unconstrained = dom == "*"
    dow_unconstrained = dow == "*"
    if dom_unconstrained and dow_unconstrained:
        return True
    if dom_unconstrained:
        return dow_ok
    if dow_unconstrained:
        return dom_ok
    return dom_ok or dow_ok


def _validate_cron_field(field: str, lo: int, hi: int) -> str | None:
    """Validate a single cron field value is within [lo, hi]."""
    if field == "*":
        return None
    if field.startswith("*/"):
        step_str = field[2:]
        if not step_str.isdigit():
            return f"Invalid step: {field}"
        step = int(step_str)
        if step <= 0:
            return f"Step must be > 0: {field}"
        return None
    if "," in field:
        for part in field.split(","):
            err = _validate_cron_field(part.strip(), lo, hi)
            if err:
                return err
        return None
    if "-" in field:
        parts = field.split("-", 1)
        if not parts[0].isdigit() or not parts[1].isdigit():
            return f"Invalid range: {field}"
        a, b = int(parts[0]), int(parts[1])
        if a < lo or a > hi or b < lo or b > hi:
            return f"Range {field} out of bounds [{lo}-{hi}]"
        if a > b:
            return f"Range start > end: {field}"
        return None
    if not field.isdigit():
        return f"Invalid field: {field}"
    val = int(field)
    if val < lo or val > hi:
        return f"Value {val} out of bounds [{lo}-{hi}]"
    return None


def validate_cron(cron_expr: str) -> None:
    """Validate a 5-field cron expression.

    Raises :class:`CronError` if the expression is malformed.
    """
    fields = cron_expr.strip().split()
    if len(fields) != 5:
        raise CronError(f"Expected 5 fields, got {len(fields)}")
    bounds = [(0, 59), (0, 23), (1, 31), (1, 12), (0, 6)]
    names = ["minute", "hour", "day-of-month", "month", "day-of-week"]
    for i, (field, (lo, hi), name) in enumerate(zip(fields, bounds, names)):
        err = _validate_cron_field(field, lo, hi)
        if err:
            raise CronError(f"{name}: {err}")


# ── Module-level shims (delegate to a default scheduler) ────────────────

_default_scheduler: CronScheduler | None = None


def _get_default() -> CronScheduler:
    global _default_scheduler
    if _default_scheduler is None:
        _default_scheduler = CronScheduler()
    return _default_scheduler


def init_cron_system() -> None:
    _get_default().start()


def schedule_job(cron: str, prompt: str, recurring: bool = True, durable: bool = True) -> CronJob:
    return _get_default().schedule(cron, prompt, recurring, durable)


def cancel_job(job_id: str) -> str:
    return _get_default().cancel(job_id)


def list_jobs() -> list[CronJob]:
    return _get_default().list_jobs()


def consume_cron_queue() -> list[CronJob]:
    return _get_default().consume_queue()


def has_cron_queue() -> bool:
    return _get_default().has_queue()


# agent_lock exposed for backward compatibility
def _get_agent_lock() -> threading.Lock:
    return _get_default().agent_lock