# GT4Py - GridTools Framework
#
# Copyright (c) 2014-2023, ETH Zurich
# All rights reserved.
#
# This file is part of the GT4Py project and the GridTools framework.
# GT4Py is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the
# Free Software Foundation, either version 3 of the License, or any later
# version. See the LICENSE.txt file at the top-level directory of this
# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
#
# SPDX-License-Identifier: GPL-3.0-or-later

"""Utilities for runtime: introspection and bytecode."""

import ast
import copy
import inspect
import operator
import platform
import textwrap
from typing import Callable, Dict, List, Tuple, Type

from packaging import version

from .base import shashed_id


def get_closure(func, *, include_globals=True, included_nonlocals=True, include_builtins=True):
    closure_vars = inspect.getclosurevars(func)
    closure = {}
    unbound = set(closure_vars.unbound)
    if include_globals:
        closure.update(closure_vars.globals)
    else:
        unbound |= set(closure_vars.globals.keys())
    if included_nonlocals:
        closure.update(closure_vars.nonlocals)
    else:
        unbound |= set(closure_vars.nonlocals.keys())
    if include_builtins:
        closure.update(closure_vars.builtins)
    else:
        unbound |= set(closure_vars.builtins.keys())

    return closure, unbound


def get_source(func):
    try:
        source = inspect.getsource(func)
    except OSError as e:
        # Dynamically generated functions using exec() do not contain references
        # to the original source code, so we can use the convention of attaching
        # the source to the function using the "__exec_source__" attribute
        source = getattr(func, "__exec_source__", None)
        if not isinstance(source, str):
            raise e

    return source


def get_ast(func_or_source_or_ast, *, feature_version: Tuple[int, int]):
    if callable(func_or_source_or_ast):
        func_or_source_or_ast = get_source(func_or_source_or_ast)
    if isinstance(func_or_source_or_ast, str):
        func_or_source_or_ast = ast.parse(
            textwrap.dedent(func_or_source_or_ast), feature_version=feature_version
        )
    if isinstance(func_or_source_or_ast, (ast.AST, list)):
        ast_root = func_or_source_or_ast
    else:
        raise ValueError("Invalid function definition ({})".format(func_or_source_or_ast))
    return ast_root


def ast_dump(
    definition,
    *,
    skip_annotations: bool = True,
    skip_decorators: bool = True,
    feature_version: Tuple[int, int],
) -> str:
    def _dump(node: ast.AST, excluded_names):
        if isinstance(node, ast.AST):
            if excluded_names:
                fields = [
                    (name, _dump(value, excluded_names))
                    for name, value in sorted(ast.iter_fields(node))
                    if name not in excluded_names
                ]
            else:
                fields = [
                    (name, _dump(value, excluded_names))
                    for name, value in sorted(ast.iter_fields(node))
                ]

            return "".join(
                [
                    node.__class__.__name__,
                    "({content})".format(
                        content=", ".join("{}={}".format(name, value) for name, value in fields)
                    ),
                ]
            )

        elif isinstance(node, list):
            lines = ["[", *[_dump(i, excluded_names) + "," for i in node], "]"]
            return "\n".join(lines)

        else:
            return repr(node)

    skip_node_names = set()
    if skip_decorators:
        skip_node_names.add("decorator_list")
    if skip_annotations:
        skip_node_names.add("annotation")

    dumped_ast = _dump(get_ast(definition, feature_version=feature_version), skip_node_names)

    return dumped_ast


def ast_unparse(ast_node):
    """Call ast.unparse, but use astunparse for Python prior to 3.9."""
    if version.parse(platform.python_version()) < version.parse("3.9"):
        import astunparse

        return astunparse.unparse(ast_node)
    else:
        return ast.unparse(ast_node)


def ast_shash(ast_node, *, skip_decorators=True):
    return shashed_id(ast_dump(ast_node, skip_decorators=skip_decorators))


def collect_decorators(func_or_source_or_ast):
    if callable(func_or_source_or_ast):
        func_or_source_or_ast = get_source(func_or_source_or_ast)
    assert isinstance(func_or_source_or_ast, (str, ast.FunctionDef))

    if isinstance(func_or_source_or_ast, str):
        source = textwrap.dedent(func_or_source_or_ast)
        lines = source.splitlines()
        result = []
        i = 0
        while not lines[i].startswith("def "):
            if lines[i].startswith("@"):
                end = lines[i].find("(")
                if end < 0:
                    end = len(lines[i]) - 1
                decorator = lines[i][1:end]
                result.append(decorator)
            i += 1
    else:
        result = func_or_source_or_ast.decorator_list

    return result


def remove_decorators(func_or_source_or_ast):
    if callable(func_or_source_or_ast):
        func_or_source_or_ast = get_source(func_or_source_or_ast)
    assert isinstance(func_or_source_or_ast, (str, ast.FunctionDef))

    if isinstance(func_or_source_or_ast, str):
        source = textwrap.dedent(func_or_source_or_ast)
        lines = source.splitlines()
        start = 0
        for n, line in enumerate(lines):
            if line.startswith("def "):
                start = n
                break
        result = "\n".join(lines[start:])
    else:
        ast_root = copy.copy(func_or_source_or_ast)
        ast_root.decorator_list = []
        result = ast_root

    return result


