"""
    Hydra configuration management
"""
import collections.abc
import os
import re
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Dict, List, Mapping, Optional, Set, Tuple, Union

import hydra
from omegaconf import DictConfig
from pydantic import BaseModel, ConfigDict, PrivateAttr, model_validator

from hyfi.core import (
    __about__,
    __hydra_config__,
    __hydra_default_config_group_value__,
    __hydra_version_base__,
)
from hyfi.utils.conf import CONFs
from hyfi.utils.logging import LOGGING
from hyfi.utils.packages import PKGs

if level := os.environ.get("HYFI_LOG_LEVEL"):
    LOGGING.setLogger(level)
logger = LOGGING.getLogger(__name__)


class SpecialKeys(str, Enum):
    """Special keys in configs used by HyFI."""

    CALL = "_call_"
    CONFIG = "_config_"
    CONFIG_GROUP = "_config_group_"
    CONFIG_NAME = "_config_name_"
    EXEC = "_exec_"
    FUNC = "_func_"
    KWARGS = "_kwargs_"
    METHOD = "_method_"
    PARTIAL = "_partial_"
    PIPE = "_pipe_"
    RECURSIVE = "_recursive_"
    RUN = "_run_"
    TARGET = "_target_"
    TYPE = "_type_"
    WITH = "_with_"


