"""Configuration management for SalTTS."""

import os
from pathlib import Path
from typing import Any, Dict, Optional

import yaml
from omegaconf import OmegaConf, DictConfig


class Config:
    """Configuration manager for SalTTS."""

    def __init__(self, config_dir: Optional[str] = None):
        """Initialize configuration manager.

        Args:
            config_dir: Directory containing config files. Defaults to ./config
        """
        if config_dir is None:
            config_dir = Path(__file__).parent.parent.parent / "config"
        self.config_dir = Path(config_dir)
        self._configs: Dict[str, DictConfig] = {}

    def load(self, config_name: str) -> DictConfig:
        """Load a configuration file.

        Args:
            config_name: Name of config file (without .yaml extension)

        Returns:
            Configuration as OmegaConf DictConfig
        """
        if config_name in self._configs:
            return self._configs[config_name]

        config_path = self.config_dir / f"{config_name}.yaml"
        if not config_path.exists():
            raise FileNotFoundError(f"Config file not found: {config_path}")

        with open(config_path, "r") as f:
            config_dict = yaml.safe_load(f)

        config = OmegaConf.create(config_dict)
        self._configs[config_name] = config
        return config

    def load_all(self) -> Dict[str, DictConfig]:
        """Load all configuration files.

        Returns:
            Dictionary mapping config names to DictConfig objects
        """
        config_files = ["model", "runtime", "training", "nmdb_channels", "deployment"]
        for config_name in config_files:
            self.load(config_name)
        return self._configs

    def get(self, config_name: str, key: str, default: Any = None) -> Any:
        """Get a configuration value.

        Args:
            config_name: Name of config file
            key: Dot-separated key path (e.g., "model.stream_transformer.num_layers")
            default: Default value if key not found

        Returns:
            Configuration value
        """
        config = self.load(config_name)
        return OmegaConf.select(config, key, default=default)

    def merge(self, *configs: DictConfig) -> DictConfig:
        """Merge multiple configurations.

        Args:
            *configs: Configuration objects to merge

        Returns:
            Merged configuration
        """
        return OmegaConf.merge(*configs)


def load_config(config_name: str, config_dir: Optional[str] = None) -> DictConfig:
    """Convenience function to load a single config file.

    Args:
        config_name: Name of config file (without .yaml extension)
        config_dir: Directory containing config files

    Returns:
        Configuration as OmegaConf DictConfig
    """
    config_manager = Config(config_dir)
    return config_manager.load(config_name)
