"""Simple benchmark registry - discovers benchmarks from directory."""

import json
from dataclasses import dataclass
from pathlib import Path

from pentestgpt.benchmark.config import DEFAULT_BENCHMARKS_DIR


@dataclass
class BenchmarkInfo:
    """Information about a benchmark."""

    id: str
    name: str
    description: str
    level: int
    tags: list[str]
    path: Path
    expected_flag: str | None = None


class BenchmarkRegistry:
    """Discovers benchmarks from a directory."""

    def __init__(self, benchmarks_dir: Path | None = None):
        self.benchmarks_dir = benchmarks_dir or DEFAULT_BENCHMARKS_DIR
        self._benchmarks: dict[str, BenchmarkInfo] = {}

    def load(self) -> None:
        """Load all benchmarks from directory."""
        self._benchmarks.clear()

        if not self.benchmarks_dir.exists():
            raise FileNotFoundError(f"Benchmarks directory not found: {self.benchmarks_dir}")

        for benchmark_path in sorted(self.benchmarks_dir.iterdir()):
            if not benchmark_path.is_dir():
                continue

            benchmark_json = benchmark_path / "benchmark.json"
            if not benchmark_json.exists():
                continue

            try:
                with open(benchmark_json) as f:
                    data = json.load(f)

                info = BenchmarkInfo(
                    id=benchmark_path.name,
                    name=data.get("name", benchmark_path.name),
                    description=data.get("description", ""),
                    level=int(data.get("level", 1)),
                    tags=data.get("tags", []),
                    path=benchmark_path,
                )
                self._benchmarks[info.id] = info
            except (json.JSONDecodeError, KeyError, ValueError) as e:
                print(f"Warning: Failed to parse {benchmark_json}: {e}")

    def get(self, benchmark_id: str) -> BenchmarkInfo | None:
        """Get benchmark by ID."""
        if not self._benchmarks:
            self.load()
        return self._benchmarks.get(benchmark_id)

    def list_all(self) -> list[BenchmarkInfo]:
        """List all benchmarks."""
        if not self._benchmarks:
            self.load()
        return sorted(self._benchmarks.values(), key=lambda b: b.id)

    def filter(
        self,
        tags: list[str] | None = None,
        levels: list[int] | None = None,
    ) -> list[BenchmarkInfo]:
        """Filter benchmarks by tags or levels."""
        if not self._benchmarks:
            self.load()

        result = list(self._benchmarks.values())

        if levels:
            result = [b for b in result if b.level in levels]

        if tags:
            tags_lower = [t.lower() for t in tags]
            result = [b for b in result if any(t.lower() in tags_lower for t in b.tags)]

        return sorted(result, key=lambda b: b.id)

    def get_all_tags(self) -> set[str]:
        """Get all unique tags."""
        if not self._benchmarks:
            self.load()
        tags: set[str] = set()
        for b in self._benchmarks.values():
            tags.update(b.tags)
        return tags