import ast
import sys
from ast import (
    Call,
    FunctionDef,
    Pass,
    UnaryOp,
    iter_child_nodes,
    literal_eval,
    NodeVisitor,
    If,
    Expr,
    Assign,
    Attribute,
    Name,
    Load,
    Store,
    Not,
    arg,
    Raise,
    Return,
)

try:
    from ast import Constant

    def Str(s):
        return Constant(value=s)

    def is_str_node(node):
        return isinstance(node, Constant) and isinstance(node.value, str)

    def catenate_str_node(node1, node2):
        node1.value += node2.value

except ImportError:
    from ast import Str

    def is_str_node(node):
        return isinstance(node, Str)

    def catenate_str_node(node1, node2):
        node1.s += node2.s


from collections.abc import Iterable

from .astalyzer import FreeVarFinder
from ..base import LanguageNode, ComplexNode, BaseGenerator
from ...helpers import StringWithLocation
from ...runtime.debug import TemplateSyntaxError

try:  # pragma: no cover
    import sysconfig

    HAS_ASSERT = bool(sysconfig.get_config_var("Py_DEBUG"))
    del sysconfig
except ImportError:  # pragma: no cover
    HAS_ASSERT = False

name_counter = 0
ALWAYS_BUILTINS = """
    False
    True
    None
""".split()


def simple_call(func, args=None):
    # Remove deprecated starargs and kwargs parameters in Python 3.13+
    call_kwargs = {
        "func": func,
        "args": args or [],
        "keywords": [],
    }

    # These parameters were deprecated in Python 3.13+
    if sys.version_info < (3, 13):
        call_kwargs.update(
            {
                "starargs": None,
                "kwargs": None,
            }
        )

    return Call(**call_kwargs)


def create_argument_list(arguments):
    return [arg(arg=id, annotation=None) for id in arguments]


def simple_function_def(name, arguments=()):
    arguments = create_argument_list(arguments)
    if sys.version_info >= (3, 8):
        extra = {"posonlyargs": []}
    else:
        extra = {}

    # Handle deprecated annotation parameters in Python 3.13+
    args_kwargs = {
        "args": arguments,
        "vararg": None,
        "kwonlyargs": [],
        "kwarg": None,
        "defaults": [],
        "kw_defaults": [],
        **extra,
    }

    # These annotation parameters were deprecated in Python 3.13+
    if sys.version_info < (3, 13):
        args_kwargs.update(
            {
                "varargannotation": None,
                "kwargannotation": None,
            }
        )

    return FunctionDef(
        name=name,
        args=ast.arguments(**args_kwargs),
        body=[Pass()],
        decorator_list=[],
        returns=None,
    )


def NameX(id, store=False):
    return Name(id=id, ctx=Load() if not store else Store())


def adjust_locations(ast_node, first_lineno, first_offset):
    """
    Adjust the locations of the ast nodes, offsetting them
    to the new lineno and column offset
    """

    line_delta = first_lineno - 1

    def _fix(node):
        for prefix in ["", "end_"]:
            if prefix + "lineno" in node._attributes:
                lineno_attr = prefix + "lineno"
                col_offset_attr = prefix + "col_offset"

                lineno = getattr(node, lineno_attr)
                col = getattr(node, col_offset_attr)

                # adjust the offset on the first line
                if lineno == 1:
                    col += first_offset

                lineno += line_delta

                setattr(node, lineno_attr, lineno)
                setattr(node, col_offset_attr, col)

        for child in iter_child_nodes(node):
            _fix(child)

    _fix(ast_node)


def get_fragment_ast(expression, mode="eval", adjust=(0, 0)):
    if not isinstance(expression, str):
        return expression

    position = getattr(expression, "position", (1, 0))
    position = position[0] + adjust[0], position[1] + adjust[1]

    tree = None
    t = None
    try:
        exp = expression
        if expression[-1:] != "\n":
            exp = expression + "\n"
        tree = ast.parse(exp, mode=mode)
    except SyntaxError as e:
        lineno = e.lineno
        lineno += position[0] - 1
        t = TemplateSyntaxError(e.msg, lineno=lineno)

    if tree is None:
        raise t

    adjust_locations(tree, position[0], position[1])
    return tree.body


