"""Per-task runner dispatch.

Neutral home for the routing that maps a task to its execution vehicle, so no
single runner module owns the cross-cutting dispatcher. ``_execution_worker_loop`` plugs
``dispatch_execute`` in as its ``ExecuteFn``; a job whose tasks span runners
(e.g. subprocess and container tasks) is served by one worker without runner
affinity rules.
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable

from sqlmodel import select

from ..docker_config import effective_image_tag
from ..models import RUNNER_DOCKER, RUNNER_KUBERNETES, RUNNER_SUBPROCESS, Job, RunnerMode, Task
from ..orch_context import get_sql_session
from ..runner_config import ENTRY_SHELL, KubernetesRunner, RunnerConfigT, parse_runner_config
from .docker_worker import _run_task_in_container
from .execution_worker import JobDispatch
from .kubernetes_worker import _run_task_in_pod
from .mp_worker import _run_shell_on_host, _run_task_in_child

ExecuteResult = tuple[bool, dict | None, str | None, str | None]


def _kube_dict(runner: RunnerConfigT | None) -> dict | None:
    """Reconstruct the kubernetes_config dict the k8s worker expects from a
    KubernetesRunner; None for any other runner."""
    if not isinstance(runner, KubernetesRunner):
        return None
    return {
        "namespace": runner.namespace,
        "service_account": runner.service_account,
        "image_pull_secret": runner.image_pull_secret,
        "resources": runner.resources,
    }


async def _resolve_dispatch(task: Task) -> JobDispatch:
    """Pick the runner for a task and snapshot its job's launch spec.

    Every task inherits the job's ``runner_mode``."""
    async with get_sql_session() as session:
        job = (await session.execute(select(Job).where(Job.id == task.job_id))).scalar_one_or_none()
    if job is None:
        return JobDispatch(RUNNER_SUBPROCESS, None, None)
    runner = parse_runner_config(job.runner) if job.runner else None
    image_tag = effective_image_tag(runner) if runner is not None else None
    image_source = getattr(runner, "image", None)
    return JobDispatch(
        job.runner_mode,
        image_tag,
        _kube_dict(runner),
        task.entry_type,
        task.command,
        task.command_env,
        image_source,
    )


# Image-based runners need the dispatch snapshot; subprocess is the default and
# needs nothing, so it stays off the registry rather than carry an unused arg.
_IMAGE_RUNNERS: dict[RunnerMode, Callable[[Task, int, JobDispatch], Awaitable[ExecuteResult]]] = {
    RUNNER_DOCKER: _run_task_in_container,
    RUNNER_KUBERNETES: _run_task_in_pod,
}


async def dispatch_execute(task: Task, execution_worker_id: int) -> ExecuteResult:
    """ExecuteFn that picks the runner per task."""
    dispatch = await _resolve_dispatch(task)
    handler = _IMAGE_RUNNERS.get(dispatch.runner_mode)
    if handler is not None:
        return await handler(task, execution_worker_id, dispatch)
    if dispatch.entry_type == ENTRY_SHELL:
        return await _run_shell_on_host(task, execution_worker_id, dispatch)
    return await _run_task_in_child(task, execution_worker_id)
