from typing import Callable, Set, Tuple, NamedTuple, Literal, Union, TYPE_CHECKING
import collections
import sys
import ast
from functools import cached_property
from dataclasses import dataclass, field

from IPython.utils.docs import GENERATING_DOCUMENTATION


if TYPE_CHECKING or GENERATING_DOCUMENTATION:
    from typing_extensions import Protocol
else:
    # do not require on runtime
    Protocol = object  # requires Python >=3.8


class HasGetItem(Protocol):
    def __getitem__(self, key) -> None:
        ...


class InstancesHaveGetItem(Protocol):
    def __call__(self) -> HasGetItem:
        ...


class HasGetAttr(Protocol):
    def __getattr__(self, key) -> None:
        ...


class DoesNotHaveGetAttr(Protocol):
    pass


# By default `__getattr__` is not explicitly implemented on most objects
MayHaveGetattr = Union[HasGetAttr, DoesNotHaveGetAttr]


def unbind_method(func: Callable) -> Union[Callable, None]:
    """Get unbound method for given bound method.

    Returns None if cannot get unbound method."""
    owner = getattr(func, "__self__", None)
    owner_class = type(owner)
    name = getattr(func, "__name__", None)
    instance_dict_overrides = getattr(owner, "__dict__", None)
    if (
        owner is not None
        and name
        and (
            not instance_dict_overrides
            or (instance_dict_overrides and name not in instance_dict_overrides)
        )
    ):
        return getattr(owner_class, name)


@dataclass
class EvaluationPolicy:
    allow_locals_access: bool = False
    allow_globals_access: bool = False
    allow_item_access: bool = False
    allow_attr_access: bool = False
    allow_builtins_access: bool = False
    allow_any_calls: bool = False
    allowed_calls: Set[Callable] = field(default_factory=set)

    def can_get_item(self, value, item):
        return self.allow_item_access

    def can_get_attr(self, value, attr):
        return self.allow_attr_access

    def can_call(self, func):
        if self.allow_any_calls:
            return True

        if func in self.allowed_calls:
            return True

        owner_method = unbind_method(func)
        if owner_method and owner_method in self.allowed_calls:
            return True


def has_original_dunder_external(
    value,
    module_name,
    access_path,
    method_name,
):
    try:
        if module_name not in sys.modules:
            return False
        member_type = sys.modules[module_name]
        for attr in access_path:
            member_type = getattr(member_type, attr)
        value_type = type(value)
        if type(value) == member_type:
            return True
        if isinstance(value, member_type):
            method = getattr(value_type, method_name, None)
            member_method = getattr(member_type, method_name, None)
            if member_method == method:
                return True
    except (AttributeError, KeyError):
        return False


def has_original_dunder(
    value, allowed_types, allowed_methods, allowed_external, method_name
):
    # note: Python ignores `__getattr__`/`__getitem__` on instances,
    # we only need to check at class level
    value_type = type(value)

    # strict type check passes → no need to check method
    if value_type in allowed_types:
        return True

    method = getattr(value_type, method_name, None)

    if not method:
        return None

    if method in allowed_methods:
        return True

    for module_name, *access_path in allowed_external:
        if has_original_dunder_external(value, module_name, access_path, method_name):
            return True

    return False


