"""
Background job manager for long-running LLM test runs.

Jobs are stored in a module-level dict so they survive Streamlit page
navigation (the Python process stays alive between rerenders).
"""

import threading
import datetime
import uuid
from typing import Any, Dict, Optional

from streamlit.runtime.scriptrunner import add_script_run_ctx, get_script_run_ctx

MAX_WORKERS = 3
_semaphore  = threading.Semaphore(MAX_WORKERS)

_lock = threading.Lock()
JOBS: Dict[str, Dict[str, Any]] = {}

STATUS_PENDING   = "pending"
STATUS_RUNNING   = "running"
STATUS_DONE      = "done"
STATUS_FAILED    = "failed"
STATUS_CANCELLED = "cancelled"



def submit_job(
    label: str,
    fn,
    args: tuple = (),
    kwargs: dict = None,
    group_id: Optional[str] = None,
) -> str:
    """Schedule *fn(*args, **kwargs)* to run in a background thread.

    Returns the new job_id.
    """
    if kwargs is None:
        kwargs = {}

    job_id = str(uuid.uuid4())[:8]
    record = {
        "job_id":     job_id,
        "group_id":   group_id,
        "label":      label,
        "status":     STATUS_PENDING,
        "created_at": datetime.datetime.now().isoformat(timespec="seconds"),
        "finished_at": None,
        "result":     None,
        "error":      None,
        "_cancel":    threading.Event(),
    }

    with _lock:
        JOBS[job_id] = record

    ctx = get_script_run_ctx()
    t = threading.Thread(target=_worker, args=(job_id, fn, args, kwargs), daemon=True)
    if ctx:
        add_script_run_ctx(t, ctx)
    t.start()
    return job_id


def cancel_job(job_id: str):
    """Request cancellation of a job.  The job's thread checks this flag."""
    with _lock:
        job = JOBS.get(job_id)
    if job and job["status"] in (STATUS_PENDING, STATUS_RUNNING):
        job["_cancel"].set()
        job["status"] = STATUS_CANCELLED
        job["finished_at"] = datetime.datetime.now().isoformat(timespec="seconds")


def cancel_group(group_id: str):
    """Cancel all jobs that belong to *group_id*."""
    with _lock:
        ids = [jid for jid, j in JOBS.items() if j.get("group_id") == group_id]
    for jid in ids:
        cancel_job(jid)


def get_job(job_id: str) -> Optional[Dict]:
    with _lock:
        return dict(JOBS.get(job_id, {}))


def get_all_jobs() -> list:
    with _lock:
        return [dict(j) for j in JOBS.values()]


def get_active_count() -> int:
    with _lock:
        return sum(1 for j in JOBS.values() if j["status"] in (STATUS_PENDING, STATUS_RUNNING))


def remove_job(job_id: str):
    with _lock:
        JOBS.pop(job_id, None)


def clear_finished():
    with _lock:
        finished = [jid for jid, j in JOBS.items()
                    if j["status"] in (STATUS_DONE, STATUS_FAILED, STATUS_CANCELLED)]
        for jid in finished:
            del JOBS[jid]



def _worker(job_id: str, fn, args: tuple, kwargs: dict):
    with _lock:
        job = JOBS.get(job_id)

    if job is None or job["_cancel"].is_set():
        return

    _semaphore.acquire()
    try:
        with _lock:
            job = JOBS.get(job_id)
        if job is None or job["_cancel"].is_set():
            return

        job["status"] = STATUS_RUNNING

        try:
            result = fn(*args, cancel_event=job["_cancel"], **kwargs)
            with _lock:
                if not job["_cancel"].is_set():
                    job["status"]      = STATUS_DONE
                    job["result"]      = result
                    job["finished_at"] = datetime.datetime.now().isoformat(timespec="seconds")
        except Exception as exc:
            with _lock:
                job["status"]      = STATUS_FAILED
                job["error"]       = str(exc)
                job["finished_at"] = datetime.datetime.now().isoformat(timespec="seconds")
    finally:
        _semaphore.release()

