# platform/deskclaw/execution/dispatcher.py
# ═══════════════════════════════════════════════════════════════
# Phase 3: Execution Engine — Test Orchestration & Dispatch
# Capability matching, topological DAG sort, async worker pool,
# timeout guard with retry, and raw_results collection.
# ═══════════════════════════════════════════════════════════════

from __future__ import annotations

import asyncio
import time
from collections.abc import Mapping
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Protocol, runtime_checkable

from deskclaw.adapters.protocol import AgentAdapter, AgentCapability, AgentResponse

# Lazy import to avoid circular dependency at module level.
# BenchmarkRegistry is imported inside methods that need it.

# ── Status Enum ──


class TestStatus(str, Enum):
    """Outcome status for an individual test execution."""

    PASSED = "PASSED"
    FAILED = "FAILED"
    ERROR = "ERROR"
    TIMEOUT = "TIMEOUT"
    SKIPPED = "SKIPPED"


# ── Data Models ──


@dataclass
class TestSpec:
    """
    Specification for a single test item within a benchmark suite.

    Attributes:
        id:                      Unique test identifier (e.g. ``ST-01``).
        suite:                   Parent benchmark suite name.
        required_capabilities:   Capabilities the agent must have to run this test.
        depends_on:              Test IDs that must complete before this one.
        severity:                ``"critical"`` | ``"high"`` | ``"normal"``.
    """

    id: str
    suite: str
    required_capabilities: set[AgentCapability] = field(default_factory=set)
    depends_on: list[str] = field(default_factory=list)
    severity: str = "normal"

    def __hash__(self) -> int:
        return hash(self.id)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, TestSpec):
            return NotImplemented
        return self.id == other.id


@runtime_checkable
class TestRunner(Protocol):
    """
    Protocol for callable test runners.

    Implementations receive an ``AgentAdapter`` and return an ``AgentResponse``.
    The dispatcher does NOT need to know test internals — only that they are
    async callables matching this signature.
    """

    async def __call__(self, agent: AgentAdapter) -> AgentResponse: ...


@dataclass
class TestResult:
    """Result of a single test execution."""

    test_id: str
    suite: str
    status: TestStatus
    response: AgentResponse | None = None
    error: str | None = None
    attempts: int = 1
    elapsed_ms: int = 0
    severity: str = "normal"

    def to_dict(self) -> dict[str, Any]:
        return {
            "test_id": self.test_id,
            "suite": self.suite,
            "status": self.status.value,
            "error": self.error,
            "attempts": self.attempts,
            "elapsed_ms": self.elapsed_ms,
            "severity": self.severity,
        }


@dataclass
class DispatchResult:
    """Aggregate output of a full dispatch run."""

    matched: int
    skipped: list[dict[str, str]]  # [{"test_id": ..., "reason": ...}]
    results: dict[str, TestResult]  # test_id → TestResult

    def to_dict(self) -> dict[str, Any]:
        return {
            "matched": self.matched,
            "skipped": self.skipped,
            "results": {k: v.to_dict() for k, v in self.results.items()},
        }


# ── Exceptions ──


class DispatchError(Exception):
    """Raised when the dispatch orchestration encounters an unrecoverable error."""


# ── Capability Matcher ──


def match_tests_to_agent(
    tests: list[TestSpec],
    agent: AgentAdapter,
) -> tuple[list[TestSpec], list[dict[str, str]]]:
    """
    Filter *tests* by agent capability matching.

    Returns:
        Tuple of (matched tests, skipped records).
        A test is matched if ``agent.capabilities() ⊇ test.required_capabilities``.
        Tests with no required capabilities are always matched.
    """
    agent_caps = agent.capabilities()
    matched: list[TestSpec] = []
    skipped: list[dict[str, str]] = []

    for test in tests:
        if test.required_capabilities and not test.required_capabilities.issubset(
            agent_caps
        ):
            missing = test.required_capabilities - agent_caps
            skipped.append(
                {
                    "test_id": test.id,
                    "reason": (
                        f"Agent lacks capabilities: "
                        f"{', '.join(c.value for c in sorted(missing, key=lambda c: c.value))}"
                    ),
                }
            )
        else:
            matched.append(test)

    return matched, skipped


# ── Topological Sort (DAG Layer Decomposition) ──


def topological_sort(tests: list[TestSpec]) -> list[list[TestSpec]]:
    """
    Decompose *tests* into execution layers via Kahn's algorithm.

    Tests within the same layer have no mutual dependencies and may run in
    parallel.  Layers are returned in dependency order (layer 0 first).

    Raises:
        DispatchError: If a dependency cycle is detected.
    """
    id_map: dict[str, TestSpec] = {t.id: t for t in tests}
    matched_ids = set(id_map.keys())

    # Build in-degree map (only count deps that are in our matched set)
    in_degree: dict[str, int] = {t.id: 0 for t in tests}
    dependents: dict[str, list[str]] = {t.id: [] for t in tests}

    for t in tests:
        for dep_id in t.depends_on:
            if dep_id in matched_ids:
                in_degree[t.id] += 1
                dependents[dep_id].append(t.id)

    layers: list[list[TestSpec]] = []
    remaining = set(id_map.keys())

    while remaining:
        # Current layer = all nodes with in_degree 0
        layer_ids = [tid for tid in remaining if in_degree[tid] == 0]
        if not layer_ids:
            raise DispatchError(f"Dependency cycle detected among: {sorted(remaining)}")

        layer = [id_map[tid] for tid in sorted(layer_ids)]
        layers.append(layer)

        for tid in layer_ids:
            remaining.remove(tid)
            for dependent_id in dependents[tid]:
                in_degree[dependent_id] -= 1

    return layers


# ── Core Dispatcher ──