@dataclass
class SelectivePolicy(EvaluationPolicy):
    allowed_getitem: Set[HasGetItem] = field(default_factory=set)
    allowed_getitem_external: Set[Tuple[str, ...]] = field(default_factory=set)
    allowed_getattr: Set[MayHaveGetattr] = field(default_factory=set)
    allowed_getattr_external: Set[Tuple[str, ...]] = field(default_factory=set)

    def can_get_attr(self, value, attr):
        has_original_attribute = has_original_dunder(
            value,
            allowed_types=self.allowed_getattr,
            allowed_methods=self._getattribute_methods,
            allowed_external=self.allowed_getattr_external,
            method_name="__getattribute__",
        )
        has_original_attr = has_original_dunder(
            value,
            allowed_types=self.allowed_getattr,
            allowed_methods=self._getattr_methods,
            allowed_external=self.allowed_getattr_external,
            method_name="__getattr__",
        )
        # Many objects do not have `__getattr__`, this is fine
        if has_original_attr is None and has_original_attribute:
            return True

        # Accept objects without modifications to `__getattr__` and `__getattribute__`
        return has_original_attr and has_original_attribute

    def get_attr(self, value, attr):
        if self.can_get_attr(value, attr):
            return getattr(value, attr)

    def can_get_item(self, value, item):
        """Allow accessing `__getiitem__` of allow-listed instances unless it was not modified."""
        return has_original_dunder(
            value,
            allowed_types=self.allowed_getitem,
            allowed_methods=self._getitem_methods,
            allowed_external=self.allowed_getitem_external,
            method_name="__getitem__",
        )

    @cached_property
    def _getitem_methods(self) -> Set[Callable]:
        return self._safe_get_methods(self.allowed_getitem, "__getitem__")

    @cached_property
    def _getattr_methods(self) -> Set[Callable]:
        return self._safe_get_methods(self.allowed_getattr, "__getattr__")

    @cached_property
    def _getattribute_methods(self) -> Set[Callable]:
        return self._safe_get_methods(self.allowed_getattr, "__getattribute__")

    def _safe_get_methods(self, classes, name) -> Set[Callable]:
        return {
            method
            for class_ in classes
            for method in [getattr(class_, name, None)]
            if method
        }


class DummyNamedTuple(NamedTuple):
    pass


class EvaluationContext(NamedTuple):
    locals_: dict
    globals_: dict
    evaluation: Literal[
        "forbidden", "minimal", "limitted", "unsafe", "dangerous"
    ] = "forbidden"
    in_subscript: bool = False


class IdentitySubscript:
    def __getitem__(self, key):
        return key


IDENTITY_SUBSCRIPT = IdentitySubscript()
SUBSCRIPT_MARKER = "__SUBSCRIPT_SENTINEL__"


class GuardRejection(ValueError):
    pass


def guarded_eval(code: str, context: EvaluationContext):
    locals_ = context.locals_

    if context.evaluation == "forbidden":
        raise GuardRejection("Forbidden mode")

    # note: not using `ast.literal_eval` as it does not implement
    # getitem at all, for example it fails on simple `[0][1]`

    if context.in_subscript:
        # syntatic sugar for ellipsis (:) is only available in susbcripts
        # so we need to trick the ast parser into thinking that we have
        # a subscript, but we need to be able to later recognise that we did
        # it so we can ignore the actual __getitem__ operation
        if not code:
            return tuple()
        locals_ = locals_.copy()
        locals_[SUBSCRIPT_MARKER] = IDENTITY_SUBSCRIPT
        code = SUBSCRIPT_MARKER + "[" + code + "]"
        context = EvaluationContext(**{**context._asdict(), **{"locals_": locals_}})

    if context.evaluation == "dangerous":
        return eval(code, context.globals_, context.locals_)

    expression = ast.parse(code, mode="eval")

    return eval_node(expression, context)


