"""Benchmark registry for discovery and instantiation.

This module provides a central registry for all available benchmarks,
supporting both programmatic registration and YAML-based configuration.
"""

from pathlib import Path
from typing import Any

import yaml

from embergraphrouter.evals.benchmarks.base import Benchmark, BenchmarkDomain


class BenchmarkRegistry:
    """Central registry for benchmark discovery and instantiation.

    Benchmarks can be registered either:
    1. Programmatically via register()
    2. Via YAML configuration in benchmarks/configs/

    Example:
        registry = BenchmarkRegistry()
        registry.load_from_config()

        # Get a specific benchmark
        gpqa = registry.get("gpqa_diamond")

        # List all fresh (OOD) benchmarks
        fresh = registry.list_fresh_benchmarks()
    """

    _instance: "BenchmarkRegistry | None" = None
    _benchmarks: dict[str, type[Benchmark]]
    _configs: dict[str, dict[str, Any]]

    def __new__(cls) -> "BenchmarkRegistry":
        """Singleton pattern for global registry access."""
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._benchmarks = {}
            cls._instance._configs = {}
        return cls._instance

    def register(
        self,
        name: str,
        benchmark_class: type[Benchmark],
        config: dict[str, Any] | None = None,
    ) -> None:
        """Register a benchmark class.

        Args:
            name: Unique identifier for this benchmark.
            benchmark_class: The Benchmark subclass to register.
            config: Optional configuration overrides.
        """
        self._benchmarks[name] = benchmark_class
        if config:
            self._configs[name] = config

    def get(self, name: str, **kwargs: Any) -> Benchmark:
        """Instantiate a benchmark by name.

        Args:
            name: Registered benchmark name.
            **kwargs: Constructor arguments passed to the benchmark.

        Returns:
            Instantiated Benchmark.

        Raises:
            KeyError: If benchmark not found in registry.
        """
        if name not in self._benchmarks:
            available = ", ".join(sorted(self._benchmarks.keys()))
            raise KeyError(
                f"Benchmark '{name}' not found. Available: {available}"
            )

        # Merge config with kwargs (kwargs take precedence)
        config = self._configs.get(name, {}).copy()
        config.update(kwargs)

        return self._benchmarks[name](**config)

    def list_all(self) -> list[str]:
        """List all registered benchmark names."""
        return sorted(self._benchmarks.keys())

    def list_by_domain(self, domain: BenchmarkDomain) -> list[str]:
        """List benchmarks in a specific domain."""
        results = []
        for name, cls in self._benchmarks.items():
            if hasattr(cls, "domain") and cls.domain == domain:
                results.append(name)
        return sorted(results)

    def list_fresh_benchmarks(self) -> list[str]:
        """List benchmarks not in EB1 training data (OOD).

        These are suitable for uncontaminated evaluation.
        """
        results = []
        for name in self._benchmarks:
            # Check against the training benchmarks set
            if name.lower() not in Benchmark.TRAINING_BENCHMARKS:
                results.append(name)
        return sorted(results)

    def list_training_benchmarks(self) -> list[str]:
        """List benchmarks that overlap with EB1 training data."""
        results = []
        for name in self._benchmarks:
            if name.lower() in Benchmark.TRAINING_BENCHMARKS:
                results.append(name)
        return sorted(results)

    def load_from_config(self, config_path: Path | None = None) -> None:
        """Load benchmark configurations from YAML.

        Args:
            config_path: Path to benchmarks.yaml. If None, uses default.
        """
        if config_path is None:
            config_path = Path(__file__).parent / "configs" / "benchmarks.yaml"

        if not config_path.exists():
            return

        with open(config_path) as f:
            config = yaml.safe_load(f) or {}

        for name, bench_config in config.get("benchmarks", {}).items():
            self._configs[name] = bench_config

    def get_config(self, name: str) -> dict[str, Any]:
        """Get configuration for a benchmark.

        Args:
            name: Benchmark name.

        Returns:
            Configuration dict (empty if no config).
        """
        return self._configs.get(name, {}).copy()

    def clear(self) -> None:
        """Clear all registrations. Primarily for testing."""
        self._benchmarks.clear()
        self._configs.clear()


# Global registry instance
_registry = BenchmarkRegistry()


def register_benchmark(
    name: str,
    config: dict[str, Any] | None = None,
):
    """Decorator to register a benchmark class.

    Example:
        @register_benchmark("gpqa_diamond")
        class GPQADiamondBenchmark(Benchmark):
            ...
    """
    def decorator(cls: type[Benchmark]) -> type[Benchmark]:
        _registry.register(name, cls, config)
        return cls
    return decorator


def get_benchmark(name: str, **kwargs: Any) -> Benchmark:
    """Get a benchmark instance by name."""
    return _registry.get(name, **kwargs)


def list_benchmarks(
    domain: BenchmarkDomain | None = None,
    fresh_only: bool = False,
) -> list[str]:
    """List available benchmarks.

    Args:
        domain: Filter by domain if specified.
        fresh_only: If True, only return OOD benchmarks.

    Returns:
        List of benchmark names.
    """
    if fresh_only:
        return _registry.list_fresh_benchmarks()
    if domain:
        return _registry.list_by_domain(domain)
    return _registry.list_all()


__all__ = [
    "BenchmarkRegistry",
    "register_benchmark",
    "get_benchmark",
    "list_benchmarks",
]