def gen_name(typename=None):
    global name_counter
    name_counter += 1
    if typename:
        return "__TK__typed__%s__%d__" % (typename, name_counter)
    else:
        return "__TK_%d__" % (name_counter)


def static_eval(expr):
    if isinstance(expr, UnaryOp) and isinstance(expr.op, Not):
        # bogus error from PyCharm
        return not static_eval(expr.operand)  # NoQA

    return literal_eval(expr)


def static_expr_to_bool(expr):
    try:
        return bool(static_eval(expr))
    except BaseException:
        return None


class PythonNode(LanguageNode):
    is_top_level = False

    def generate_output_ast(self, code, generator, parent, escape=False, position=None):
        func = Name(id="__TK__output", ctx=Load())

        if not isinstance(code, list):
            code = [code]

        rv = []
        for i in code:
            if position is not None:
                i.lineno, i.col_offset = position

            e = Expr(simple_call(func, [i]))
            e.output_args = [i]
            rv.append(e)

        return rv

    def make_buffer_frame(self, body):
        new_body = []
        new_body.append(
            Assign(
                targets=[
                    NameX("__TK__output", store=True),
                ],
                value=simple_call(NameX("__TK__mkbuffer")),
            )
        )

        new_body.extend(body)
        new_body.append(Return(value=NameX("__TK__output")))
        return new_body

    def make_function(self, name, body, add_buffer=False, arguments=()):
        # ensure that the function name is an str
        func = simple_function_def(str(name), arguments=arguments)
        new_body = func.body = []

        if add_buffer:
            new_body.extend(self.make_buffer_frame(body))

        else:
            new_body.extend(body)

        if not new_body:
            new_body.append(Pass())

        return func

    def generate_varscope(self, body):
        name = gen_name("variable_scope")
        rv = [
            self.make_function(name, body, arguments=["__TK__output", "__TK__escape"]),
            Expr(
                simple_call(NameX(name), [NameX("__TK__output"), NameX("__TK__escape")])
            ),
        ]
        return rv


class PyOutputNode(PythonNode):
    def __init__(self, text):
        super(PyOutputNode, self).__init__()
        self.text = text

    def get_expressions(self):
        return [self.get_expression()]

    def get_expression(self):
        return Str(s=self.text)

    def generate_ast(self, generator, parent):
        return self.generate_output_ast(self.get_expression(), generator, parent)


class PyTranslatableOutputNode(PyOutputNode):
    def __init__(self, text, needs_escape=False):
        super(PyTranslatableOutputNode, self).__init__(text)
        self.needs_escape = needs_escape

    def get_expressions(self):
        return [self.get_expression()]

    def get_expression(self):
        name = "gettext"
        if self.needs_escape:
            name = "egettext"

        expr = simple_call(
            NameX(name),
            [Str(s=self.text)],
        )
        return expr


class PyExpressionNode(PythonNode):
    def __init__(self, expression):
        super(PyExpressionNode, self).__init__()
        self.expr = expression

    def get_expressions(self):
        return [self.get_expression()]

    def get_expression(self):
        return simple_call(NameX("__TK__escape"), [self.get_unescaped_expression()])

    def get_unescaped_expression(self):
        return get_fragment_ast(self.expr)

    def generate_ast(self, generator, parent):
        return self.generate_output_ast([self.get_expression()], generator, parent)


class PyCodeNode(PythonNode):
    def __init__(self, source):
        super(PyCodeNode, self).__init__()
        self.source = source

    def generate_ast(self, generator, parent):
        return get_fragment_ast(self.source, mode="exec")


