import yaml
from pathlib import Path
from typing import Union, Dict, Any
from emissary_ml_toolkit.config.schemas import TrainingJobConfig


class ConfigLoader:
    """Handles loading and parsing YAML configuration files"""

    @staticmethod
    def load_yaml(path: Union[str, Path]) -> Dict[str, Any]:
        """Load YAML file and return as dictionary"""
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"Configuration file not found: {path}")

        with open(path, 'r') as f:
            try:
                return yaml.safe_load(f)
            except yaml.YAMLError as e:
                raise ValueError(f"Invalid YAML format: {e}")

    @classmethod
    def from_yaml(cls, path:Union[str, Path]) -> TrainingJobConfig:
        """Load and validate training job configuration from YAML"""
        config_dict = cls.load_yaml(path)
        return TrainingJobConfig(**config_dict)

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> TrainingJobConfig:
        """Load and validate training job configuration from DICT"""
        return TrainingJobConfig(**config_dict)

    @classmethod
    def from_json(cls, path: Union[str, Path]) -> TrainingJobConfig:
        # We can later support multiple type of training configuration if needed
        raise NotImplementedError("Currently not supporting configuration in json format")