class Composer(BaseModel, CONFs):
    """
    Compose a configuration by applying overrides
    """

    config_group: Union[str, None] = None
    overrides: Union[List[str], None] = None
    config_data: Union[Dict[str, Any], DictConfig, None] = None
    throw_on_resolution_failure: bool = True
    throw_on_missing: bool = False
    config_name: Union[str, None] = None
    config_module: Union[str, None] = None
    global_package: bool = False
    verbose: bool = False

    _cfg_: DictConfig = PrivateAttr({})

    model_config = ConfigDict(
        arbitrary_types_allowed=True,
        validate_assignment=True,
    )  # type: ignore

    def __init__(self, **args):
        super().__init__(**args)
        self._cfg_ = self.compose(
            config_group=self.config_group,
            overrides=self.overrides,
            config_data=self.config_data,
            throw_on_resolution_failure=self.throw_on_resolution_failure,
            throw_on_missing=self.throw_on_missing,
            root_config_name=self.config_name,
            config_module=self.config_module,
            global_package=self.global_package,
            verbose=self.verbose,
        )

    def compose(
        self,
        config_group: Union[str, None] = None,
        overrides: Union[List[str], None] = None,
        config_data: Union[Dict[str, Any], DictConfig, None] = None,
        throw_on_resolution_failure: bool = True,
        throw_on_missing: bool = False,
        root_config_name: Union[str, None] = None,
        config_module: Union[str, None] = None,
        global_package: bool = False,
        verbose: bool = False,
    ) -> DictConfig:
        """
        Compose a configuration by applying overrides

        Args:
            config_group: Name of the config group to compose (`config_group=name`)
            overrides: List of config groups to apply overrides to (`overrides=["override_name"]`)
            config_data: Keyword arguments to override config group values (will be converted to overrides of the form `config_group.key=value`)
            return_as_dict: Return the result as a dict
            throw_on_resolution_failure: If True throw an exception if resolution fails
            throw_on_missing: If True throw an exception if config_group doesn't exist
            root_config_name: Name of the root config to be used (e.g. `hconf`)
            config_module: Module of the config to be used (e.g. `hyfi.conf`)
            global_package: If True, the config assumed to be a global package
            verbose: If True print configuration to stdout

        Returns:
            A config object or a dictionary with the composed config
        """
        self._cfg_ = Composer._compose(
            config_group=config_group,
            overrides=overrides,
            config_data=config_data,
            throw_on_resolution_failure=throw_on_resolution_failure,
            throw_on_missing=throw_on_missing,
            root_config_name=root_config_name,
            config_module=config_module,
            global_package=global_package,
            verbose=verbose,
        )
        return self._cfg_

    @property
    def config(self) -> DictConfig:
        """
        Returns the composed configuration.
        """
        return self._cfg_

    @property
    def config_as_dict(self) -> Dict:
        """
        Return the configuration as a dictionary.
        """
        return Composer.to_dict(self._cfg_)

    def __call__(
        self,
        config_group: Union[str, None] = None,
        overrides: Union[List[str], None] = None,
        config_data: Union[Dict[str, Any], DictConfig, None] = None,
        throw_on_resolution_failure: bool = True,
        throw_on_missing: bool = False,
        root_config_name: Union[str, None] = None,
        config_module: Union[str, None] = None,
        global_package: bool = False,
        verbose: bool = False,
    ) -> DictConfig:
        return self.compose(
            config_group=config_group,
            overrides=overrides,
            config_data=config_data,
            throw_on_resolution_failure=throw_on_resolution_failure,
            throw_on_missing=throw_on_missing,
            root_config_name=root_config_name,
            config_module=config_module,
            global_package=global_package,
            verbose=verbose,
        )

    def __getitem__(self, key):
        return self._cfg_[key]

    def __iter__(self):
        return iter(self._cfg_)

    def __len__(self):
        return len(self._cfg_)

    def __contains__(self, key):
        return key in self._cfg_

    def __eq__(self, other):
        return self._cfg_ == other

    def __ne__(self, other):
        return self._cfg_ != other

    def __bool__(self):
        return bool(self._cfg_)

    def __hash__(self):
        return hash(self._cfg_)

    def __getstate__(self):
        return self._cfg_

    @staticmethod
    def hydra_compose(
        root_config_name: Union[str, None] = None,
        config_module: Union[str, None] = None,
        overrides: Union[List[str], None] = None,
    ):
        is_initialized = hydra.core.global_hydra.GlobalHydra.instance().is_initialized()  # type: ignore
        config_module = config_module or __hydra_config__.hyfi_config_module
        logger.debug("config_module: %s", config_module)
        if is_initialized:
            # Hydra is already initialized.
            logger.debug("Hydra is already initialized")
            cfg = hydra.compose(config_name=root_config_name, overrides=overrides)
        else:
            with hydra.initialize_config_module(
                config_module=config_module, version_base=__hydra_version_base__
            ):
                cfg = hydra.compose(config_name=root_config_name, overrides=overrides)
        if is_initialized:
            cfg = hydra.compose(config_name=root_config_name, overrides=overrides)
        else:
            with hydra.initialize_config_module(
                config_module=config_module, version_base=__hydra_version_base__
            ):
                cfg = hydra.compose(config_name=root_config_name, overrides=overrides)
        return cfg

    @staticmethod
    def split_config_group(
        config_group: Union[str, None] = None,
    ) -> Tuple[str, str, str]:
        if config_group:
            group_ = config_group.split("=")
            # group_key group_value group_key group_value group_key group_value default
            if len(group_) == 2:
                group_key, group_value = group_
            else:
                group_key = group_[0]
                group_value = __hydra_default_config_group_value__
            config_group = f"{group_key}={group_value}"
        else:
            group_key = ""
            group_value = ""
            config_group = ""
        return config_group, group_key, group_value

    @staticmethod
    def _compose_as_dict(
        config_group: Union[str, None] = None,
        overrides: Union[List[str], None] = None,
        config_data: Union[Dict[str, Any], DictConfig, None] = None,
        throw_on_resolution_failure: bool = True,
        throw_on_missing: bool = False,
        root_config_name: Union[str, None] = None,
        config_module: Union[str, None] = None,
        global_package: bool = False,
        **kwargs,
    ) -> Dict:
        return Composer.to_dict(
            Composer._compose(
                config_group=config_group,
                overrides=overrides,
                config_data=config_data,
                throw_on_resolution_failure=throw_on_resolution_failure,
                throw_on_missing=throw_on_missing,
                root_config_name=root_config_name,
                config_module=config_module,
                global_package=global_package,
                **kwargs,
            )
        )

    @staticmethod
    def _compose(
        config_group: Union[str, None] = None,
        overrides: Union[List[str], None] = None,
        config_data: Union[Dict[str, Any], DictConfig, None] = None,
        throw_on_resolution_failure: bool = True,
        throw_on_missing: bool = False,
        root_config_name: Union[str, None] = None,
        config_module: Union[str, None] = None,
        global_package: bool = False,
        **kwargs,
    ) -> DictConfig:
        if isinstance(config_data, DictConfig):
            logger.debug("returning config_group_kwargs without composing")
            return config_data
        # Set overrides to the empty list if None
        if overrides is None:
            overrides = []
        # Set the group key and value of the config group.
        config_group, group_key, group_value = Composer.split_config_group(config_group)
        # If group_key and group_value are specified in the configuration file.
        if group_key and group_value:
            # Initialize hydra configuration module.
            cfg = Composer.hydra_compose(
                root_config_name=root_config_name,
                config_module=config_module,
                overrides=overrides,
            )
            cfg = Composer.select(
                cfg,
                key=group_key,
                default=None,
                throw_on_missing=False,
                throw_on_resolution_failure=False,
            )
            override = config_group if cfg is not None else f"+{config_group}"
            # Add override to overrides list.
            if isinstance(override, str):
                if overrides:
                    overrides.append(override)
                else:
                    overrides = [override]
        # if verbose:
        logger.debug(f"compose config with overrides: {overrides}")
        # Initialize hydra and return the configuration.
        cfg = Composer.hydra_compose(
            root_config_name=root_config_name,
            config_module=config_module,
            overrides=overrides,
        )
        # Add config group overrides to overrides list.
        group_overrides: List[str] = []
        group_cfg = Composer.select(
            cfg,
            key=group_key,
            default=None,
            throw_on_missing=False,
            throw_on_resolution_failure=False,
        )
        if config_data and group_cfg:
            group_overrides.extend(
                f"{group_key}.{k}={v}"
                for k, v in config_data.items()
                if isinstance(v, (str, int, float, bool)) and k in group_cfg
            )
        if group_overrides:
            overrides.extend(group_overrides)
            cfg = Composer.hydra_compose(
                root_config_name=root_config_name,
                config_module=config_module,
                overrides=overrides,
            )

        # Select the group_key from the configuration.
        if group_key and not global_package:
            cfg = Composer.select(
                cfg,
                key=group_key,
                default=None,
                throw_on_missing=throw_on_missing,
                throw_on_resolution_failure=throw_on_resolution_failure,
            )
        # logger.debug("Composed config: %s", OmegaConf.to_yaml(_to_dict(cfg)))
        return cfg

    @staticmethod
    def is_instantiatable(cfg: Any):
        """
        Determines whether the input configuration object is instantiatable.

        Args:
            cfg (Any): The configuration object to check.

        Returns:
            bool: True if the configuration object is instantiatable, False otherwise.
        """
        return Composer.is_config(cfg) and SpecialKeys.TARGET in cfg

    @staticmethod
    def replace_special_keys(_dict: Mapping[str, Any]) -> Mapping:
        """
        Replace special keys in a dictionary.

        Args:
            _dict (Mapping[str, Any]): The dictionary to update.

        Returns:
            Mapping: The updated dictionary.
        """
        _new_dict = {}
        for k, v in _dict.items():
            key = Composer.generate_alias_for_special_keys(k)
            if isinstance(v, collections.abc.Mapping):
                _new_dict[key] = Composer.replace_special_keys(v)
            else:
                _new_dict[key] = v
        return _new_dict

    @staticmethod
    def generate_alias_for_special_keys(key: str) -> str:
        """
        Generate an alias for special keys.
        _with_ -> run_with
        _pipe_ -> run_pipe
        _run_ -> run

        Args:
            key (str): The special key to generate an alias for.

        Returns:
            str: The alias for the special key.
        """
        # replace the exact `with`, `pipe` with `run_with`, `run_pipe`
        key_ = re.sub(r"^with$", "run_with", key)
        # replace the prefix `_` with `run_`
        key_ = re.sub(r"^_with_$", "run_with", key_)
        key_ = re.sub(r"^_pipe_$", "pipe_target", key_)
        key_ = re.sub(r"^_run_$", "run", key_)
        return key_

    @staticmethod
    def partial(
        config: Union[str, Dict],
        *args: Any,
        **kwargs: Any,
    ) -> Callable:
        """
        Returns a callable object that is a partial version of the function or class specified in the config object.

        Args:
            config: An config object describing what to call and what params to use.
                    In addition to the parameters, the config must contain:
                    _target_ : target class or callable name (str)
                    And may contain:
                    _partial_: If True, return functools.partial wrapped method or object
                                False by default. Configure per target.
            args: Optional positional parameters pass-through
            kwargs: Optional named parameters to override
                    parameters in the config object. Parameters not present
                    in the config objects are being passed as is to the target.
                    IMPORTANT: dataclasses instances in kwargs are interpreted as config
                                and cannot be used as passthrough
        Returns:
            A callable object that is a partial version of the function or class specified in the config object.
        """
        if isinstance(config, str):
            config = {SpecialKeys.TARGET.value: config}
        else:
            config = Composer.to_dict(config)
        if not isinstance(config, dict):
            raise ValueError("config must be a dict or a str")
        config[SpecialKeys.PARTIAL.value] = True
        rc_kwargs_ = config.pop(SpecialKeys.KWARGS, {})
        if rc_kwargs_ and kwargs:
            kwargs.update(rc_kwargs_)
        return Composer.instantiate(config, *args, **kwargs)

    @staticmethod
    def instantiate(config: Any, *args: Any, **kwargs: Any) -> Any:
        """
        Instantiates an object using the provided config object.

        Args:
            config: An config object describing what to call and what params to use.
                    In addition to the parameters, the config must contain:
                    _target_ : target class or callable name (str)
                    And may contain:
                    _args_: List-like of positional arguments to pass to the target
                    _recursive_: Construct nested objects as well (bool).
                                    False by default.
                                    may be overridden via a _recursive_ key in
                                    the kwargs
                    _convert_: Conversion strategy
                            none    : Passed objects are DictConfig and ListConfig, default
                            partial : Passed objects are converted to dict and list, with
                                    the exception of Structured Configs (and their fields).
                            all     : Passed objects are dicts, lists and primitives without
                                    a trace of OmegaConf containers
                    _partial_: If True, return functools.partial wrapped method or object
                                False by default. Configure per target.
                    _args_: List-like of positional arguments
            args: Optional positional parameters pass-through
            kwargs: Optional named parameters to override
                    parameters in the config object. Parameters not present
                    in the config objects are being passed as is to the target.
                    IMPORTANT: dataclasses instances in kwargs are interpreted as config
                                and cannot be used as passthrough

        Returns:
            if _target_ is a class name: the instantiated object
            if _target_ is a callable: the return value of the call
        """
        verbose = config.get("verbose", False)
        if not Composer.is_instantiatable(config):
            if verbose:
                logger.info("Config is not instantiatable, returning config")
            return config
        _recursive_ = config.get(SpecialKeys.RECURSIVE, False)
        if SpecialKeys.RECURSIVE not in kwargs:
            kwargs[SpecialKeys.RECURSIVE.value] = _recursive_
        if verbose:
            logger.info("instantiating %s ...", config.get(SpecialKeys.TARGET))
        return hydra.utils.instantiate(config, *args, **kwargs)

    @staticmethod
    def getsource(obj: Any) -> str:
        """
        Return the source code of the object.

        Args:
            obj: The object to get the source code of.

        Returns:
            The source code of the object as a string.

        """
        try:
            target_string = ""
            if Composer.is_config(obj):
                if SpecialKeys.TARGET in obj:
                    target_string = obj[SpecialKeys.TARGET]
            elif isinstance(obj, str):
                target_string = obj
            return PKGs.getsource(target_string) if target_string else ""
        except Exception as e:
            logger.error(f"Error getting source: {e}")
            return ""

    @staticmethod
    def viewsource(obj: Any):
        """
        Print the source code of the object.

        Args:
            obj: The object to print the source code of.

        """
        print(Composer.getsource(obj))


