"""General-purpose helper functions for MLWiz.

Includes dotted-path resolution, config parsing helpers, and atomic dill serialization utilities.
"""

from pydoc import locate
from typing import Tuple, Callable

import os
import dill
from mlwiz.static import ATOMIC_SAVE_EXTENSION


def return_class_and_args(
    config: dict, key: str, return_class_name: bool = False
) -> Tuple[Callable[..., object], dict]:
    r"""
    Returns the class and arguments associated to a specific key in the
    configuration file.

    Args:
        config (dict): the configuration dictionary
        key (str): a string representing a particular class in the
            configuration dictionary
        return_class_name (bool): if ``True``, returns the class name as a
            string rather than the class object

    Returns:
        a tuple (class, dict of arguments), or (None, None) if the key
        is not present in the config dictionary
    """
    if key not in config or config[key] is None:
        return None, None
    elif isinstance(config[key], str):
        return s2c(config[key]), {}
    elif isinstance(config[key], dict):
        return (
            (
                s2c(config[key]["class_name"])
                if not return_class_name
                else config[key]["class_name"]
            ),
            config[key]["args"] if "args" in config[key] else {},
        )
    else:
        raise NotImplementedError(
            f"Parameter {key} has not been formatted properly"
        )


def s2c(class_name: str) -> Callable[..., object]:
    r"""
    Converts a dotted path to the corresponding class

    Args:
         class_name (str): dotted path to class name

    Returns:
        the class to be used
    """
    result = locate(class_name)
    if result is None:
        raise ImportError(
            f"The (dotted) path '{class_name}' is unknown. Check your configuration."
        )
    return result


def atomic_dill_save(data: object, filepath: str) -> object:
    """
    Saves a dill object to a file.
    :param data: the dill object to save
    :param filepath: the path to the dill object to save
    """
    tmp_path = str(filepath) + ATOMIC_SAVE_EXTENSION
    try:
        with open(tmp_path, "wb") as f:
            dill.dump(data, f)
        os.replace(tmp_path, filepath)
    finally:
        # Best-effort cleanup of temp file (if replace failed)
        try:
            os.remove(tmp_path)
        except FileNotFoundError:
            pass


def dill_load(filepath: str) -> object:
    """
    Loads a dill file.
    :param filepath: the path to the dill file
    :return: an object
    """
    with open(filepath, "rb") as file:
        return dill.load(file)
