from __future__ import annotations

import importlib
from collections.abc import Callable
from typing import Any

from kernbench.runtime_api.context import RuntimeContext

BenchFn = Callable[[RuntimeContext], Any]


def _load_module(bench_id: str):
    bench_id = bench_id.strip()
    if not bench_id:
        raise ValueError("Bench id is empty.")
    module_path = f"benches.{bench_id}"
    try:
        return importlib.import_module(module_path)
    except ModuleNotFoundError as e:
        raise ValueError(
            f"Unknown bench '{bench_id}'. Expected module {module_path}.py"
        ) from e


def resolve_bench(bench_id: str) -> BenchFn:
    """Resolve a bench id into its ``run(torch)`` callable.

    Expected layout (repo root):
        benches/<bench_id>.py
            def run(torch: RuntimeContext) -> Any
    """
    mod = _load_module(bench_id)
    run_fn = getattr(mod, "run", None)
    if run_fn is None:
        raise ValueError(
            f"Bench module benches.{bench_id} must define 'run(torch)'."
        )
    if not callable(run_fn):
        raise ValueError(f"'run' in benches.{bench_id} is not callable.")
    return run_fn