"""Private module with AST parser for safe evaluation."""
import ast
from typing import Any, Callable, Dict, List

import yaml


def _process_node(node: Any, flat_dict: dict) -> Any:
    """Compute an AST from the root by replacing param name by their values.

    The AST can contain any parameter name of the configuration.
    The most usefull operators and built-in functions are supported,
    the random and math packages are also supported as well as some (safe)
    numpy functions. If/else statements and comprehension lists are also
    supported.
    """
    # Case None, bool or number
    if isinstance(node, ast.Constant):
        return node.n
    functions: Dict[Any, Callable[[Any, dict], Any]] = {
        ast.BinOp: _process_binop,  # binary operation
        ast.BoolOp: _process_boolop,  # boolean operation
        ast.Compare: _process_comparator,  # comparison
        ast.Name: _process_param_name,  # parameter name
        ast.Attribute: _process_subconfig,  # sub-config
        ast.IfExp: _process_ifexp,  # if/else
        ast.List: _process_ltd,  # list
        ast.Tuple: _process_ltd,  # tuple
        ast.Dict: _process_ltd,  # dict
        ast.Call: _process_call,  # function
        ast.ListComp: _process_listcomp,  # comprehension list
        ast.comprehension: _process_comprehension,  # comprehension
    }
    if isinstance(node, tuple(functions.keys())):
        return functions[type(node)](node, flat_dict)
    raise ValueError(f"Not supported node in expression (of type {type(node)}).")


def _process_binop(node: Any, flat_dict: dict) -> Any:
    """Process a binary operator node."""
    left_val = _process_node(node=node.left, flat_dict=flat_dict)
    right_val = _process_node(node=node.right, flat_dict=flat_dict)
    operators: Dict[Any, Callable[[Any, Any], Any]] = {
        ast.Add: lambda x, y: x + y,
        ast.Sub: lambda x, y: x - y,
        ast.Mult: lambda x, y: x * y,
        ast.Div: lambda x, y: x / y,
        ast.Pow: lambda x, y: x**y,
        ast.FloorDiv: lambda x, y: x // y,
        ast.Mod: lambda x, y: x % y,
        ast.BitOr: lambda x, y: x or y,
        ast.BitAnd: lambda x, y: x and y,
        ast.MatMult: lambda x, y: x @ y,
    }
    if isinstance(node.op, tuple(operators.keys())):
        return operators[type(node.op)](left_val, right_val)
    raise ValueError(
        f"Invalid operator detected: {node.op}."
        "Please use only these ops: '+', '-', '*', '/', '**', "
        "'//', '%', '|', '&'."
    )


def _process_boolop(node: Any, flat_dict: dict) -> Any:
    """Process bool operator node."""
    values = [
        _process_node(node=nodeval, flat_dict=flat_dict) for nodeval in node.values
    ]
    operators: Dict[Any, Callable[[Any, Any], Any]] = {
        ast.And: lambda x, y: x and y,
        ast.Or: lambda x, y: x or y,
    }
    operator = operators[type(node.op)]
    result = isinstance(node.op, ast.And)  # Neutral for the operation
    for val in values:
        result = operator(result, val)  # type: ignore
    return result


def _process_comparator(node: Any, flat_dict: dict) -> Any:
    """Process a comparator node."""
    left_val = _process_node(node=node.left, flat_dict=flat_dict)
    right_val = _process_node(node=node.comparators[0], flat_dict=flat_dict)
    operators: Dict[Any, Callable[[bool, bool], bool]] = {
        ast.Eq: lambda x, y: x == y,
        ast.NotEq: lambda x, y: x != y,
        ast.Lt: lambda x, y: x < y,
        ast.LtE: lambda x, y: x <= y,
        ast.Gt: lambda x, y: x > y,
        ast.GtE: lambda x, y: x >= y,
    }
    return operators[type(node.ops[0])](left_val, right_val)  # type: ignore


def _process_param_name(node: Any, flat_dict: dict) -> Any:
    """Process a parameter name."""
    name = node.id
    if name in flat_dict:
        if isinstance(flat_dict[name], (bool, int, float, complex, list, type(None))):
            return flat_dict[name]
        raise ValueError(
            f"Invalid value in expression for parameter '{name}', "
            f"found type {type(flat_dict[name])}, expected, None, bool, int, "
            "float, complex or list."
        )
    raise ValueError(f"Unknown parameter '{name}'.")


def _process_subconfig(node: Any, flat_dict: dict) -> Any:
    """Process a sub-config."""
    # Get recursively the attribute names
    param_name = ""
    while isinstance(node, ast.Attribute):
        param_name = f"{node.attr}.{param_name}"
        node = node.value
    # Add the global subconfig name and process the node containing the
    # full flat parameter name
    param_name = f"{node.id}.{param_name}"[:-1]  # Remove last dot
    node = ast.Name(id=param_name, ctx=ast.Load())
    return _process_node(node=node, flat_dict=flat_dict)