def coalesce_strings(args):
    rv = []
    str_on = None

    for i in args:
        if is_str_node(i):
            if str_on:
                catenate_str_node(str_on, i)
                continue

            str_on = i

        else:
            str_on = None

        rv.append(i)

    return rv


class PyComplexNode(ComplexNode, PythonNode):
    def generate_child_ast(self, generator, parent_for_children):
        rv = []
        for i in self.children:
            rv.extend(i.generate_ast(generator, parent_for_children))

        return rv


class PyIfNode(PyComplexNode):
    def __init__(self, expression):
        super(PyIfNode, self).__init__()
        self.expression = expression

    def generate_ast(self, generator, parent):
        test = get_fragment_ast(self.expression)
        boolean = static_expr_to_bool(test)

        if boolean is False:
            return []

        if boolean is True:
            return self.generate_child_ast(generator, parent)

        node = If(test=test, body=self.generate_child_ast(generator, self), orelse=[])
        return [node]


def PyUnlessNode(self, expression):
    expression = get_fragment_ast(expression)
    expression = UnaryOp(op=Not(), operand=expression)
    return PyIfNode(expression)


class PyImportNode(PythonNode):
    def __init__(self, href, alias):
        super(PyImportNode, self).__init__()
        self.href = str(href)
        self.alias = str(alias)

    def generate_ast(self, generator, parent):
        node = Assign(
            targets=[NameX(str(self.alias), store=True)],
            value=simple_call(
                func=Attribute(
                    value=NameX("__TK__runtime", store=False),
                    attr="import_defs",
                    ctx=Load(),
                ),
                args=[NameX("__TK__original_context"), Str(s=self.href)],
            ),
        )

        if parent.is_top_level:
            generator.add_top_level_import(str(self.alias), node)
            return []

        return [node]


class PyAttributeNode(PyComplexNode):
    def __init__(self, name, value):
        super(PyAttributeNode, self).__init__()
        self.name = name

    def get_expressions(self):
        rv = []
        for i in self.children:
            rv.extend(i.get_expressions())

        return rv

    def generate_ast(self, generator, parent):
        if len(self.children) == 1 and isinstance(self.children[0], PyExpressionNode):
            # special case, the attribute contains a single
            # expression, these are handled by
            # __TK__output.output_boolean_attr,
            # given the name, and unescaped expression!
            return [
                Expr(
                    simple_call(
                        func=Attribute(
                            value=NameX("__TK__output"),
                            attr="output_boolean_attr",
                            ctx=Load(),
                        ),
                        args=[
                            Str(s=self.name),
                            self.children[0].get_unescaped_expression(),
                        ],
                    )
                )
            ]

        # otherwise just return the output for the attribute code
        # like before
        return self.generate_output_ast(
            [Str(s=' %s="' % self.name)] + self.get_expressions() + [Str(s='"')],
            generator,
            parent,
        )


class PyAttrsNode(PythonNode):
    def __init__(self, expression):
        super(PyAttrsNode, self).__init__()
        self.expression = expression

    def generate_ast(self, generator, parent):
        expression = get_fragment_ast(self.expression)

        output = simple_call(NameX("__TK__output_attrs"), args=[expression])
        return self.generate_output_ast([output], generator, parent)


class PyForNode(PyComplexNode):
    def __init__(self, target_and_expression, parts):
        super(PyForNode, self).__init__()
        self.target_and_expression = target_and_expression

    def generate_contents(self, generator, parent):
        lineno, col = getattr(self.target_and_expression, "position", (1, 0))

        body = get_fragment_ast(
            StringWithLocation(
                "for %s: pass" % self.target_and_expression, lineno, col - 4
            ),
            "exec",
        )
        for_node = body[0]
        for_node.body = self.generate_child_ast(generator, self)
        return [for_node]

    def generate_ast(self, generator, parent):
        # TODO: this could be needed to be reinstantiated
        # return self.generate_varscope(self.generate_contents())
        return self.generate_contents(generator, parent)


