# platform/deskclaw/execution/benchmark_registry.py
# ═══════════════════════════════════════════════════════════════
# Benchmark Registry — Discovery, Registration & Test Plan Builder
#
# Three-level discovery:
#   1. deskclaw-plugin.yaml execution.benchmarks (explicit)
#   2. plugins/benchmarks/*.yaml manifest files (third-party)
#   3. Subdirectories with deskclaw-benchmark.yaml (auto-discovery)
#
# Bridges the gap between benchmark definitions and Dispatcher's
# TestSpec + TestRunner interface.
# ═══════════════════════════════════════════════════════════════

from __future__ import annotations

import importlib
import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Literal

from deskclaw.adapters.protocol import AgentAdapter, AgentCapability, AgentResponse
from deskclaw.execution.dispatcher import TestRunner, TestSpec

logger = logging.getLogger("deskclaw.execution.benchmark_registry")


_CAPABILITY_MAP: dict[str, AgentCapability] = {
    "text_generation": AgentCapability.TEXT_GENERATION,
    "code_execution": AgentCapability.CODE_EXECUTION,
    "tool_use": AgentCapability.TOOL_USE,
    "multimodal": AgentCapability.MULTIMODAL,
    "desktop_control": AgentCapability.DESKTOP_CONTROL,
    "multi_turn": AgentCapability.MULTI_TURN,
}


@dataclass
class TestDefinition:
    """Single test within a benchmark manifest."""

    id: str
    severity: str = "normal"
    runner: str = ""
    depends_on: list[str] = field(default_factory=list)


@dataclass
class BenchmarkManifest:
    """Self-describing metadata for a benchmark suite."""

    name: str
    version: str = "0.0.0"
    type: Literal["builtin", "external", "plugin"] = "plugin"
    description: str = ""
    required_capabilities: list[str] = field(default_factory=list)
    tests: list[TestDefinition] = field(default_factory=list)
    path: str = ""
    entry_point: str = ""
    config_schema: dict[str, Any] | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> BenchmarkManifest:
        bench = data.get("benchmark", data)

        tests = [
            TestDefinition(
                id=t["id"],
                severity=t.get("severity", "normal"),
                runner=t.get("runner", ""),
                depends_on=t.get("depends_on", []),
            )
            for t in bench.get("tests", [])
        ]

        return cls(
            name=bench.get("name", "unknown"),
            version=bench.get("version", "0.0.0"),
            type=bench.get("type", "plugin"),
            description=bench.get("description", ""),
            required_capabilities=bench.get("required_capabilities", []),
            tests=tests,
            path=bench.get("path", ""),
            entry_point=bench.get("entry_point", ""),
            config_schema=bench.get("config_schema"),
        )

    @classmethod
    def from_plugin_yaml_entry(cls, entry: dict[str, Any]) -> BenchmarkManifest:
        """Build from a deskclaw-plugin.yaml execution.benchmarks entry."""
        test_ids = entry.get("tests", [])
        tests = [TestDefinition(id=tid, severity="normal") for tid in test_ids]
        return cls(
            name=entry.get("name", "unknown"),
            type=entry.get("type", "external"),
            required_capabilities=entry.get("required_capabilities", []),
            tests=tests,
            path=entry.get("path", ""),
        )


def _parse_yaml_or_json(path: Path) -> dict[str, Any]:
    text = path.read_text(encoding="utf-8")
    try:
        import yaml  # type: ignore[import-untyped]

        data = yaml.safe_load(text)
    except ModuleNotFoundError:
        data = json.loads(text)

    if not isinstance(data, dict):
        raise ValueError(f"Expected dict at top level of {path}")
    return data


def _resolve_runner(runner_path: str) -> Callable[[AgentAdapter], Any] | None:
    """
    Resolve a dotted Python path to a callable.

    Format: "module.path:callable_name"
    Example: "custom_stress_tests.runners.run_st01:run"
    """
    if not runner_path or ":" not in runner_path:
        return None

    module_path, _, func_name = runner_path.rpartition(":")
    try:
        module = importlib.import_module(module_path)
        return getattr(module, func_name)
    except (ImportError, AttributeError) as exc:
        logger.warning("Failed to resolve runner '%s': %s", runner_path, exc)
        return None


def _parse_caps(caps: list[str]) -> set[AgentCapability]:
    result: set[AgentCapability] = set()
    for c in caps:
        cap = _CAPABILITY_MAP.get(c.lower().strip())
        if cap:
            result.add(cap)
    return result


