"""Benchmark configuration management."""

from typing import Literal, Optional, Dict, Any, List
from pydantic import BaseModel, Field
from pathlib import Path
import yaml


class ModelConfig(BaseModel):
    """Configuration for a single model."""
    provider: Literal["stub", "openai", "anthropic", "local"]
    model: str
    weight: float = 1.0
    api_key: Optional[str] = None
    
    
class DatasetConfig(BaseModel):
    """Configuration for benchmark dataset."""
    type: Literal["synthetic", "file"]
    count: int = 100
    categories: List[str] = ["self_referential", "neutral"]
    path: Optional[Path] = None
    

class OutputConfig(BaseModel):
    """Configuration for benchmark output."""
    format: Literal["json", "yaml"] = "json"
    path: str = "outputs/{timestamp}_{name}.json"
    include_traces: bool = True
    

class EvaluatorConfig(BaseModel):
    """Configuration for evaluator."""
    type: Literal["multi_model", "rv_telos", "cuda_gates", "stub"]
    models: List[ModelConfig]
    cuda_enabled: bool = False
    batch_size: int = 1
    

class BenchmarkConfig(BaseModel):
    """Complete benchmark configuration."""
    name: str
    version: str = "v0.1"
    evaluator: EvaluatorConfig
    dataset: DatasetConfig
    output: OutputConfig
    metadata: Dict[str, Any] = Field(default_factory=dict)
    
    @classmethod
    def from_yaml(cls, path: Path) -> "BenchmarkConfig":
        """Load configuration from YAML file."""
        with open(path) as f:
            data = yaml.safe_load(f)
        return cls(**data)
    
    def to_yaml(self, path: Path) -> None:
        """Save configuration to YAML file."""
        with open(path, 'w') as f:
            yaml.dump(self.model_dump(), f, default_flow_style=False)
