"""Read-only evaluation helpers for application job execution statistics."""

from typing import TYPE_CHECKING

from ..components.job import BackgroundJob, BaseJob, CronJob, StreamJob
from .counter import global_counter_get

if TYPE_CHECKING:
    from ..components.application_context import ApplicationContext


_JOB_ENTRY_CLASSES = {CronJob, StreamJob, BackgroundJob, BaseJob}


def check_job_count(job_name: str, app_context: "ApplicationContext") -> int:
    """Return the application-lifetime execution count for a registered job.

    ``app_context`` scopes the lookup because ReMe does not maintain a global
    current Application instance. Unknown job names use the same ``KeyError``
    contract as :meth:`Application.run_job`.
    """
    job = app_context.jobs.get(job_name)
    if job is None:
        raise KeyError(f"Job '{job_name}' not found")

    entry_class = next((cls for cls in type(job).mro() if cls in _JOB_ENTRY_CLASSES), None)
    if entry_class is None:
        raise TypeError(f"Job '{job_name}' does not inherit from a supported job implementation")
    # pylint: disable-next=protected-access
    return global_counter_get(app_context.metadata, job._counter_key(entry_class))


class JobCountTracker:
    """Measure registered job calls made while this context is active."""

    def __init__(self, job_names: list[str], app_context: "ApplicationContext") -> None:
        self.job_names = list(dict.fromkeys(job_names))
        self.app_context = app_context
        self._start_counts: dict[str, int] = {}
        self.counts: dict[str, int] = {}

    def __enter__(self) -> dict[str, int]:
        self._start_counts = {name: check_job_count(name, self.app_context) for name in self.job_names}
        return self.counts

    def __exit__(self, exc_type, exc_value, traceback) -> bool:
        self.counts.update(
            {
                name: check_job_count(name, self.app_context) - start_count
                for name, start_count in self._start_counts.items()
            },
        )
        return False


def track_job_counts(job_names: list[str], app_context: "ApplicationContext") -> JobCountTracker:
    """Return a context manager that reports call deltas for ``job_names``.

    Example:

    .. code-block:: python

        with track_job_counts(["search"], app.context) as counts:
            await app.run_job("agentic_answer", query="...")
        assert counts == {"search": 2}
    """
    return JobCountTracker(job_names, app_context)