from typing import Type

from benchmark.draw_bench import DrawBenchPrompts
from benchmark.genai_bench import GenAIBenchPrompts
from benchmark.geneval import GenEvalPrompts
from benchmark.hps import HPSPrompts
from benchmark.parti import PartiPrompts


def create_benchmark(
    benchmark_type: str,
) -> Type[
    DrawBenchPrompts | GenAIBenchPrompts | GenEvalPrompts | HPSPrompts | PartiPrompts
]:
    """
    Factory function to create benchmark instances.

    Args:
        benchmark_type (str): The type of benchmark to create. Must be one of:
            - "draw_bench"
            - "genai_bench"
            - "geneval"
            - "hps"
            - "parti"

    Returns:
        An instance of the requested benchmark implementation

    Raises:
        ValueError: If an invalid benchmark type is provided
    """
    benchmark_map = {
        "draw_bench": DrawBenchPrompts,
        "genai_bench": GenAIBenchPrompts,
        "geneval": GenEvalPrompts,
        "hps": HPSPrompts,
        "parti": PartiPrompts,
    }

    if benchmark_type not in benchmark_map:
        raise ValueError(
            f"Invalid benchmark type: {benchmark_type}. Must be one of {list(benchmark_map.keys())}"
        )

    return benchmark_map[benchmark_type]()