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]()