"""
Configuration for nyx-probing framework.
"""
from pathlib import Path
from pydantic import BaseModel
from typing import Optional
import os


class ModelConfig(BaseModel):
    """Model configuration."""
    name: str = "Qwen/Qwen2.5-7B"
    device: str = "cuda"
    dtype: str = "float16"
    cache_dir: Optional[Path] = None


class ProbeConfig(BaseModel):
    """Probe configuration."""
    max_new_tokens: int = 50
    temperature: float = 0.8
    do_sample: bool = True
    num_runs: int = 5  # For distribution sampling


class StorageConfig(BaseModel):
    """Storage configuration."""
    results_dir: Path = Path("results")
    experiments_dir: Path = Path("experiments")


class Config(BaseModel):
    """Main configuration."""
    model: ModelConfig = ModelConfig()
    probe: ProbeConfig = ProbeConfig()
    storage: StorageConfig = StorageConfig()

    # Paths
    project_root: Path = Path(__file__).parent.parent

    class Config:
        arbitrary_types_allowed = True


# Default config instance
config = Config()


def get_config() -> Config:
    """Get the current configuration."""
    return config