class Dispatcher:
    """
    Phase 3 execution engine.

    Orchestrates test execution with:
    - Capability-based filtering
    - DAG-aware topological layer execution (layers run sequentially,
      tests within a layer run concurrently)
    - Semaphore-bounded worker pool
    - Per-test timeout guard with configurable retry
    """

    def __init__(
        self,
        agent: AgentAdapter,
        *,
        max_workers: int = 4,
        timeout_per_test_s: float = 600,
        retry_on_timeout: int = 1,
    ) -> None:
        self.agent = agent
        self.max_workers = max_workers
        self.timeout_per_test_s = timeout_per_test_s
        self.retry_on_timeout = retry_on_timeout

    @classmethod
    def from_registry(
        cls,
        agent: AgentAdapter,
        registry: object,
        *,
        max_workers: int = 4,
        timeout_per_test_s: float = 600,
        retry_on_timeout: int = 1,
        suite: str | None = None,
        test_ids: list[str] | None = None,
    ) -> tuple["Dispatcher", list[TestSpec], dict[str, TestRunner]]:
        """
        Build a Dispatcher wired to a :class:`BenchmarkRegistry`.

        Returns the dispatcher together with the test plan so callers
        can invoke ``dispatcher.dispatch(specs, runners)`` immediately.
        """
        from deskclaw.execution.benchmark_registry import BenchmarkRegistry

        if not isinstance(registry, BenchmarkRegistry):
            raise TypeError(
                f"Expected BenchmarkRegistry, got {type(registry).__name__}"
            )

        specs, runners = registry.build_test_plan(suite=suite, test_ids=test_ids)
        dispatcher = cls(
            agent,
            max_workers=max_workers,
            timeout_per_test_s=timeout_per_test_s,
            retry_on_timeout=retry_on_timeout,
        )
        return dispatcher, specs, runners

    async def dispatch(
        self,
        tests: list[TestSpec],
        runners: Mapping[str, TestRunner],
    ) -> DispatchResult:
        """
        Run all *tests* through capability matching, topological ordering,
        and parallel execution.

        Args:
            tests:   Full list of :class:`TestSpec` items to consider.
            runners: Mapping of ``test_id → async callable(agent) → AgentResponse``.

        Returns:
            :class:`DispatchResult` with matched/skipped/results.
        """
        # Step 1: Capability match
        matched, skipped = match_tests_to_agent(tests, self.agent)

        # Step 2: Topological sort
        layers = topological_sort(matched)

        # Step 3: Layer-by-layer parallel execution
        results: dict[str, TestResult] = {}
        semaphore = asyncio.Semaphore(self.max_workers)

        for layer in layers:
            tasks = [self._run_with_guard(spec, runners, semaphore) for spec in layer]
            layer_results = await asyncio.gather(*tasks, return_exceptions=True)

            for spec, result in zip(layer, layer_results):
                if isinstance(result, BaseException):
                    results[spec.id] = TestResult(
                        test_id=spec.id,
                        suite=spec.suite,
                        status=TestStatus.ERROR,
                        error=f"{type(result).__name__}: {result}",
                        severity=spec.severity,
                    )
                else:
                    results[spec.id] = result

        return DispatchResult(
            matched=len(matched),
            skipped=skipped,
            results=results,
        )

    async def _run_with_guard(
        self,
        spec: TestSpec,
        runners: Mapping[str, TestRunner],
        semaphore: asyncio.Semaphore,
    ) -> TestResult:
        """
        Execute a single test with semaphore, timeout, and retry.

        If the test_id has no corresponding runner, it is marked ERROR.
        """
        runner = runners.get(spec.id)
        if runner is None:
            return TestResult(
                test_id=spec.id,
                suite=spec.suite,
                status=TestStatus.ERROR,
                error=f"No runner registered for test '{spec.id}'",
                severity=spec.severity,
            )

        async with semaphore:
            last_error: str | None = None
            total_attempts = 1 + self.retry_on_timeout

            for attempt in range(total_attempts):
                start = time.monotonic()
                try:
                    response = await asyncio.wait_for(
                        runner(self.agent),
                        timeout=self.timeout_per_test_s,
                    )
                    elapsed_ms = int((time.monotonic() - start) * 1000)

                    return TestResult(
                        test_id=spec.id,
                        suite=spec.suite,
                        status=TestStatus.PASSED,
                        response=response,
                        attempts=attempt + 1,
                        elapsed_ms=elapsed_ms,
                        severity=spec.severity,
                    )

                except asyncio.TimeoutError:
                    elapsed_ms = int((time.monotonic() - start) * 1000)
                    last_error = (
                        f"Timeout after {self.timeout_per_test_s}s "
                        f"(attempt {attempt + 1}/{total_attempts})"
                    )
                    if attempt < total_attempts - 1:
                        continue

                    return TestResult(
                        test_id=spec.id,
                        suite=spec.suite,
                        status=TestStatus.TIMEOUT,
                        error=last_error,
                        attempts=total_attempts,
                        elapsed_ms=elapsed_ms,
                        severity=spec.severity,
                    )

                except Exception as exc:
                    elapsed_ms = int((time.monotonic() - start) * 1000)
                    return TestResult(
                        test_id=spec.id,
                        suite=spec.suite,
                        status=TestStatus.ERROR,
                        error=f"{type(exc).__name__}: {exc}",
                        attempts=attempt + 1,
                        elapsed_ms=elapsed_ms,
                        severity=spec.severity,
                    )

        # Should never be reached, but satisfy type checker
        return TestResult(  # pragma: no cover
            test_id=spec.id,
            suite=spec.suite,
            status=TestStatus.ERROR,
            error="Unexpected dispatch state",
            severity=spec.severity,
        )
