AirlockTools/Server/scheduler_async.py

#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.

import asyncio
import json
import logging
import os
from typing import Any, Callable, Dict, List

logger = logging.getLogger(__name__)

# Registry of functions that can be scheduled
FUNCTION_MAP: Dict[str, Callable] = {}

# Dictionary to manually track scheduled jobs by ID
scheduled_jobs: Dict[str, asyncio.TimerHandle] = {}

# Path to the JSON file for job persistence TODO - pin this to the correct place
JOBS_FILE = os.path.join(os.getcwd(), "jobs.json")


def register_function(name: str, func: Callable):
    """
    Register a function so it can be called by name later.
    Example:
        register_function("say_hello", say_hello)
    """
    FUNCTION_MAP[name] = func


def load_jobs() -> List[Dict[str, Any]]:
    """
    Load jobs from the JSON file, or return [] if none exist.
    """
    if not os.path.exists(JOBS_FILE):
        return []
    with open(JOBS_FILE, "r") as f:
        return json.load(f)


def save_jobs(jobs: List[Dict[str, Any]]):
    """
    Save jobs to the JSON file (overwrite).
    """
    with open(JOBS_FILE, "w") as f:
        json.dump(jobs, f, indent=4)


def cancel_job(job_id: str):
    """
    Cancel a scheduled job by ID and remove it from the registry and persistence.
    """
    handle = scheduled_jobs.pop(job_id, None)
    if handle:
        handle.cancel()
        logger.info(f"Cancelled job '{job_id}'")

    jobs = [j for j in load_jobs() if j.get("id") != job_id]
    save_jobs(jobs)


def run_once_job(
    job_id: str,
    func_name: str,
    delay_seconds: float,
    args=None,
    kwargs=None,
    persist=True,
):
    """
    Schedule a job to run once after a delay (in seconds).
    """
    args = args or []
    kwargs = kwargs or {}

    def job_wrapper():
        func = FUNCTION_MAP.get(func_name)
        if func is None:
            logger.error(f"Function '{func_name}' is not registered.")
            return
        func(*args, **kwargs)
        cancel_job(job_id)

    loop = asyncio.get_event_loop()
    handle = loop.call_later(delay_seconds, job_wrapper)
    scheduled_jobs[job_id] = handle

    if persist:
        jobs = [j for j in load_jobs() if j.get("id") != job_id]
        jobs.append(
            {
                "id": job_id,
                "type": "once",
                "delay": delay_seconds,
                "function": func_name,
                "args": args,
                "kwargs": kwargs,
            }
        )
        save_jobs(jobs)
        logger.info(
            f"Scheduled one-time job '{job_id}' to run in {delay_seconds} seconds."
        )


def recurring_job(
    job_id: str, func_name: str, interval: float, args=None, kwargs=None, persist=True
):
    """
    Schedule a recurring job.
    """
    args = args or []
    kwargs = kwargs or {}

    def job_wrapper():
        func = FUNCTION_MAP.get(func_name)
        if func is None:
            logger.error(f"Function '{func_name}' is not registered.")
            return
        func(*args, **kwargs)
        # Reschedule the job
        handle = asyncio.get_event_loop().call_later(interval, job_wrapper)
        scheduled_jobs[job_id] = handle

    cancel_job(job_id)
    handle = asyncio.get_event_loop().call_later(interval, job_wrapper)
    scheduled_jobs[job_id] = handle

    if persist:
        jobs = [j for j in load_jobs() if j.get("id") != job_id]
        jobs.append(
            {
                "id": job_id,
                "type": "recurring",
                "interval": interval,
                "function": func_name,
                "args": args,
                "kwargs": kwargs,
            }
        )
        save_jobs(jobs)
        logger.info(f"Scheduled recurring job '{job_id}' every {interval} seconds.")


def reload_jobs():
    """
    Reload jobs from JSON and reschedule them.
    """
    jobs = load_jobs()
    for job in jobs:
        if job["type"] == "once":
            run_once_job(
                job["id"],
                job["function"],
                job["delay"],
                job.get("args"),
                job.get("kwargs"),
                persist=False,
            )
        elif job["type"] == "recurring":
            recurring_job(
                job["id"],
                job["function"],
                job["interval"],
                job.get("args"),
                job.get("kwargs"),
                persist=False,
            )


async def start_scheduler():
    """
    Start the asynchronous scheduler loop.

    This function is a placeholder to keep the event loop alive.
    Jobs are scheduled using asyncio.call_later and do not require polling.
    """
    try:
        await asyncio.Event().wait()
    except asyncio.CancelledError:
        logger.critical("Scheduler stopped.")

    """
    Start the asynchronous scheduler loop.

    This function is a placeholder for compatibility. Since we use asyncio.call_later,
    jobs are scheduled directly on the event loop and no polling is required.

    Usage:
        # In an async app (e.g., Textual)
        asyncio.create_task(start_scheduler())

        # Or in a standalone script
        async def main():
            await start_scheduler()

        asyncio.run(main())
    """
    try:
        while True:
            await asyncio.sleep(3600)  # Sleep indefinitely; jobs run via call_later
    except asyncio.CancelledError:
        logger.critical("Scheduler stopped.")