class PyDefineNode(PyComplexNode):
    def __init__(self, funcspec):
        super(PyDefineNode, self).__init__()

        self.position = getattr(funcspec, "position", (1, 0))

        if "(" not in funcspec:
            funcspec += "()"

        self.funcspec = funcspec

    def generate_ast(self, generator, parent):
        body = get_fragment_ast(
            StringWithLocation(
                "def %s: pass" % self.funcspec, self.position[0], self.position[1] - 4
            ),
            "exec",
        )
        def_node = body[0]
        def_node.body = self.make_buffer_frame(self.generate_child_ast(generator, self))

        # move the function out of the closure
        if parent.is_top_level:
            generator.add_top_def(def_node.name, def_node)
            return []

        return [def_node]


class PyComplexExprNode(PyComplexNode):
    def get_expressions(self):
        rv = []
        for i in self.children:
            if hasattr(i, "get_expression"):
                rv.append(i.get_expression())

            else:
                rv.extend(i.get_expressions())

        return rv

    def generate_ast(self, generator, parent=None):
        return self.generate_output_ast(self.get_expressions(), generator, parent)


class PyBlockNode(PyComplexNode):
    def __init__(self, name):
        super(PyBlockNode, self).__init__()
        self.name = name

    def generate_ast(self, generator, parent):
        is_extended = isinstance(parent, PyExtendsNode)

        name = self.name
        blockfunc_name = f"__TK__block__{name}"
        position = getattr(name, "position", (1, 0))
        body = get_fragment_ast(
            StringWithLocation(
                "def %s():pass" % blockfunc_name, position[0], position[1] - 4
            ),
            "exec",
        )
        def_node = body[0]
        def_node.body = self.make_buffer_frame(self.generate_child_ast(generator, self))

        generator.add_block(str(name), def_node, blockfunc_name)

        if not is_extended:
            # call the block in place
            return self.generate_output_ast(
                [simple_call(NameX(str(self.name)), [])],
                self,
                parent,
                position=position,
            )

        else:
            return []


class PyWithNode(PyComplexNode):
    def __init__(self, vars):
        super(PyWithNode, self).__init__()
        self.vars = vars

    def generate_ast(self, generator, parent=None):
        var_defs = get_fragment_ast(self.vars, "exec")
        body = var_defs + self.generate_child_ast(generator, self)
        return self.generate_varscope(body)


class PyExtendsNode(PyComplexNode):
    is_top_level = True

    def __init__(self, href):
        super(PyExtendsNode, self).__init__()
        self.href = href

    def generate_ast(self, generator, parent=None):
        generator.make_extended_template(self.href)
        return self.generate_child_ast(generator, self)


def ast_equals(tree1, tree2):
    x1 = ast.dump(tree1)
    x2 = ast.dump(tree2)
    return x1 == x2