def split_def_decorators(func_or_source):
    if callable(func_or_source):
        func_or_source = get_source(func_or_source)
    assert isinstance(func_or_source, str)

    source = textwrap.dedent(func_or_source)
    lines = source.splitlines()
    start = 0
    for n, line in enumerate(lines):
        if line.startswith("def "):
            start = n
            break
    definition = "\n".join(lines[start:])
    decorators = "\n".join(lines[:start])

    return definition, decorators


def get_qualified_name_from_node(name_or_attribute, *, as_list=False):
    assert isinstance(name_or_attribute, (ast.Attribute, ast.Name))

    node = name_or_attribute
    if isinstance(node, ast.Name):
        components = [node.id]
    else:
        assert isinstance(node, ast.Attribute)
        components = get_qualified_name_from_node(node.value, as_list=True)
        components = components + [node.attr]

    return components if as_list else ".".join(components)


class ASTPass:
    """Clone of the ast.NodeVisitor that supports forwarding kwargs."""

    def visit(self, node, **kwargs):
        """Visit a node."""
        method = "visit_" + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node, **kwargs)

    def generic_visit(self, node, **kwargs):
        """Fallback if no explicit visitor function exists for a node."""
        for _, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        self.visit(item, **kwargs)
            elif isinstance(value, ast.AST):
                self.visit(value, **kwargs)


class ASTTransformPass(ASTPass):
    """Clone of the ast.NodeTransformer that supports forwarding kwargs."""

    def generic_visit(self, node, **kwargs):
        for field, old_value in ast.iter_fields(node):
            if isinstance(old_value, list):
                new_values = []
                for value in old_value:
                    if isinstance(value, ast.AST):
                        value = self.visit(value, **kwargs)
                        if value is None:
                            continue
                        elif not isinstance(value, ast.AST):
                            new_values.extend(value)
                            continue
                    new_values.append(value)
                old_value[:] = new_values
            elif isinstance(old_value, ast.AST):
                new_node = self.visit(old_value, **kwargs)
                if new_node is None:
                    delattr(node, field)
                else:
                    setattr(node, field, new_node)
        return node


class ASTEvaluator(ASTPass):
    AST_OP_TO_OP: Dict[Type, Callable] = {
        # Arithmetic operations
        ast.UAdd: operator.pos,
        ast.USub: operator.neg,
        ast.Add: operator.add,
        ast.Sub: operator.sub,
        ast.Mult: operator.mul,
        ast.Div: operator.truediv,
        ast.Mod: operator.mod,
        ast.Pow: operator.pow,
        # Bitwise operations
        ast.Invert: operator.invert,
        ast.BitOr: operator.or_,
        ast.BitXor: operator.xor,
        ast.BitAnd: operator.and_,
        # Logical operations
        ast.Not: operator.not_,
        # ast.And: lambda a, b: a and b,
        # ast.Or: lambda a, b: a or b,
        ast.Eq: operator.eq,
        ast.NotEq: operator.ne,
        ast.Lt: operator.lt,
        ast.LtE: operator.le,
        ast.Gt: operator.gt,
        ast.GtE: operator.ge,
    }

    @classmethod
    def apply(cls, ast_root, context, default=None):
        assert isinstance(ast_root, ast.AST)
        try:
            result = cls(context).visit(ast_root)

        except (KeyError, ValueError):
            result = default
        return result

    def __init__(self, context: dict):
        self.context = copy.deepcopy(context)

    def visit_Name(self, node):
        return self.context[node.id]

    def visit_Num(self, node):
        return node.n

    def visit_Constant(self, node):
        return node.n

    def visit_NameConstant(self, node):
        return node.value

    def visit_Tuple(self, node: ast.Tuple):
        return tuple(self.visit(elem) for elem in node.elts)

    def visit_UnaryOp(self, node: ast.UnaryOp):
        val = self.visit(node.operand)
        return self.AST_OP_TO_OP[type(node.op)](val)

    def visit_BinOp(self, node: ast.BinOp):
        return self.AST_OP_TO_OP[type(node.op)](self.visit(node.left), self.visit(node.right))

    def visit_BoolOp(self, node: ast.BoolOp):
        # Use short-circuited evaluation of logical expressions
        condition = True if isinstance(node.op, ast.And) else False
        for value in node.values:
            if bool(self.visit(value)) is not condition:
                return not condition

        return condition

    def visit_Attribute(self, node: ast.Attribute):
        qualified_name = get_qualified_name_from_node(node)
        if qualified_name not in self.context:
            raise ValueError(f"{qualified_name} not found in context")
        return self.context[qualified_name]

    def visit_Compare(self, node: ast.Compare):
        values = [self.visit(node.left)] + [self.visit(cmp) for cmp in node.comparators]
        comparisons = [
            self.AST_OP_TO_OP[type(op)](values[i], values[i + 1]) for i, op in enumerate(node.ops)
        ]
        return all(comparisons)

    def generic_visit(self, node):
        raise ValueError("Invalid AST node for evaluation: {}".format(repr(node)))


