"""Executes functions in background threads with automatic tqdm progress capture."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/job_runner.ipynb.

# %% auto 0
__all__ = ['JobRunner']

# %% ../nbs/job_runner.ipynb 3
import threading
from typing import Callable, Any, Optional, Dict
from .progress_info import ProgressInfo
from .progress_monitor import ProgressMonitor
from .patch_tqdm import patch_tqdm

# %% ../nbs/job_runner.ipynb 5
class JobRunner:
    """
    Runs a callable in a background thread, patches tqdm inside the job,
    and forwards ProgressInfo updates to a ProgressMonitor under job_id.
    """
    def __init__(
        self,
        monitor: ProgressMonitor  # Progress monitor instance to receive updates
    ):
        "Initialize a job runner with a progress monitor"
        self.monitor = monitor
        self._threads: Dict[str, threading.Thread] = {}

    def start(
        self,
        job_id: str,  # Unique identifier for this job
        fn: Callable[..., Any],
        *args,
        patch_kwargs: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> threading.Thread:  # The thread running the job
        "Start a job in a background thread with automatic tqdm patching"
        patch_kwargs = patch_kwargs or {}
        def _target():
            "Internal thread target that wraps the job function with tqdm patching"
            def cb(
                info: ProgressInfo  # Progress update from the patched tqdm
            ):
                "Callback to forward progress updates to the monitor"
                self.monitor.update(job_id, info)
            try:
                with patch_tqdm(cb, **patch_kwargs):
                    fn(*args, **kwargs)
            except Exception as e:
                # Record an error as a final update
                self.monitor.update(job_id, ProgressInfo(progress=100.0, description=f"ERROR: {e}"))
        t = threading.Thread(target=_target, name=f"job-{job_id}", daemon=True)
        t.start()
        self._threads[job_id] = t
        return t

    def is_alive(
        self,
        job_id: str  # Unique identifier of the job to check
    ) -> bool:  # True if the job thread is still running
        "Check if a job's thread is still running"
        t = self._threads.get(job_id)
        return bool(t and t.is_alive())

    def join(
        self,
        job_id: str,  # Unique identifier of the job to wait for
        timeout: Optional[float] = None  # Maximum seconds to wait (None for indefinite)
    ) -> None:  # Returns when thread completes or timeout expires
        "Wait for a job's thread to complete"
        t = self._threads.get(job_id)
        if t:
            t.join(timeout)