def coalesce_outputs(tree):
    """
    Coalesce the constant output expressions

        __output__('foo')
        __output__('bar')
        __output__(baz)
        __output__('xyzzy')

    into

        __output__('foobar', baz, 'xyzzy')
    """

    coalesce_all_outputs = True
    if coalesce_all_outputs:

        def should_coalesce(node):
            return True
    else:

        def should_coalesce(node):
            return node.output_args[0].__class__ is Str

    class OutputCoalescer(NodeVisitor):
        def visit(self, node):
            # if - else expression also has a body! it is not we want, though.
            if hasattr(node, "body") and isinstance(node.body, Iterable):
                # coalesce continuous string output nodes
                new_body = []
                output_node = None

                def coalesce_strs():
                    if output_node:
                        output_node.value.args[:] = coalesce_strings(
                            output_node.value.args
                        )

                for i in node.body:
                    if hasattr(i, "output_args") and should_coalesce(i):
                        if output_node:
                            if len(output_node.value.args) + len(i.output_args) > 250:
                                coalesce_strs()
                                output_node = i
                            else:
                                output_node.value.args.extend(i.output_args)
                                continue

                        output_node = i

                    else:
                        coalesce_strs()
                        output_node = None

                    new_body.append(i)

                coalesce_strs()
                node.body[:] = new_body

            NodeVisitor.visit(self, node)

        def check(self, node):
            """
            Coalesce __TK__output(__TK__escape(literal(x))) into
            __TK__output(x).
            """
            if not ast_equals(node.func, NameX("__TK__output")):
                return

            for i in range(len(node.args)):
                arg1 = node.args[i]
                if not arg1.__class__.__name__ == "Call":
                    continue

                if not ast_equals(arg1.func, NameX("__TK__escape")):
                    continue

                if len(arg1.args) != 1:
                    continue

                arg2 = arg1.args[0]
                if not arg2.__class__.__name__ == "Call":
                    continue

                if not ast_equals(arg2.func, NameX("literal")):
                    continue

                if len(arg2.args) != 1:
                    continue

                node.args[i] = arg2.args[0]

        def visit_Call(self, node):
            self.check(node)
            self.generic_visit(node)

    OutputCoalescer().visit(tree)


def remove_locations(node):
    """
    Removes locations from the given AST tree completely
    """

    def fix(node):
        if "lineno" in node._attributes and hasattr(node, "lineno"):
            del node.lineno

        if "col_offset" in node._attributes and hasattr(node, "col_offset"):
            del node.col_offset

        for child in iter_child_nodes(node):
            fix(child)

    fix(node)


class PyRootNode(PyComplexNode):
    def __init__(self):
        super(PyRootNode, self).__init__()

    is_top_level = True

    def generate_ast(self, generator, parent=None):
        main_body = self.generate_child_ast(generator, self)

        extended = generator.extended_href

        toplevel_funcs = generator.blocks + generator.top_defs
        # do not generate __main__ for extended templates
        if not extended:
            main_func = self.make_function("__main__", main_body, add_buffer=True)
            generator.add_bind_decorator(main_func)

            # In Python 3.13+, functions must be defined before they're referenced in closures
            # Put block and top-level functions before the main function
            toplevel_funcs = toplevel_funcs + [main_func]

        # analyze the set of free variables
        free_variables = set()
        for i in toplevel_funcs:
            fv_info = FreeVarFinder.for_ast(i)
            free_variables.update(fv_info.get_free_variables())

        # discard __TK__ variables, always builtin names True, False, None
        # from free variables.
        for i in list(free_variables):
            if i.startswith("__TK__") or i in ALWAYS_BUILTINS:
                free_variables.discard(i)

        # discard the names of toplevel funcs from free variables
        free_variables.difference_update(generator.top_level_names)

        code = "__TK__mkbuffer = __TK__runtime.Buffer\n"
        code += "__TK__escape = __TK__escape_g = __TK__runtime.escape\n"
        code += "__TK__output_attrs = __TK__runtime.output_attrs\n"

        if extended:
            code += "__TK__parent_template = __TK__runtime.load(%r)\n" % extended

        code += "def __TK__binder(__TK__context):\n"
        code += "    __TK__original_context = __TK__context.copy()\n"
        code += "    __TK__bind = __TK__runtime.bind(__TK__context)\n"
        code += (
            "    __TK__bindblock = __TK__runtime.bind(__TK__context, " "block=True)\n"
        )

        # bind gettext early!
        for i in ["egettext"]:
            if i in free_variables:
                free_variables.add("gettext")
                free_variables.discard(i)

        if "gettext" in free_variables:
            code += "    def egettext(msg):\n"
            code += "        return __TK__escape(gettext(msg))\n"
            code += '    gettext = __TK__context["gettext"]\n'
            free_variables.discard("gettext")

        code += "    raise\n"  # a placeholder

        if extended:
            # an extended template does not have a __main__ (it is inherited)
            code += "    __TK__parent_template.binder_func(__TK__context)\n"

        # Bind block functions to context in non-extended templates
        if not extended:
            for block_func in generator.blocks:
                block_name = block_func.name
                if block_name.startswith("__TK__block__"):
                    # Extract the user-facing block name
                    user_block_name = block_name[len("__TK__block__") :]
                    code += f'    __TK__context["{user_block_name}"] = {block_name}\n'

        for i in free_variables:
            code += '    if "%s" in __TK__context:\n' % i
            code += '        %s = __TK__context["%s"]\n' % (i, i)

        code += "    return __TK__context\n"

        tree = ast.parse(code)
        remove_locations(tree)

        class LocatorAndTransformer(ast.NodeTransformer):
            binder = None

            def visit_FunctionDef(self, node):
                if node.name == "__TK__binder" and not self.binder:
                    self.binder = node

                self.generic_visit(node)
                return node

        locator = LocatorAndTransformer()
        locator.visit(tree)

        # inject the other top level funcs in the binder
        binder = locator.binder
        for i, e in enumerate(binder.body):
            if isinstance(e, Raise):
                break

        binder.body[i : i + 1] = toplevel_funcs
        binder.body[i:i] = generator.imports

        coalesce_outputs(tree)
        return tree


