# coding=utf-8
__author__ = "Dimitris Karkalousos"

# Taken and adapted from: https://github.com/NVIDIA/NeMo/blob/main/nemo/utils/env_var_parsing.py

import decimal
import json
import os

from dateutil import parser  # type: ignore

__all__ = [
    "get_env",
    "get_envbool",
    "get_envint",
    "get_envfloat",
    "get_envdecimal",
    "get_envdate",
    "get_envdatetime",
    "get_envlist",
    "get_envdict",
    "CoercionError",
    "RequiredSettingMissingError",
]


class CoercionError(Exception):
    """Custom error raised when a value cannot be coerced."""

    def __init__(self, key, value, func):
        """Inits :class:`CoercionError`.

        Parameters
        ----------
        key : str
            The name of the env var that is missing.
        value : str
            The value of the env var that is missing.
        func : function
            The function that was used to try and coerce the value.
        """
        msg = f"Unable to coerce '{key}={value}' using {func.__name__}."
        super().__init__(msg)


class RequiredSettingMissingError(Exception):
    """Custom error raised when a required env var is missing."""

    def __init__(self, key):
        """Inits :class:`RequiredSettingMissingError`.

        Parameters
        ----------
        key : str
            The name of the env var that is missing.
        """
        msg = f"Required env var '{key}' is missing."
        super().__init__(msg)


def _get_env(key, default=None, coerce=lambda x: x, required=False):
    """Return env var coerced into a type other than string. This function extends the standard os.getenv function to
    enable the coercion of values into data types other than string (all env vars are strings by default).

    Parameters
    ----------
    key : str
        The name of the env var to retrieve.
    default : str, (optional)
        The default value to return if the env var is not set. NB the default value is **not** coerced, and is assumed
        to be of the correct type.
    coerce : function, (optional)
        A function that takes a string and returns a value of the desired type.
    required : bool, (optional)
        If True, raises a RequiredSettingMissingError if the env var is not set.

    Returns
    -------
    str
        The value of the env var coerced into the desired type.
    """
    try:
        value = os.environ[key]
    except KeyError as e:
        if required is True:
            raise RequiredSettingMissingError(key) from e
        return default

    try:
        return coerce(value)
    except Exception as exc:
        raise CoercionError(key, value, coerce) from exc


# standard type coercion functions
def _bool(value):
    """Return env var cast as boolean."""
    if isinstance(value, bool):
        return value

    return value is not None and value.lower() not in (
        "false",
        "0",
        "no",
        "n",
        "f",
        "none",
    )


def _int(value):
    """Return env var cast as integer."""
    return int(value)


def _float(value):
    """Return env var cast as float."""
    return float(value)


def _decimal(value):
    """Return env var cast as Decimal."""
    return decimal.Decimal(value)


def _dict(value):
    """Return env var as a dict."""
    return json.loads(value)


def _datetime(value):
    """Return env var as a datetime."""
    return parser.parse(value)


def _date(value):
    """Return env var as a date."""
    return parser.parse(value).date()


def get_env(key, *default, **kwargs):
    """Return env var. This is the parent function of all other get_foo functions, and is responsible for unpacking
    args/kwargs into the values that _get_env expects (it is the root function that actually interacts with environ).

    Parameters
    ----------
    key : str
        The env var name to look up.
    default : str, (optional)
        The value to use if the env var does not exist. If this value is not supplied, then the env var is considered
        to be required, and a RequiredSettingMissingError error will be raised if it does not exist.
    kwargs : dict, (optional)
        Additional keyword arguments to pass to _get_env.

    Returns
    -------
    str
        The env var, coerced if required, and a default if supplied.
    """
    if len(default) not in (0, 1):
        raise AssertionError("Too many args supplied.")
    func = kwargs.get("coerce", lambda x: x)
    required = len(default) == 0
    default = None if required else default[0]
    return _get_env(key, default=default, coerce=func, required=required)


def get_envbool(key, *default):
    """Return env var cast as boolean."""
    return get_env(key, *default, coerce=_bool)


def get_envint(key, *default):
    """Return env var cast as integer."""
    return get_env(key, *default, coerce=_int)


def get_envfloat(key, *default):
    """Return env var cast as float."""
    return get_env(key, *default, coerce=_float)


def get_envdecimal(key, *default):
    """Return env var cast as Decimal."""
    return get_env(key, *default, coerce=_decimal)


def get_envdate(key, *default):
    """Return env var as a date."""
    return get_env(key, *default, coerce=_date)


def get_envdatetime(key, *default):
    """Return env var as a datetime."""
    return get_env(key, *default, coerce=_datetime)


def get_envlist(key, *default, **kwargs):
    """Return env var as a list."""
    separator = kwargs.get("separator", " ")
    return get_env(key, *default, coerce=lambda x: x.split(separator))


def get_envdict(key, *default):
    """Return env var as a dict."""
    return get_env(key, *default, coerce=_dict)
