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.")