#Standard Libary Imports:
import json
import os
import time
from typing import Callable, Any, List, Dict

#3rd Party Imports:
import schedule

# File where all jobs are persisted
JOBS_FILE = "scheduling\\jobs.json"

# Ensure directory exists
os.makedirs(os.path.dirname(JOBS_FILE), exist_ok=True)

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

# -------------------------------
# Function Registration
# -------------------------------

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

# -------------------------------
# Persistence Helpers
# -------------------------------

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 _atomic_save(path: str, data: Any):
    """Write JSON atomically to avoid partial writes."""
    tmp = f"{path}.tmp"
    with open(tmp, "w") as f:
        json.dump(data, f, indent=4)
    os.replace(tmp, path)

def save_jobs(jobs: List[Dict[str, Any]]):
    """Save jobs to the JSON file (overwrite)."""
    _atomic_save(JOBS_FILE, jobs)

# -------------------------------
# Uniqueness Helpers
# -------------------------------

def job_in_store(job_id: str) -> bool:
    """Check if a job id exists in the persisted JSON file."""
    return any(j.get("id") == job_id for j in load_jobs())

def job_in_scheduler(job_id: str) -> bool:
    """
    Check if a job with this tag exists in the in-memory scheduler.
    Uses schedule.get_jobs(tag=...) if available, otherwise scans tags.
    """
    try:
        jobs = schedule.get_jobs(tag=job_id)  # schedule >= 1.2.0
        return len(jobs) > 0
    except TypeError:
        # Fallback for older versions
        return any(job_id in getattr(j, "tags", set()) for j in schedule.jobs)

def ensure_unique(job_id: str, on_conflict: str = "skip") -> bool:
    """
    Ensure the job_id is unique across persistence and in-memory schedule.
    on_conflict:
        - "error": raise ValueError if exists.
        - "skip" : print and return False.
        - "replace": remove existing (in-memory + JSON), then continue.
    """
    exists = job_in_store(job_id) or job_in_scheduler(job_id)
    if not exists:
        return True
    if on_conflict == "error":
        raise ValueError(f"Job id '{job_id}' already exists.")
    elif on_conflict == "skip":
        print(f"[INFO] Job '{job_id}' already exists. Skipping creation.")
        return False
    elif on_conflict == "replace":
        # Clear from scheduler
        schedule.clear(job_id)
        # Remove from persistence
        jobs = [j for j in load_jobs() if j.get("id") != job_id]
        save_jobs(jobs)
        return True
    else:
        raise ValueError(f"Unsupported on_conflict policy: {on_conflict}")

# -------------------------------
# Internal scheduling (no persistence)
# -------------------------------

def _schedule_once(job_id: str, func_name: str, run_at_timestamp: float, args=None, kwargs=None):
    args = args or []
    kwargs = kwargs or {}

    def job_wrapper():
        """Executes the job once, then removes it."""
        if func_name not in FUNCTION_MAP:
            print(f"[ERROR] Function '{func_name}' is not registered.")
            return
        FUNCTION_MAP[func_name](*args, **kwargs)
        # Remove from persistence
        jobs = load_jobs()
        jobs = [j for j in jobs if j["id"] != job_id]
        save_jobs(jobs)
        # Clear from in-memory schedule
        schedule.clear(job_id)

    delay_seconds = run_at_timestamp - time.time()
    if delay_seconds <= 0:
        print(f"[WARN] Job {job_id} scheduled in the past. Skipping.")
        return
    # Schedule via schedule library
    schedule.every(int(delay_seconds)).seconds.do(job_wrapper).tag(job_id)

def _schedule_recurring(job_id: str, func_name: str, interval: int, unit: str, args=None, kwargs=None):
    args = args or []
    kwargs = kwargs or {}

    def job_wrapper():
        if func_name not in FUNCTION_MAP:
            print(f"[ERROR] Function '{func_name}' is not registered.")
            return
        FUNCTION_MAP[func_name](*args, **kwargs)

    if unit == "seconds":
        schedule.every(interval).seconds.do(job_wrapper).tag(job_id)
    elif unit == "minutes":
        schedule.every(interval).minutes.do(job_wrapper).tag(job_id)
    elif unit == "hours":
        schedule.every(interval).hours.do(job_wrapper).tag(job_id)
    elif unit == "days":
        schedule.every(interval).days.do(job_wrapper).tag(job_id)
    else:
        raise ValueError(f"Unsupported unit: {unit}")