# noinspection PyProtectedMember
class LocationMapper(object):
    def __init__(self):
        self.lineno_map = {1: 1}
        self.prev_original_line = 1
        self.prev_mapped_line = 1
        self.prev_column = 0

    def map_linenos(self, node):
        if "lineno" in node._attributes:
            if hasattr(node, "lineno"):
                if node.lineno != self.prev_original_line:
                    self.prev_mapped_line += 1
                    self.lineno_map[self.prev_mapped_line] = node.lineno
                    self.prev_original_line = node.lineno

            node.lineno = self.prev_mapped_line

        if "col_offset" in node._attributes:
            if hasattr(node, "col_offset"):
                self.prev_column = node.col_offset

            node.col_offset = self.prev_column

        for child in iter_child_nodes(node):
            self.map_linenos(child)


class Generator(BaseGenerator):
    OutputNode = PyOutputNode
    TranslatableOutputNode = PyTranslatableOutputNode

    IfNode = PyIfNode
    ForNode = PyForNode
    DefineNode = PyDefineNode
    ComplexExprNode = PyComplexExprNode
    ExpressionNode = PyExpressionNode
    ImportNode = PyImportNode
    RootNode = PyRootNode
    AttributeNode = PyAttributeNode
    AttrsNode = PyAttrsNode
    UnlessNode = PyUnlessNode
    ExtendsNode = PyExtendsNode
    BlockNode = PyBlockNode
    CodeNode = PyCodeNode
    WithNode = PyWithNode

    def __init__(self, ir_tree):
        super(Generator, self).__init__(ir_tree)
        self.blocks = []
        self.top_defs = []
        self.top_level_names = set()
        self.extended_href = None
        self.imports = []
        self.lnotab = None

    def add_bind_decorator(self, func, block=True):
        binder_call = NameX("__TK__bind" + ("block" if block else ""))
        decors = [binder_call]
        func.decorator_list = decors

    def add_block(self, name, blockfunc, blockfunc_name):
        self.top_level_names.add(blockfunc_name)
        self.add_bind_decorator(blockfunc, block=True)
        self.blocks.append(blockfunc)

    def add_top_def(self, name, defblock):
        self.top_level_names.add(name)
        self.add_bind_decorator(defblock)
        self.top_defs.append(defblock)

    def add_top_level_import(self, name, node):
        self.top_level_names.add(name)
        self.imports.append(node)

    def make_extended_template(self, href):
        self.extended_href = href

    def lnotab_info(self):
        return self.lnotab

    def generate_ast(self):
        tree = super(Generator, self).generate_ast()
        self.lnotab = {}
        ast.fix_missing_locations(tree)
        return tree