def eval_node(node: Union[ast.AST, None], context: EvaluationContext):
    """
    Evaluate AST node in provided context.

    Applies evaluation restrictions defined in the context.

    Currently does not support evaluation of functions with arguments.

    Does not evaluate actions which always have side effects:
    - class definitions (``class sth: ...``)
    - function definitions (``def sth: ...``)
    - variable assignments (``x = 1``)
    - augumented assignments (``x += 1``)
    - deletions (``del x``)

    Does not evaluate operations which do not return values:
    - assertions (``assert x``)
    - pass (``pass``)
    - imports (``import x``)
    - control flow
       - conditionals (``if x:``) except for terenary IfExp (``a if x else b``)
       - loops (``for`` and `while``)
       - exception handling

    The purpose of this function is to guard against unwanted side-effects;
    it does not give guarantees on protection from malicious code execution.
    """
    policy = EVALUATION_POLICIES[context.evaluation]
    if node is None:
        return None
    if isinstance(node, ast.Expression):
        return eval_node(node.body, context)
    if isinstance(node, ast.BinOp):
        # TODO: add guards
        left = eval_node(node.left, context)
        right = eval_node(node.right, context)
        if isinstance(node.op, ast.Add):
            return left + right
        if isinstance(node.op, ast.Sub):
            return left - right
        if isinstance(node.op, ast.Mult):
            return left * right
        if isinstance(node.op, ast.Div):
            return left / right
        if isinstance(node.op, ast.FloorDiv):
            return left // right
        if isinstance(node.op, ast.Mod):
            return left % right
        if isinstance(node.op, ast.Pow):
            return left**right
        if isinstance(node.op, ast.LShift):
            return left << right
        if isinstance(node.op, ast.RShift):
            return left >> right
        if isinstance(node.op, ast.BitOr):
            return left | right
        if isinstance(node.op, ast.BitXor):
            return left ^ right
        if isinstance(node.op, ast.BitAnd):
            return left & right
        if isinstance(node.op, ast.MatMult):
            return left @ right
    if isinstance(node, ast.Constant):
        return node.value
    if isinstance(node, ast.Index):
        return eval_node(node.value, context)
    if isinstance(node, ast.Tuple):
        return tuple(eval_node(e, context) for e in node.elts)
    if isinstance(node, ast.List):
        return [eval_node(e, context) for e in node.elts]
    if isinstance(node, ast.Set):
        return {eval_node(e, context) for e in node.elts}
    if isinstance(node, ast.Dict):
        return dict(
            zip(
                [eval_node(k, context) for k in node.keys],
                [eval_node(v, context) for v in node.values],
            )
        )
    if isinstance(node, ast.Slice):
        return slice(
            eval_node(node.lower, context),
            eval_node(node.upper, context),
            eval_node(node.step, context),
        )
    if isinstance(node, ast.ExtSlice):
        return tuple([eval_node(dim, context) for dim in node.dims])
    if isinstance(node, ast.UnaryOp):
        # TODO: add guards
        value = eval_node(node.operand, context)
        if isinstance(node.op, ast.USub):
            return -value
        if isinstance(node.op, ast.UAdd):
            return +value
        if isinstance(node.op, ast.Invert):
            return ~value
        if isinstance(node.op, ast.Not):
            return not value
        raise ValueError("Unhandled unary operation:", node.op)
    if isinstance(node, ast.Subscript):
        value = eval_node(node.value, context)
        slice_ = eval_node(node.slice, context)
        if policy.can_get_item(value, slice_):
            return value[slice_]
        raise GuardRejection(
            "Subscript access (`__getitem__`) for",
            type(value),  # not joined to avoid calling `repr`
            f" not allowed in {context.evaluation} mode",
        )
    if isinstance(node, ast.Name):
        if policy.allow_locals_access and node.id in context.locals_:
            return context.locals_[node.id]
        if policy.allow_globals_access and node.id in context.globals_:
            return context.globals_[node.id]
        if policy.allow_builtins_access and node.id in __builtins__:
            return __builtins__[node.id]
        if not policy.allow_globals_access and not policy.allow_locals_access:
            raise GuardRejection(
                f"Namespace access not allowed in {context.evaluation} mode"
            )
        else:
            raise NameError(f"{node.id} not found in locals nor globals")
    if isinstance(node, ast.Attribute):
        value = eval_node(node.value, context)
        if policy.can_get_attr(value, node.attr):
            return getattr(value, node.attr)
        raise GuardRejection(
            "Attribute access (`__getattr__`) for",
            type(value),  # not joined to avoid calling `repr`
            f"not allowed in {context.evaluation} mode",
        )
    if isinstance(node, ast.IfExp):
        test = eval_node(node.test, context)
        if test:
            return eval_node(node.body, context)
        else:
            return eval_node(node.orelse, context)
    if isinstance(node, ast.Call):
        func = eval_node(node.func, context)
        print(node.keywords)
        if policy.can_call(func) and not node.keywords:
            args = [eval_node(arg, context) for arg in node.args]
            return func(*args)
        raise GuardRejection(
            "Call for",
            func,  # not joined to avoid calling `repr`
            f"not allowed in {context.evaluation} mode",
        )
    raise ValueError("Unhandled node", node)