ast_eval = ASTEvaluator.apply


class AssignTargetsCollector(ASTPass):
    @classmethod
    def apply(cls, ast_root, *, allow_multiple_targets=True):
        assert isinstance(ast_root, ast.AST)
        collector = cls(allow_multiple_targets=allow_multiple_targets)
        return collector(ast_root)

    def __init__(self, *, allow_multiple_targets):
        self.assign_targets = []
        self.allow_multiple_targets = allow_multiple_targets

    def __call__(self, ast_root):
        self.defs = []
        self.visit(ast_root)
        return self.assign_targets

    def visit_Assign(self, node: ast.Assign):
        if len(node.targets) > 1 and not self.allow_multiple_targets:
            raise RuntimeError(f"Multiple targets found in assignment ({node})")

        for target in node.targets:
            if isinstance(target, ast.Tuple):
                for t in target.elts:
                    assert isinstance(t, ast.Name)
                    self.assign_targets.append(t)
            else:
                self.assign_targets.append(target)


collect_assign_targets = AssignTargetsCollector.apply


class QualifiedNameCollector(ASTPass):
    @classmethod
    def apply(cls, ast_root, prefixes=None, *, skip_decorators=True, skip_annotations=True):
        assert isinstance(ast_root, ast.AST)
        collector = cls(
            prefixes=prefixes, skip_decorators=skip_decorators, skip_annotations=skip_annotations
        )
        return collector(ast_root)

    def __init__(self, prefixes=None, *, skip_decorators=True, skip_annotations=True):
        self.prefixes = set(prefixes) if prefixes else None
        self.skip_decorators = skip_decorators
        self.skip_annotations = skip_annotations
        self.name_nodes = None

    def __call__(self, ast_root):
        self.name_nodes = {}
        self.visit(ast_root)
        return self.name_nodes

    def generic_visit(self, node):
        if self.skip_decorators:
            node._fields = [name for name in node._fields if name != "decorator_list"]
        if self.skip_annotations:
            node._fields = [name for name in node._fields if name != "annotation"]
        super().generic_visit(node)

    def visit_Attribute(self, node: ast.Attribute):
        components, valid = self._get_name_components(node)
        if valid and not any(item is None for item in components):
            qualified_name = ".".join(components)
            self.name_nodes.setdefault(qualified_name, [])
            self.name_nodes[qualified_name].append(node)

    def visit_Name(self, node: ast.Name):
        if self.prefixes is None or node.id in self.prefixes:
            self.name_nodes.setdefault(node.id, [])
            self.name_nodes[node.id].append(node)

    def _get_name_components(self, node: ast.AST):
        components: List
        if isinstance(node, ast.Name):
            components = [node.id]
            valid = self.prefixes is None or node.id in self.prefixes
        elif isinstance(node, ast.Attribute):
            components, valid = self._get_name_components(node.value)
            components = components + [node.attr]
            valid = valid or ".".join(components) in self.prefixes
        else:
            components = [None]
            valid = False

        return components, valid


collect_names = QualifiedNameCollector.apply


class ImportsCollector(ASTPass):
    @classmethod
    def apply(cls, ast_root):
        assert isinstance(ast_root, ast.AST)
        collector = cls()
        return collector(ast_root)

    def __init__(self):
        self.bare_imports = None
        self.from_imports = None
        self.relative_imports = None

    def __call__(self, ast_root):
        self.bare_imports = {}
        self.from_imports = {}
        self.relative_imports = {}
        self.visit(ast_root)
        return self.bare_imports, self.from_imports, self.relative_imports

    def visit_Import(self, node: ast.Import):
        for alias in node.names:
            name = alias.name
            as_name = alias.asname if alias.asname else name
            self.bare_imports[name] = as_name

    def visit_ImportFrom(self, node: ast.ImportFrom):
        if node.module is None:
            imports_dict = self.relative_imports
            module = "".join(["."] * node.level)
        else:
            imports_dict = self.from_imports
            module = node.module

        for alias in node.names:
            name = ".".join([module, alias.name])
            as_name = alias.asname if alias.asname else name
            imports_dict[name] = as_name


collect_imported_symbols = ImportsCollector.apply


class SymbolsNameMapper(ASTTransformPass):
    @classmethod
    def apply(cls, ast_root, mapping, template_fmt=None, skip_names=None):
        assert isinstance(ast_root, ast.AST)
        collector = cls(mapping=mapping, template_fmt=template_fmt, skip_names=skip_names)
        return collector(ast_root)

    def __init__(self, mapping, template_fmt, skip_names):
        self.mapping = mapping
        self.template_fmt = template_fmt
        self.skip_names = set(skip_names)

    def __call__(self, ast_root):
        self.returns = []
        self.visit(ast_root)
        return self.returns

    def visit_Name(self, node: ast.Name):
        self.generic_visit(node)
        if not self.skip_names or node.id not in self.skip_names:
            if node.id in self.mapping:
                node.id = self.mapping[node.id]
            else:
                node.id = self.template_fmt.format(name=node.id)
        return node


map_symbol_names = SymbolsNameMapper.apply
