"""Bench registry: @bench decorator + name/index resolution.

Each bench module under ``kernbench.benches`` MUST register its callable
via ``@bench(name=..., description=...)``. Indices are assigned
alphabetically by name after eager import; they are a CLI convenience,
not a stable API.
"""
from __future__ import annotations

import re
from collections.abc import Callable
from dataclasses import dataclass
from importlib import import_module
from pkgutil import iter_modules

BenchFn = Callable[..., object]

_NAME_RE = re.compile(r"^[a-z][a-z0-9]*(-[a-z0-9]+)*$")


@dataclass(frozen=True)
class BenchSpec:
    index: int
    name: str
    description: str
    run: BenchFn


_PENDING: list[tuple[str, str, BenchFn]] = []
_REGISTERED_MODULES: set[str] = set()
_REGISTRY: dict[str, BenchSpec] = {}


def bench(*, name: str, description: str) -> Callable[[BenchFn], BenchFn]:
    if not isinstance(name, str) or not _NAME_RE.match(name):
        raise ValueError(
            f"bench name {name!r} must be kebab-case (lowercase, digits, dashes; "
            f"starts with a letter)."
        )
    if not isinstance(description, str) or not description.strip():
        raise ValueError(f"bench {name!r}: description must be a non-empty string.")

    def deco(fn: BenchFn) -> BenchFn:
        _PENDING.append((name, description, fn))
        _REGISTERED_MODULES.add(fn.__module__)
        return fn

    return deco


def _finalize() -> None:
    if _REGISTRY:
        return
    seen: set[str] = set()
    for n, _, _ in _PENDING:
        if n in seen:
            raise RuntimeError(f"duplicate bench name: {n!r}")
        seen.add(n)
    for i, (n, d, f) in enumerate(sorted(_PENDING, key=lambda t: t[0]), start=1):
        _REGISTRY[n] = BenchSpec(index=i, name=n, description=d, run=f)


def list_all() -> list[BenchSpec]:
    _finalize()
    return sorted(_REGISTRY.values(), key=lambda s: s.index)


def resolve(identifier: str) -> BenchSpec:
    _finalize()
    if not isinstance(identifier, str) or not identifier.strip():
        raise ValueError("bench identifier must be a non-empty string.")
    ident = identifier.strip()
    if ident.isdigit():
        idx = int(ident)
        for s in _REGISTRY.values():
            if s.index == idx:
                return s
        raise ValueError(
            f"No bench with index {idx}. Use 'kernbench list' to see options."
        )
    if ident in _REGISTRY:
        return _REGISTRY[ident]
    raise ValueError(
        f"Unknown bench {ident!r}. Use 'kernbench list' to see options."
    )


def _audit_modules(imported: list[str], registered: set[str]) -> None:
    missing = sorted(m for m in imported if m not in registered)
    if missing:
        raise RuntimeError(
            f"Bench module(s) missing @bench decorator: {missing}. "
            f"Each file under kernbench.benches/ must register at least one bench "
            f"via @bench(...), or be renamed with a leading underscore if it is a "
            f"helper."
        )


def _eager_import_and_audit(pkg_path: list[str], pkg_name: str) -> None:
    imported: list[str] = []
    for m in iter_modules(pkg_path):
        if m.name == "registry" or m.name.startswith("_"):
            continue
        mod = import_module(f"{pkg_name}.{m.name}")
        imported.append(mod.__name__)
    _audit_modules(imported, _REGISTERED_MODULES)