class BenchmarkRegistry:
    """
    Central registry for benchmark suites.

    Manages discovery, registration, and conversion of benchmark manifests
    into the TestSpec + TestRunner pairs that the Dispatcher consumes.
    """

    def __init__(self) -> None:
        self._manifests: dict[str, BenchmarkManifest] = {}

    @property
    def registered(self) -> list[str]:
        return sorted(self._manifests.keys())

    def register(self, manifest: BenchmarkManifest, *, upgrade: bool = False) -> bool:
        """
        Register a benchmark manifest.

        If *upgrade* is True and the manifest name already exists, the new
        manifest replaces the old one when it carries richer test definitions
        (i.e. tests with non-empty ``runner`` fields whereas the existing
        manifest has none).  This allows filesystem-level manifests discovered
        in Level 3 to supersede stub entries from Level 1 plugin YAML.

        Returns True if newly registered or upgraded, False if skipped.
        """
        existing = self._manifests.get(manifest.name)
        if existing is not None:
            if not upgrade:
                logger.debug("Benchmark '%s' already registered.", manifest.name)
                return False
            existing_has_runners = any(t.runner for t in existing.tests)
            new_has_runners = any(t.runner for t in manifest.tests)
            if new_has_runners and not existing_has_runners:
                logger.info(
                    "Upgrading benchmark '%s': filesystem manifest overrides plugin stub.",
                    manifest.name,
                )
                self._manifests[manifest.name] = manifest
                return True
            logger.debug(
                "Benchmark '%s' already registered with runners — skipping upgrade.",
                manifest.name,
            )
            return False
        self._manifests[manifest.name] = manifest
        logger.info("Registered benchmark: %s (type=%s)", manifest.name, manifest.type)
        return True

    def discover(
        self,
        plugin_benchmarks: list[dict[str, Any]] | None = None,
        project_root: Path | None = None,
    ) -> list[str]:
        """
        Three-level benchmark discovery.

        Args:
            plugin_benchmarks: Entries from deskclaw-plugin.yaml execution.benchmarks.
            project_root:      Project root for filesystem scanning.

        Returns:
            List of newly registered benchmark names.
        """
        newly_registered: list[str] = []

        # Level 1: Plugin YAML declarations
        for entry in plugin_benchmarks or []:
            manifest = BenchmarkManifest.from_plugin_yaml_entry(entry)
            if self.register(manifest):
                newly_registered.append(manifest.name)

        if project_root is None:
            return newly_registered

        # Level 2: plugins/benchmarks/*.yaml  (upgrade=True to override stubs)
        benchmarks_dir = project_root / "plugins" / "benchmarks"
        if benchmarks_dir.is_dir():
            for manifest_file in sorted(benchmarks_dir.glob("*.yaml")):
                try:
                    data = _parse_yaml_or_json(manifest_file)
                    manifest = BenchmarkManifest.from_dict(data)
                    if self.register(manifest, upgrade=True):
                        newly_registered.append(manifest.name)
                except Exception as exc:
                    logger.warning(
                        "Failed to load benchmark manifest %s: %s",
                        manifest_file,
                        exc,
                    )

        # Level 3: Subdirectories with deskclaw-benchmark.yaml  (upgrade=True)
        for subdir in sorted(project_root.iterdir()):
            if not subdir.is_dir():
                continue
            manifest_path = subdir / "deskclaw-benchmark.yaml"
            if not manifest_path.exists():
                continue
            try:
                data = _parse_yaml_or_json(manifest_path)
                manifest = BenchmarkManifest.from_dict(data)
                if not manifest.path:
                    manifest.path = str(subdir)
                if self.register(manifest, upgrade=True):
                    newly_registered.append(manifest.name)
            except Exception as exc:
                logger.warning(
                    "Failed to load benchmark from %s: %s", manifest_path, exc
                )

        return newly_registered

    def build_test_plan(
        self,
        suite: str | None = None,
        test_ids: list[str] | None = None,
    ) -> tuple[list[TestSpec], dict[str, TestRunner]]:
        """
        Convert registered benchmarks into Dispatcher-compatible TestSpec + Runner pairs.

        Args:
            suite: If given, only build test plan for this specific suite.
                   If None, build for all registered benchmarks.
            test_ids: If given, only include tests whose IDs are in this list.

        Returns:
            Tuple of (test_specs, runners_mapping).
        """
        specs: list[TestSpec] = []
        runners: dict[str, TestRunner] = {}
        selected_test_ids = {tid.strip() for tid in (test_ids or []) if tid.strip()}

        target_manifests = (
            {suite: self._manifests[suite]}
            if suite and suite in self._manifests
            else self._manifests
        )

        for name, manifest in target_manifests.items():
            required_caps = _parse_caps(manifest.required_capabilities)

            for test_def in manifest.tests:
                if selected_test_ids and test_def.id not in selected_test_ids:
                    continue
                spec = TestSpec(
                    id=test_def.id,
                    suite=name,
                    required_capabilities=required_caps,
                    depends_on=test_def.depends_on,
                    severity=test_def.severity,
                )
                specs.append(spec)

                runner_fn = _resolve_runner(test_def.runner)
                if runner_fn is not None:
                    runners[test_def.id] = runner_fn

        return specs, runners

    def get_manifest(self, name: str) -> BenchmarkManifest | None:
        return self._manifests.get(name)

    def list_manifests(self) -> list[BenchmarkManifest]:
        return list(self._manifests.values())