SUPPORTED_EXTERNAL_GETITEM = {
    ("pandas", "core", "indexing", "_iLocIndexer"),
    ("pandas", "core", "indexing", "_LocIndexer"),
    ("pandas", "DataFrame"),
    ("pandas", "Series"),
    ("numpy", "ndarray"),
    ("numpy", "void"),
}

BUILTIN_GETITEM = {
    dict,
    str,
    bytes,
    list,
    tuple,
    collections.defaultdict,
    collections.deque,
    collections.OrderedDict,
    collections.ChainMap,
    collections.UserDict,
    collections.UserList,
    collections.UserString,
    DummyNamedTuple,
    IdentitySubscript,
}


def _list_methods(cls, source=None):
    """For use on immutable objects or with methods returning a copy"""
    return [getattr(cls, k) for k in (source if source else dir(cls))]


dict_non_mutating_methods = ("copy", "keys", "values", "items")
list_non_mutating_methods = ("copy", "index", "count")
set_non_mutating_methods = set(dir(set)) & set(dir(frozenset))


dict_keys = type({}.keys())
method_descriptor = type(list.copy)

ALLOWED_CALLS = {
    bytes,
    *_list_methods(bytes),
    dict,
    *_list_methods(dict, dict_non_mutating_methods),
    dict_keys.isdisjoint,
    list,
    *_list_methods(list, list_non_mutating_methods),
    set,
    *_list_methods(set, set_non_mutating_methods),
    frozenset,
    *_list_methods(frozenset),
    range,
    str,
    *_list_methods(str),
    tuple,
    *_list_methods(tuple),
    collections.deque,
    *_list_methods(collections.deque, list_non_mutating_methods),
    collections.defaultdict,
    *_list_methods(collections.defaultdict, dict_non_mutating_methods),
    collections.OrderedDict,
    *_list_methods(collections.OrderedDict, dict_non_mutating_methods),
    collections.UserDict,
    *_list_methods(collections.UserDict, dict_non_mutating_methods),
    collections.UserList,
    *_list_methods(collections.UserList, list_non_mutating_methods),
    collections.UserString,
    *_list_methods(collections.UserString, dir(str)),
    collections.Counter,
    *_list_methods(collections.Counter, dict_non_mutating_methods),
    collections.Counter.elements,
    collections.Counter.most_common,
}

EVALUATION_POLICIES = {
    "minimal": EvaluationPolicy(
        allow_builtins_access=True,
        allow_locals_access=False,
        allow_globals_access=False,
        allow_item_access=False,
        allow_attr_access=False,
        allowed_calls=set(),
        allow_any_calls=False,
    ),
    "limitted": SelectivePolicy(
        # TODO:
        # - should reject binary and unary operations if custom methods would be dispatched
        allowed_getitem=BUILTIN_GETITEM,
        allowed_getitem_external=SUPPORTED_EXTERNAL_GETITEM,
        allowed_getattr={
            *BUILTIN_GETITEM,
            set,
            frozenset,
            object,
            type,  # `type` handles a lot of generic cases, e.g. numbers as in `int.real`.
            dict_keys,
            method_descriptor,
        },
        allowed_getattr_external={
            # pandas Series/Frame implements custom `__getattr__`
            ("pandas", "DataFrame"),
            ("pandas", "Series"),
        },
        allow_builtins_access=True,
        allow_locals_access=True,
        allow_globals_access=True,
        allowed_calls=ALLOWED_CALLS,
    ),
    "unsafe": EvaluationPolicy(
        allow_builtins_access=True,
        allow_locals_access=True,
        allow_globals_access=True,
        allow_attr_access=True,
        allow_item_access=True,
        allow_any_calls=True,
    ),
}