PhoneWork/agent/scheduler.py

"""Simple scheduler for reminders and recurring tasks."""

from __future__ import annotations

import asyncio
import json
import logging
import time
import uuid
from dataclasses import dataclass, asdict
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Dict, Optional

logger = logging.getLogger(__name__)

PERSISTENCE_FILE = Path(__file__).parent.parent / "scheduled_jobs.json"


class JobStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    CANCELLED = "cancelled"


@dataclass
class ScheduledJob:
    job_id: str
    description: str
    scheduled_at: float
    delay_seconds: float
    status: JobStatus = JobStatus.PENDING
    is_recurring: bool = False
    interval_seconds: Optional[float] = None
    max_runs: int = 1
    runs_completed: int = 0
    notify_chat_id: Optional[str] = None

    def to_dict(self) -> dict:
        data = asdict(self)
        data["status"] = self.status.value
        return data

    @classmethod
    def from_dict(cls, data: dict) -> "ScheduledJob":
        data = data.copy()
        data["status"] = JobStatus(data.get("status", "pending"))
        return cls(**data)


class Scheduler:
    """Singleton that manages scheduled jobs with Feishu notifications."""

    def __init__(self) -> None:
        self._jobs: dict[str, ScheduledJob] = {}
        self._tasks: dict[str, asyncio.Task] = {}
        self._lock = asyncio.Lock()
        self._started = False

    async def start(self) -> None:
        """Load persisted jobs and reschedule pending ones."""
        self._load()
        self._started = True
        now = time.time()
        for job in list(self._jobs.values()):
            if job.status == JobStatus.PENDING:
                elapsed = now - job.scheduled_at
                if job.is_recurring:
                    remaining = job.interval_seconds - (elapsed % job.interval_seconds) if job.interval_seconds else 0
                    task = asyncio.create_task(self._run_recurring(job.job_id, skip_initial=True, initial_delay=remaining))
                    self._tasks[job.job_id] = task
                else:
                    remaining = max(0, job.delay_seconds - elapsed)
                    if remaining <= 0:
                        asyncio.create_task(self._run_once(job.job_id))
                    else:
                        task = asyncio.create_task(self._run_once(job.job_id, initial_delay=remaining))
                        self._tasks[job.job_id] = task
        logger.info("Scheduler started with %d jobs", len(self._jobs))

    def _load(self) -> None:
        """Load jobs from persistence file."""
        if not PERSISTENCE_FILE.exists():
            return
        try:
            with open(PERSISTENCE_FILE, "r", encoding="utf-8") as f:
                data = json.load(f)
            for job_id, job_data in data.items():
                job = ScheduledJob.from_dict(job_data)
                if job.status not in (JobStatus.COMPLETED, JobStatus.CANCELLED):
                    self._jobs[job_id] = job
            logger.info("Loaded %d jobs from %s", len(self._jobs), PERSISTENCE_FILE)
        except Exception:
            logger.exception("Failed to load scheduled jobs")

    def _save(self) -> None:
        """Save jobs to persistence file."""
        try:
            data = {jid: job.to_dict() for jid, job in self._jobs.items()}
            with open(PERSISTENCE_FILE, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
        except Exception:
            logger.exception("Failed to save scheduled jobs")

    async def schedule_once(
        self,
        delay_seconds: float,
        message: str,
        notify_chat_id: Optional[str] = None,
    ) -> str:
        """Schedule a one-time reminder."""
        job_id = str(uuid.uuid4())[:8]
        job = ScheduledJob(
            job_id=job_id,
            description=message,
            scheduled_at=time.time(),
            delay_seconds=delay_seconds,
            notify_chat_id=notify_chat_id,
        )

        async with self._lock:
            self._jobs[job_id] = job
            self._save()

        task = asyncio.create_task(self._run_once(job_id))
        self._tasks[job_id] = task

        logger.info("Scheduled job %s: %s (in %ds)", job_id, message[:50], delay_seconds)
        return job_id

    async def schedule_recurring(
        self,
        interval_seconds: float,
        message: str,
        max_runs: int = 10,
        notify_chat_id: Optional[str] = None,
    ) -> str:
        """Schedule a recurring reminder."""
        job_id = str(uuid.uuid4())[:8]
        job = ScheduledJob(
            job_id=job_id,
            description=message,
            scheduled_at=time.time(),
            delay_seconds=interval_seconds,
            is_recurring=True,
            interval_seconds=interval_seconds,
            max_runs=max_runs,
            notify_chat_id=notify_chat_id,
        )

        async with self._lock:
            self._jobs[job_id] = job
            self._save()

        task = asyncio.create_task(self._run_recurring(job_id))
        self._tasks[job_id] = task

        logger.info("Scheduled recurring job %s: %s (every %ds, %d runs)", 
                    job_id, message[:50], interval_seconds, max_runs)
        return job_id

    async def _run_once(self, job_id: str, initial_delay: Optional[float] = None) -> None:
        """Execute a one-time job after delay."""
        job = self._jobs.get(job_id)
        if not job:
            return

        delay = initial_delay if initial_delay is not None else job.delay_seconds
        await asyncio.sleep(delay)

        async with self._lock:
            job.status = JobStatus.RUNNING
            self._save()

        await self._send_notification(job, job.description)

        async with self._lock:
            job.status = JobStatus.COMPLETED
            job.runs_completed = 1
            self._save()

        logger.info("Job %s completed", job_id)

    async def _run_recurring(self, job_id: str, skip_initial: bool = False, initial_delay: Optional[float] = None) -> None:
        """Execute a recurring job."""
        job = self._jobs.get(job_id)
        if not job:
            return

        interval = job.interval_seconds or 60

        if skip_initial and initial_delay is not None:
            await asyncio.sleep(initial_delay)

        for run in range(job.max_runs):
            if not (skip_initial and run == 0):
                await asyncio.sleep(interval)

            async with self._lock:
                if job.status == JobStatus.CANCELLED:
                    break
                job.status = JobStatus.RUNNING
                self._save()

            await self._send_notification(job, f"[{job.runs_completed + 1}/{job.max_runs}] {job.description}")

            async with self._lock:
                job.runs_completed += 1
                if job.runs_completed < job.max_runs:
                    job.status = JobStatus.PENDING
                else:
                    job.status = JobStatus.COMPLETED
                self._save()

        logger.info("Recurring job %s finished (%d runs)", job_id, job.runs_completed)

    async def _send_notification(self, job: ScheduledJob, message: str) -> None:
        """Send Feishu notification."""
        if not job.notify_chat_id:
            return

        from bot.feishu import send_text

        try:
            await send_text(job.notify_chat_id, "chat_id", f"⏰ **Reminder**\n\n{message}")
        except Exception:
            logger.exception("Failed to send notification for job %s", job.job_id)

    def cancel(self, job_id: str) -> bool:
        """Cancel a scheduled job."""
        job = self._jobs.get(job_id)
        if not job:
            return False

        job.status = JobStatus.CANCELLED
        self._save()
        task = self._tasks.get(job_id)
        if task:
            task.cancel()

        logger.info("Cancelled job %s", job_id)
        return True

    async def stop(self) -> None:
        """Stop all tasks and clear state."""
        for task in self._tasks.values():
            task.cancel()
        self._tasks.clear()
        async with self._lock:
            self._jobs.clear()
        if PERSISTENCE_FILE.exists():
            PERSISTENCE_FILE.unlink()
        logger.info("Scheduler stopped")

    def list_jobs(self, limit: int = 20) -> list[dict]:
        jobs = sorted(
            [j for j in self._jobs.values() if j.status != JobStatus.COMPLETED],
            key=lambda j: j.scheduled_at,
        )[:limit]
        return [
            {
                "job_id": j.job_id,
                "description": j.description[:50],
                "status": j.status.value,
                "is_recurring": j.is_recurring,
                "runs_completed": j.runs_completed,
                "max_runs": j.max_runs,
            }
            for j in jobs
        ]


scheduler = Scheduler()