def _process_ifexp(node: Any, flat_dict: dict) -> Any:
    """Process a if/exp statement."""
    test_val = _process_node(node=node.test, flat_dict=flat_dict)
    return (
        _process_node(node=node.body, flat_dict=flat_dict)
        if test_val
        else _process_node(node=node.orelse, flat_dict=flat_dict)
    )


def _process_ltd(node: Any, flat_dict: dict) -> Any:
    """Process a list, a tuple or a dict node."""
    if isinstance(node, ast.List):
        # List
        return [
            _process_node(node=element, flat_dict=flat_dict) for element in node.elts
        ]
    if isinstance(node, ast.Tuple):
        # Tuple
        return tuple(
            _process_node(node=element, flat_dict=flat_dict) for element in node.elts
        )
    # Dict
    return {
        _process_node(node=key, flat_dict=flat_dict): _process_node(
            node=value, flat_dict=flat_dict
        )
        for (key, value) in zip(node.keys, node.values)
    }


def _process_call(node: Any, flat_dict: dict) -> Any:
    """Process a function node."""
    functions: Dict[str, Callable] = {
        "len": len,
        "sum": sum,
        "max": max,
        "min": min,
        "abs": abs,
        "round": round,
        "all": all,
        "any": any,
        "range": range,
        "bool": bool,
        "int": int,
        "float": float,
        "str": str,
        "list": list,
        "tuple": tuple,
        "dict": dict,
        "set": set,
    }
    args = [_process_node(node=arg, flat_dict=flat_dict) for arg in node.args]
    kwargs = {
        kwarg.arg: _process_node(node=kwarg.value, flat_dict=flat_dict)
        for kwarg in node.keywords
    }
    if isinstance(node.func, ast.Name):
        func_name = node.func.id
        value = functions[func_name](*args, **kwargs)
        return value

    # ast.Attribute
    func = _find_function(node=node.func, flat_dict=flat_dict)
    return func(*args, **kwargs)


def _find_function(node: Any, flat_dict: dict) -> Callable:
    """Find a function from node."""
    list_names: List[str] = []
    while isinstance(node, ast.Attribute):
        list_names = [str(node.attr)] + list_names
        node = node.value
    list_names = [node.id] + list_names

    aliases = {"np": "numpy", "tf": "tensorflow"}
    if list_names[0] in aliases:
        list_names[0] = aliases[list_names[0]]
    list_names = (
        ["jax", "numpy"] + list_names[1:] if list_names[0] == "jnp" else list_names
    )

    if list_names[0] in flat_dict:
        obj = flat_dict[list_names[0]]
    elif _filter_allowed(list_names=list_names):
        obj = __import__(list_names[0])
    else:
        raise ValueError(
            f"Package or function not allowed or not supported: {'.'.join(list_names)}"
        )
    for name in list_names[1:-1]:
        obj = getattr(obj, name)
    func = getattr(obj, list_names[-1])
    return func


def _filter_allowed(list_names: List[str]) -> bool:
    """Filter the allowed functions."""
    with open(
        "cliconfig/processing/allowed_functions.yaml", encoding="utf-8"
    ) as yaml_file:
        allowed_funcs = yaml.safe_load(yaml_file)
    list_names = list_names.copy()
    # jax and numpy share the same allowed functions, jax.random is also allowed
    list_names = list_names[1:] if list_names[0] == "jax" else list_names

    if any(name.startswith("_") for name in list_names):
        return False
    if list_names[0] in ("random", "math"):
        return True
    if list_names[0] in ("numpy", "torch", "tensorflow"):
        return list_names[1] in allowed_funcs[list_names[0]]
    return False


def _process_listcomp(node: Any, flat_dict: dict) -> Any:
    """Process comprehension list node."""
    elt = node.elt
    generator = node.generators[0]
    result = []
    for variables in _process_node(generator, flat_dict=flat_dict):
        result.append(_process_node(elt, flat_dict=variables))
    return result


def _process_comprehension(node: Any, flat_dict: dict) -> Any:
    """Process comprehension node."""
    target = node.target
    iterator = node.iter
    tests = node.ifs
    if isinstance(target, ast.Tuple):
        target = tuple(elt for elt in target.elts)

    variables = flat_dict.copy()
    for val in _process_node(node=iterator, flat_dict=flat_dict):
        if isinstance(target, tuple):
            variables.update({target.id: val[i] for i, target in enumerate(target)})
        else:
            variables[target.id] = val
        condition = all(_process_node(test, variables) for test in tests)
        if condition:
            yield variables
