import importlib
import logging
import os
import sys
from datetime import datetime
from pathlib import Path

from screensuite.basebenchmark import BaseBenchmark
from screensuite.utils import check_dependency

logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
datetime_str: str = datetime.now().strftime("%Y%m%d@%H%M%S")
stdout_handler = logging.StreamHandler(sys.stdout)
stdout_handler.setLevel(logging.INFO)
formatter = logging.Formatter(fmt="\x1b[1;33m[%(asctime)s \x1b[31m%(levelname)s\x1b[1;33m] \x1b[0m%(message)s")
stdout_handler.setFormatter(formatter)
logger.addHandler(stdout_handler)


class BenchmarkRegistry:
    """
    A registry for storing and managing benchmarks.

    This registry allows for easy registration and retrieval of benchmarks.
    """

    def __init__(self):
        self._benchmarks: dict[str, BaseBenchmark] = {}

    def register(self, benchmarks: list[BaseBenchmark] | BaseBenchmark) -> None:
        """
        Register a benchmark in the registry.

        Args:
            benchmark: The benchmark to register

        Raises:
            ValueError: If a benchmark with the same name already exists
        """
        benchmarks_to_register = benchmarks if isinstance(benchmarks, list) else [benchmarks]
        for benchmark in benchmarks_to_register:
            if benchmark.name in self._benchmarks:
                raise ValueError(f"Benchmark with name '{benchmark.name}' already registered")
            self._benchmarks[benchmark.name] = benchmark

    def get(self, name: str) -> BaseBenchmark | None:
        """
        Get a benchmark by name.

        Args:
            name: The name of the benchmark to retrieve

        Returns:
            The benchmark if found, None otherwise
        """
        return self._benchmarks.get(name)

    def get_by_tags(self, tags: list[str], match_all: bool = False) -> list[BaseBenchmark]:
        """
        Get all benchmarks that have all the specified tags.

        Args:
            tags: The tags to filter by

        Returns:
            A list of benchmarks that have all the specified tags
        """
        if match_all:
            return list(
                set(benchmark for benchmark in self._benchmarks.values() if all(tag in benchmark.tags for tag in tags))
            )
        else:
            benchmarks = set(
                benchmark for benchmark in self._benchmarks.values() if any(tag in benchmark.tags for tag in tags)
            )
            return list(benchmarks)

    def list_all(self) -> list[BaseBenchmark]:
        """
        Get all registered benchmarks.

        Returns:
            A list of all registered benchmarks
        """
        return list(self._benchmarks.values())

    def list_names(self) -> list[str]:
        """
        Get the names of all registered benchmarks.

        Returns:
            A list of benchmark names
        """
        return list(self._benchmarks.keys())

    def remove(self, name: str) -> bool:
        """
        Remove a benchmark from the registry.

        Args:
            name: The name of the benchmark to remove

        Returns:
            True if the benchmark was removed, False if it didn't exist
        """
        if name in self._benchmarks:
            del self._benchmarks[name]
            return True
        return False

    def clear(self) -> None:
        """
        Remove all benchmarks from the registry.
        """
        self._benchmarks.clear()


def discover_benchmarks():
    """
    Automatically discover and import all benchmarks in the benchmarks directory.

    This function imports benchmark.py files in the benchmarks directory and its immediate subdirectories.
    """
    BENCHMARKS_DIR = Path(__file__).parent / "benchmarks"

    for item in os.listdir(BENCHMARKS_DIR):
        item_path = BENCHMARKS_DIR / item

        if not item_path.is_dir() or item.startswith("__") or item.startswith("."):
            continue

        # Check for benchmark.py in subdirectories (one level deeper)
        for subitem in os.listdir(item_path):
            subitem_path = item_path / subitem

            if not subitem_path.is_dir() or subitem.startswith("__") or subitem.startswith("."):
                continue

            if subitem in ["osworld", "android_world"]:
                if not check_dependency(subitem):
                    continue

            sub_benchmark_file = subitem_path / "benchmark.py"
            if sub_benchmark_file.exists():
                # Import the module
                module_name = f"screensuite.benchmarks.{item}.{subitem}.benchmark"
                try:
                    importlib.import_module(module_name)
                except (ImportError, AttributeError) as e:
                    print(f"Error importing benchmark from {module_name}: {e}")


# Create a singleton instance of the registry
registry = BenchmarkRegistry()
discover_benchmarks()
