"""Task dispatcher: composes TaskRunner / WorkflowRunner to execute assignments.

Capability: ``task-scheduler`` (closing the loop from ``assigned`` to ``executed``).

Composes the existing single-process execution entry points (``runtime.task.TaskRunner``,
``workflow.runner.WorkflowRunner``) without editing them, strictly through their
public ``run(task) -> Task`` / ``run(definition, device_id) -> WorkflowRun`` contracts.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING

from core.models import Task

if TYPE_CHECKING:
    from cloud.store import CloudStore


@dataclass(frozen=True)
class Assignment:
    """A scheduler-produced binding of a queued task to a specific device+host."""

    task_id: str
    device_id: str
    host_id: str
    goal: str | None
    workflow_definition_id: str | None


class RemoteDispatchNotSupportedError(RuntimeError):
    """Raised when an assignment targets a host other than the local process."""


class UnknownWorkflowDefinitionError(RuntimeError):
    """Raised when a workflow-based assignment references an unknown definition id."""


# Callable type aliases; kept lazy so importing this module never imports runtime/workflow.
TaskRunnerFactory = Callable[[], "object"]
WorkflowRunnerFactory = Callable[[], "object"]


class TaskDispatcher:
    """Executes assignments via the existing task / workflow runners.

    Local-only in this change: an assignment whose ``host_id`` does not match
    the dispatcher's own ``local_host_id`` raises ``RemoteDispatchNotSupportedError``
    rather than attempting any execution over the network.
    """

    def __init__(
        self,
        local_host_id: str,
        task_runner_factory: TaskRunnerFactory,
        workflow_runner_factory: WorkflowRunnerFactory,
        store: "CloudStore",
    ) -> None:
        self.local_host_id = local_host_id
        self.task_runner_factory = task_runner_factory
        self.workflow_runner_factory = workflow_runner_factory
        self.store = store

    def dispatch(self, assignment: Assignment) -> None:
        if assignment.host_id != self.local_host_id:
            raise RemoteDispatchNotSupportedError(
                f"assignment {assignment.task_id} targets host "
                f"{assignment.host_id!r}, but this dispatcher owns {self.local_host_id!r}"
            )

        if assignment.workflow_definition_id:
            status = self._dispatch_workflow(assignment)
        else:
            status = self._dispatch_goal(assignment)

        self.store.update_task(assignment.task_id, status=status)

    def _dispatch_goal(self, assignment: Assignment) -> str:
        task = Task(
            goal=assignment.goal or "",
            device_id=assignment.device_id,
        )
        result = self.task_runner_factory().run(task)  # type: ignore[attr-defined]
        result_status = getattr(result, "status", None)
        return "done" if result_status == "completed" else "failed"

    def _dispatch_workflow(self, assignment: Assignment) -> str:
        runner = self.workflow_runner_factory()
        store = getattr(runner, "store", None)
        if store is None:
            raise UnknownWorkflowDefinitionError(
                "workflow runner has no store to load definitions from"
            )
        definition = store.get_definition(assignment.workflow_definition_id)
        if definition is None:
            raise UnknownWorkflowDefinitionError(
                f"unknown workflow definition {assignment.workflow_definition_id!r}"
            )
        run = runner.run(definition, device_id=assignment.device_id)  # type: ignore[attr-defined]
        run_status = getattr(run, "status", None)
        return "done" if run_status == "completed" else "failed"