# -------------------------------
# Public APIs (with uniqueness + persistence)
# -------------------------------

def run_once_job(
    job_id: str,
    func_name: str,
    run_at_timestamp: float,
    args=None,
    kwargs=None,
    *,
    replace: bool = False,
    persist: bool = True,
):
    """
    Schedule a job to run once at a specific timestamp.
    replace: if True, replace existing job with same id; otherwise print and skip.
    persist: if False, do not write to JSON (used by reload_jobs()).
    """
    if persist:
        policy = "replace" if replace else "skip"
        if not ensure_unique(job_id, on_conflict=policy):
            return
    else:
        if job_in_scheduler(job_id):
            schedule.clear(job_id)
    _schedule_once(job_id, func_name, run_at_timestamp, args, kwargs)
    if persist:
        jobs = [j for j in load_jobs() if j["id"] != job_id]
        jobs.append({
            "id": job_id,
            "type": "once",
            "run_at": run_at_timestamp,
            "function": func_name,
            "args": args or [],
            "kwargs": kwargs or {}
        })
        save_jobs(jobs)

def recurring_job(
    job_id: str,
    func_name: str,
    interval: int,
    unit: str,
    args=None,
    kwargs=None,
    *,
    replace: bool = False,
    persist: bool = True,
):
    """
    Schedule a recurring job.
    replace: if True, replace existing job with same id; otherwise print and skip.
    persist: if False, do not write to JSON (used by reload_jobs()).
    """
    if persist:
        policy = "replace" if replace else "skip"
        if not ensure_unique(job_id, on_conflict=policy):
            return
    else:
        if job_in_scheduler(job_id):
            schedule.clear(job_id)
    _schedule_recurring(job_id, func_name, interval, unit, args, kwargs)
    if persist:
        jobs = [j for j in load_jobs() if j["id"] != job_id]
        jobs.append({
            "id": job_id,
            "type": "recurring",
            "interval": interval,
            "unit": unit,
            "function": func_name,
            "args": args or [],
            "kwargs": kwargs or {}
        })
        save_jobs(jobs)


def find_and_prioritize_jobs_by_pid(pid_substring: str, new_delay_seconds: float = 1.0):
    """
    Find all jobs whose ID contains the given PID substring and reschedule them to run sooner.
    """
    jobs = load_jobs()
    matched_jobs = [job for job in jobs if pid_substring in job.get("id", "")]

    if not matched_jobs:
        print(f"[INFO] No jobs found containing PID substring '{pid_substring}'.")
        return

    print(f"[INFO] Found {len(matched_jobs)} job(s) containing '{pid_substring}':")

    for job in matched_jobs:
        job_id = job["id"]
        print(f"  - Prioritizing job: {job_id}")

        # Clear existing job from scheduler
        schedule.clear(job_id)

        # Reschedule based on job type
        if job["type"] == "once":
            run_once_job(
                job_id,
                job["function"],
                time.time() + new_delay_seconds,
                job.get("args"),
                job.get("kwargs"),
                replace=True,
                persist=True
            )
        elif job["type"] == "recurring":
            recurring_job(
                job_id,
                job["function"],
                job["interval"],
                job["unit"],
                job.get("args"),
                job.get("kwargs"),
                replace=True,
                persist=True
            )
        else:
            print(f"[WARN] Unknown job type for job '{job_id}'")


# -------------------------------
# Reload Saved Jobs
# -------------------------------

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

# -------------------------------
# Scheduler Loop
# -------------------------------

def start_scheduler():
    """
    Start the scheduler loop (blocking).
    Call this once in main to begin.
    """
    try:
        while True:
            schedule.run_pending()
            time.sleep(0.5)
    except KeyboardInterrupt:
        print("[INFO] Scheduler stopped.")