class BaseConfig(BaseModel):
    """
    Base class for all config classes.
    """

    _config_name_: str = "__init__"
    _config_group_: str = ""
    verbose: bool = False

    _init_args_: Dict[str, Any] = {}
    _exclude_: Set[str] = set()
    _property_set_methods_: Dict[str, str] = {}
    _subconfigs_: Dict[str, Any] = {}

    model_config = ConfigDict(
        arbitrary_types_allowed=True,
        extra="allow",
        validate_assignment=False,
    )  # type: ignore

    def __init__(self, **config_kwargs):
        logger.debug(
            "init %s with %s args", self.__class__.__name__, len(config_kwargs)
        )
        super().__init__(**config_kwargs)
        self.initialize_subconfigs(config_kwargs)

    def __setattr__(self, key, val):
        """
        Overrides the default __setattr__ method to allow for custom property set methods.

        Args:
            key (str): The name of the attribute to set.
            val (Any): The value to set the attribute to.
        """
        if method := self._property_set_methods_.get(key):  # type: ignore
            logger.info(
                "Setting %s to %s",
                key,
                val if isinstance(val, (str, int)) else type(val),
            )
            getattr(self, method)(val)
        super().__setattr__(key, val)

    @model_validator(mode="before")
    def validate_model_config_before(cls, data):
        # logger.debug("Validating model config before validating each field.")
        _config_name_ = data.get("_config_name_", getattr(cls._config_name_, "default", "__init__"))  # type: ignore
        _config_group_ = data.get("_config_group_", getattr(cls._config_group_, "default"))  # type: ignore
        _class_name_ = cls.__name__  # type: ignore
        if not _config_group_:
            logger.debug("There is no config group specified.")
            return data
        # Initialize the config with the given config_name.
        logger.info(
            "Composing `%s` class with `%s` config in `%s` group.",
            _class_name_,
            _config_name_,
            _config_group_,
        )
        cfg = Composer(
            config_group=f"{_config_group_}={_config_name_}",
            config_data=data,
        ).config_as_dict
        data = Composer.update(cfg, data)
        # Exclude any attributes specified in the class's `exclude` list.
        exclude = getattr(cls._exclude_, "default", set())  # type: ignore
        for name in exclude:
            if name in data:
                del data[name]
        return data

    # @model_validator(mode="after")  # type: ignore
    # def validate_model_config_after(cls, model):
    #     logger.debug("validate_model_config_after")
    #     return model

    def initialize_subconfigs(self, config_kwargs: Dict[str, Any]):
        """
        Initializes subconfigs with the given config data.
        The function updates the object's dictionary with the given config data,
        after excluding any attributes specified in the object's `exclude` list.

        Args:
            config_kwargs (Dict[str, Any]): The config data to initialize the subconfigs with.

        Returns:
            None
        """
        self._subconfigs_ = self._subconfigs_ or {}
        for name, config in self._subconfigs_.items():
            if name in config_kwargs and isinstance(config_kwargs[name], dict):
                cfg = config_kwargs[name]
                logger.debug("Initializing subconfig %s with %s", name, cfg)
                setattr(self, name, config.model_validate(cfg))

    def export_config(
        self,
        exclude: Optional[Union[str, List[str], Set[str], None]] = None,
        exclude_none: bool = True,
        only_include: Optional[Union[str, List[str], Set[str], None]] = None,
    ) -> Dict[str, Any]:
        """
        Export the configuration to a dictionary.

        Args:
            exclude (Optional[Union[str, List[str], Set[str], None]]): Keys to exclude from the saved configuration.
                Defaults to None.
            exclude_none (bool): Whether to exclude keys with None values. Defaults to True.
            only_include (Optional[Union[str, List[str], Set[str], None]]): Keys to include in the saved configuration.
                Defaults to None.

        Returns:
            Dict[str, Any]: The configuration dictionary.
        """
        if not exclude:
            exclude = self._exclude_  # type: ignore
        if isinstance(exclude, str):
            exclude = [exclude]
        if exclude is None:
            exclude = []
        if isinstance(only_include, str):
            only_include = [only_include]
        if only_include is None:
            only_include = []

        config = self.model_dump(exclude=exclude, exclude_none=exclude_none)  # type: ignore
        if only_include:
            config = {key: config[key] for key in only_include if key in config}

        return config

    def save_config(
        self,
        filepath: Union[str, Path],
        exclude: Optional[Union[str, List[str], Set[str], None]] = None,
        exclude_none: bool = True,
        only_include: Optional[Union[str, List[str], Set[str], None]] = None,
    ) -> str:
        """
        Save the batch configuration to file.

        Args:
            filepath ([Union[str, Path]): The filepath to save the configuration to.
            exclude (Optional[Union[str, List[str], Set[str], None]]): Keys to exclude from the saved configuration.
                Defaults to None.
            exclude_none (bool): Whether to exclude keys with None values. Defaults to True.
            only_include (Optional[Union[str, List[str], Set[str], None]]): Keys to include in the saved configuration.
                Defaults to None.

        Returns:
            str: The filename of the saved configuration.
        """
        logger.info("Saving config to %s", filepath)

        config_to_save = self.export_config(
            exclude=exclude, exclude_none=exclude_none, only_include=only_include
        )

        Composer.save(config_to_save, filepath)
        return str(filepath)

    def save_config_as_json(
        self,
        filepath: Union[str, Path],
        exclude: Optional[Union[str, List[str], Set[str], None]] = None,
        exclude_none: bool = True,
        only_include: Optional[Union[str, List[str], Set[str], None]] = None,
    ) -> str:
        def dumper(obj):
            return Composer.to_dict(obj) if isinstance(obj, DictConfig) else str(obj)

        config_to_save = self.export_config(
            exclude=exclude, exclude_none=exclude_none, only_include=only_include
        )
        logger.info("Saving config to %s", filepath)
        Composer.save_json(config_to_save, filepath, default=dumper)
        return str(filepath)

    def print_config(
        self,
    ):
        Composer.print(self.model_dump())
