"""
This module contains various static transformations and analyses of Python ASTs
that are used for executing Python functions in continuation-passing style (CPS).
"""

import ast
from functools import cached_property
from typing import Callable, TYPE_CHECKING
import textwrap
import copy
import contextlib
import collections
import dataclasses

from flippy.hashable import hashabledict
from flippy.types import CPSCallable

if TYPE_CHECKING:
    from flippy.interpreter import CPSInterpreter

class _ValueWrapper:
    '''
    A simple wrapper class that we use to indicate when a scope is being returned below.
    '''
    def __init__(self, value):
        self.value = value

class Placeholder(ast.NodeTransformer):
    '''
    Use this class to add placeholders in code constructed by string, so that
    placeholders can later be filled by AST instances.
    '''
    prefix = '__PLACEHOLDER_'
    def __init__(self, replacement_map, *, skip_missing=False):
        self.replacement_map = replacement_map
        self.skip_missing = skip_missing

    @classmethod
    def new(cls, tag):
        return f'{cls.prefix}{tag}'
    @classmethod
    def fill(cls, node, replacement_map, **kwargs):
        return cls(replacement_map, **kwargs).visit(node)
    @classmethod
    def fill_from_source(cls, src, *args, **kwargs):
        node = ast.parse(textwrap.dedent(src))
        return cls.fill(node, *args, **kwargs)

    def visit_Name(self, node):
        if node.id.startswith(__class__.prefix):
            tag = node.id[len(__class__.prefix):]
            if tag in self.replacement_map:
                return self.replacement_map[tag]
            elif not self.skip_missing:
                assert tag in self.replacement_map, f'Could not find tag {tag} in replacement map.'
        return self.generic_visit(node)
    def visit_Expr(self, node):
        new_node = self.generic_visit(node)
        if hasattr(new_node, 'value'):
            if isinstance(new_node.value, ast.expr):
                return new_node
            else:
                # This case is when replacement of var was a statement
                return new_node.value
        else:
            # When replacement of var was None, this node should be deleted.
            return None


@dataclasses.dataclass
class Scope:
    # The node this scope is associated with.
    node: ast.AST
    # The parent or outer scope
    parent: 'Scope'
    # This namespace (ns) tracks variable names and the number of times they were written to.
    ns: dict[str, int] = dataclasses.field(default_factory=collections.Counter)

    def resolve(self, name):
        scope = self
        while scope is not None:
            if name in scope.ns:
                return scope
            scope = scope.parent

    def clone(self):
        return __class__(self.node, self.parent, self.ns.copy())

    def merge_with_branched_scope(self, other):
        '''
        This is used to combine the scopes generated by analyzing each branch
        in an if/else statement. Scopes are combined by taking the union of all
        keys, and returning the maximum number of writes across the branches.
        This ensures that we track all variables that could have been set, and
        are tracking the maximum number of times they were set.
        '''
        assert self.node == other.node and self.parent == other.parent
        return __class__(self.node, self.parent, collections.Counter({
            # Return the maximum number of writes
            key: max(self.ns[key], other.ns[key])
            # Take the union of keys
            for key in self.ns.keys() | other.ns.keys()
        }))

class ClosureScopeAnalysis(ast.NodeVisitor):
    '''
    This analysis enforces some constraints on scopes that are necessary for
    our CPS transform.

    1. Only the local scope can be mutated. Enforced by prohibiting nonlocal and global keywords.
    2. If a nonlocal variable is loaded, we require that it be immutable in its original scope.

    We implement this by doing two passes over the input AST. The first establishes the scope
    associated with each variable reference. The second validates that variables follow rule #2 above.
    Two passes are necessary because names are assigned to scopes statically based on whether it is
    ever possible for them to be defined in that scope.
    '''

    def_pass = 'def_pass'
    use_pass = 'use_pass'

    def __call__(self, node, source):
        self.node = node
        self.source = source
        self.scopes: list[Scope] = []
        self.scope_map = {}
        '''
        We analyze code in two passes. The first pass tracks all stores to
        variables in order to establish scopes. We hold onto those fully
        analyzed scopes in `complete_scope_map`.
        '''
        self.analysis_pass = __class__.def_pass
        with self.in_scope(node):
            self.visit(node)
        self.complete_scope_map = dict(self.scope_map)
        '''
        The main difference in this second pass is that we validate loaded
        variables against statically-determined scopes to ensure closure
        vars are immutable. We still keep the definition machinery to validate
        that variables are defined before closures, in order to avoid issues
        with CPS.
        '''
        self.analysis_pass = __class__.use_pass
        with self.in_scope(node):
            self.visit(node)

    def analyze_scope(self, node):
        '''
        This returns a mapping from scope-defining nodes to a `Scope`.
        '''
        # HACK: Copying above invocation.
        self.scopes: list[Scope] = []
        self.scope_map = {}
        self.analysis_pass = __class__.def_pass
        with self.in_scope(node):
            self.visit(node)
        return self.scope_map

    def generic_visit(self, node):
        # Refactored from ASTVisitor
        for field, value in ast.iter_fields(node):
            self.generic_visit_field(value)

    def generic_visit_field(self, value):
        # Refactored from ASTVisitor
        if isinstance(value, list):
            for item in value:
                if isinstance(item, ast.AST):
                    self.visit(item)
        elif isinstance(value, ast.AST):
            self.visit(value)

    def visit_Lambda(self, node):
        with self.in_scope(node):
            return self.generic_visit(node)

    def visit_FunctionDef(self, node):
        self.add_name_to_scope(node.name)
        self.generic_visit_field(node.decorator_list)
        with self.in_scope(node):
            self.generic_visit_field(node.args)
            self.generic_visit_field(node.body)

    @contextlib.contextmanager
    def in_scope(self, node):
        self.scopes.append(Scope(
            node=node,
            parent=self.current_scope if len(self.scopes) else None,
        ))
        self.scope_map[node] = self.current_scope
        try:
            yield
        finally:
            self.scopes.pop()

    def add_name_to_scope(self, name):
        self.current_scope.ns[name] += 1

    @property
    def current_scope(self):
        return self.scopes[-1]

    def visit_Assign(self, node):
        # We make sure to visit values first.
        self.generic_visit_field(node.value)
        # Then we visit targets.
        self.generic_visit_field(node.targets)

    def visit_AnnAssign(self, node):
        # Value first then target, like above.
        self.generic_visit_field(node.value)
        self.generic_visit_field(node.target)

    def visit_AugAssign(self, node):
        # Value first then target, like above.
        # NOTE: We avoid separately analying a Load() to the target because
        # the case for Load() in visit_Name checks for nonlocals that were mutated
        # but we know this variable is local (because the nonlocal scope can't be changed).
        # So, there's no way a Load() here could detect an issue.
        self.generic_visit_field(node.value)
        self.generic_visit_field(node.target)

    def visit_arg(self, node):
        self.add_name_to_scope(node.arg)

    # Handling comprehensions
    def visit_ListComp(self, node):
        with self.in_scope(node):
            self.generic_visit_field(node.generators)
            self.generic_visit_field(node.elt)
    visit_SetComp = visit_ListComp
    def visit_DictComp(self, node):
        with self.in_scope(node):
            self.generic_visit_field(node.generators)
            self.generic_visit_field(node.key)
            self.generic_visit_field(node.value)
    def visit_comprehension(self, node):
        self.generic_visit_field(node.iter)
        self.generic_visit_field(node.target)
        self.generic_visit_field(node.ifs)

    def visit_If(self, node):
        def _replace_current_scope(scope):
            prev_scope = self.current_scope
            # We want to make sure we're replacing a comparable scope, since we
            # only use this function to change the namespace, but still want
            # to ensure they are associated with the same node and parent scope.
            assert (
                prev_scope.node == scope.node and
                prev_scope.parent == scope.parent and
                prev_scope == self.scope_map[scope.node]
            ), 'Ensures that new scope can replace current scope.'
            self.scopes[-1] = scope
            self.scope_map[scope.node] = scope
            return prev_scope

        # 1) Analyze the test.
        self.generic_visit_field(node.test)

        # 2) Analyze the body.
        # Copy starting scope
        starting = self.current_scope.clone()
        # Using original starting scope, analyze body
        self.generic_visit_field(node.body)

        # 3) Analyze the orelse.
        # Replace scope with starting scope and analyze orelse
        body_scope = _replace_current_scope(starting)
        self.generic_visit_field(node.orelse)

        # 4) Combine the scopes from each branch.
        _replace_current_scope(self.current_scope.merge_with_branched_scope(body_scope))

    # Case for handling names
    def visit_Name(self, node):
        if isinstance(node.ctx, ast.Store):
            self.add_name_to_scope(node.id)

        elif self.analysis_pass == __class__.use_pass:
            if isinstance(node.ctx, ast.Load):
                local_scope = self.complete_scope_map[self.current_scope.node]
                scope = local_scope.resolve(node.id)
                is_global = scope is None
                is_local = scope == local_scope
                if not is_global and not is_local:
                    write_count = scope.ns[node.id]
                    if write_count > 1:
                        raise _error_at_node(f'Variable from outer scope must be immutable, but is set in {write_count} locations.', node, self.source)
                    if node.id not in self.scope_map[scope.node].ns:
                        raise _error_at_node('Variable from outer scope must be defined before being referenced.', node, self.source)

    # We block unhandled cases that change scope or bind names.
    # https://realpython.com/python-scope-legb-rule/

    def not_implemented(self, node):
        raise NotImplementedError()

    # Blocked because they make outer scopes mutable.
    visit_Global = not_implemented
    visit_Nonlocal = not_implemented

    # Blocked because they bind names.
    visit_Try = not_implemented
    visit_TryStar = not_implemented
    visit_NamedExpr = not_implemented

    # Blocked because they cause new scopes.
    visit_ClassDef = not_implemented
    visit_GeneratorExp = not_implemented

class NodeTransformer(ast.NodeTransformer):
    def generic_visit(self, node):
        for field, _ in ast.iter_fields(node):
            self.generic_visit_field(node, field)
        return node

    def generic_visit_field(self, node, field):
        # This is the inner loop over fields implemented by `ast.NodeTransformer.generic_visit`.
        # It's useful when a specific order over fields needs to be guaranteed, so we refactor
        # it out here.
        old_value = getattr(node, field)
        if isinstance(old_value, list):
            # This case and next one are both copied from original.
            new_values = []
            for value in old_value:
                if isinstance(value, ast.AST):
                    value = self.visit(value)
                    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)
            if new_node is None:
                delattr(node, field)
            else:
                setattr(node, field, new_node)

def _error_at_node(msg, node, source, *, filename='tmp.py'):
    # Subtract 1 b/c node.lineno is 1-indexed
    code = source.splitlines()[node.lineno - 1]
    # Add 1 b/c node.col_offset is 0-indexed
    offset = node.col_offset + 1
    return SyntaxError(msg, (filename, node.lineno, offset, code))

class PythonSubsetValidator(ast.NodeVisitor):
    def __call__(self, node, source):
        self.in_loop_flag = False
        self.errors = []
        self.visit(node)
        if self.errors:
            # HACK For now, we just report the first error.
            node = self.errors[0]
            raise _error_at_node('Found unsupported Python feature.', node, source)

    def add_error(self, node):
        self.errors.append(node)

    def visit_error(self, node):
        self.add_error(node)
        super().generic_visit(node)

    visit_Global = visit_error
    visit_Nonlocal = visit_error

    visit_ClassDef = visit_error
    visit_AsyncFunctionDef = visit_error

    visit_AsyncFor = visit_error
    visit_AsyncWith = visit_error
    visit_With = visit_error

    visit_Match = visit_error

    visit_Try = visit_error
    visit_TryStar = visit_error

    visit_Import = visit_error
    visit_ImportFrom = visit_error

    visit_Await = visit_error
    visit_Yield = visit_error
    visit_YieldFrom = visit_error
    visit_NamedExpr = visit_error

    # We prohibit del/set of attributes and subscripts
    def visit_Attribute(self, node):
        super().generic_visit(node)
        if isinstance(node.ctx, (ast.Store, ast.Del)):
            self.add_error(node)
    visit_Subscript = visit_Attribute

    # We prohibit nested functions in loops
    def visit_While(self, node):
        prev_flag = self.in_loop_flag
        self.in_loop_flag = True
        super().generic_visit(node)
        self.in_loop_flag = prev_flag
    visit_For = visit_While

    def visit_Lambda(self, node):
        super().generic_visit(node)
        if self.in_loop_flag:
            self.add_error(node)
    visit_FunctionDef = visit_Lambda

class DesugaringTransform(NodeTransformer):
    """
    This "desugars" the AST by
    - separating out function calls to individual assignments
    - converting <exp> if <cond> else <exp> statements to explicit
    if-else blocks
    - converting lambda expressions to function def
    - converting logical and/or statements to equivalent if/else blocks
    - make None function return explicit
    - converting branching expressions into function calls

    Currently not implemented desugaring:
    - de-decorating function definitions
    """
    temporary_var_prefix = '__v'
    def __call__(self, rootnode):
        self.new_stmt_stack = [[],]
        self.in_expression = False
        self.n_temporary_vars = 0
        self.visit(rootnode)
        rootnode = ast.parse(ast.unparse(rootnode)) #HACK: this might be slow
        return rootnode

    def visit(self, node):
        if isinstance(node, ast.stmt):
            method = "visit_stmt"
        elif (
            isinstance(node, ast.expr) and \
            self.is_branching_expr(node) and \
            not self.is_assigned_to(node)
        ):
            method = "visit_expr"
        else:
            method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)

    def visit_stmt(self, node):
        self.new_stmt_stack.append([])
        node = ast.NodeTransformer.visit(self, node)
        if node is not None:
            self.add_statement(node)
        return self.new_stmt_stack.pop()

    def visit_expr(self, node):
        """
        Since expressions won't change values in the namespace
        (aside from side-effects), we
        wrap *outermost* expression into a thunk to
        avoid needing to branch in the current function.
        """
        if not self.in_expression:
            self.in_expression = True
            thunk_node = ast.parse("(lambda : __expr__)()").body[0].value
            thunk_node.func.body = node
            node = ast.NodeTransformer.visit(self, thunk_node)
            self.in_expression = False
        else:
            node = ast.NodeTransformer.visit(self, node)
        return node

    @classmethod
    def is_assigned_to(cls, node):
        return hasattr(node, "ctx") and isinstance(node.ctx, ast.Store)


    @classmethod
    def is_branching_expr(cls, node):
        if isinstance(node, (ast.IfExp, ast.BoolOp)):
            return True
        for child in ast.iter_child_nodes(node):
            if cls.is_branching_expr(child):
                return True
        return False

    def generate_name(self):
        self.n_temporary_vars += 1
        return f"{DesugaringTransform.temporary_var_prefix}{self.n_temporary_vars - 1}"

    # def visit_Call(self, node):
    #     node = ast.NodeTransformer.generic_visit(self, node)
    #     return_name = self.generate_name()
    #     assn_node = ast.Assign(
    #         targets=[ast.Name(id=return_name, ctx=ast.Store())],
    #         value=node,
    #         lineno=0
    #     )
    #     self.add_statement(assn_node)
    #     return ast.Name(id=return_name, ctx=ast.Load())

    def visit_IfExp(self, node):
        """
        Convert <exp> if <cond> else <exp> to explicit if then else blocks
        """
        return self.desugar_to_IfElse_block(
            test_expr=node.test,
            if_expr=node.body,
            else_expr=node.orelse,
            test_name = self.generate_name(),
            return_name = self.generate_name()
        )

    def visit_Lambda(self, node):
        def_name = self.generate_name()
        def_node = ast.parse(textwrap.dedent(f"""
        def {def_name}():
            return None
        """)).body[0]
        def_node.args = node.args
        def_node.body[0].value = node.body
        def_node = ast.NodeTransformer.generic_visit(self, def_node)
        self.add_statement(def_node)
        return ast.Name(id=def_name, ctx=ast.Load())

    def visit_Return(self, node):
        if node.value is None:
            node.value = ast.Constant(value=None)
        node = self.generic_visit(node)
        return node

    def visit_BoolOp(self, node):
        if len(node.values) > 2:
            leftmost = ast.BoolOp(
                op=type(node.op)(),
                values=node.values[:-1]
            )
            node.values = [leftmost, node.values[-1]]
            return self.visit_BoolOp(node)
        test_name = self.generate_name()
        return_name = self.generate_name()
        if isinstance(node.op, ast.And):
            return self.desugar_to_IfElse_block(
                test_expr=node.values[0],
                if_expr=node.values[1],
                else_expr=ast.Name(id=test_name),
                test_name=test_name,
                return_name=return_name
            )
        elif isinstance(node.op, ast.Or):
            return self.desugar_to_IfElse_block(
                test_expr=node.values[0],
                if_expr=ast.Name(id=test_name),
                else_expr=node.values[1],
                test_name=test_name,
                return_name=return_name
            )
        raise ValueError("BoolOp is neither And nor Or")

    def visit_FunctionDef(self, node):
        node = self.generic_visit(node)
        # make return value of None function explicit
        # Note: this is required for CPSTransform to work
        if not isinstance(node.body[-1], ast.Return):
            return_none = ast.parse("return None").body[0]
            node.body.append(return_none)
        return node

    # def de_decorate_FunctionDef(self, node):
    #     decorator_list, node.decorator_list = node.decorator_list, []
    #     new_block = [self.visit(node)]
    #     for decorator in decorator_list[::-1]:
    #         decorator_stmt = ast.parse(f"{node.name} = __dec__({node.name})").body[0]
    #         decorator_stmt.value.func = decorator
    #         decorator_stmt = self.visit(decorator_stmt)
    #         new_block.append(decorator_stmt)
    #     return new_block

    def visit_comprehension(self, node):
        '''
        Convert comprehensions into recursive_reduce calls. Later generators correspond to
        nested recursive_reduce calls. For the following example input:

        ```
        [
            (x, y)
            for x in range(4)
            if x % 2 == 0
            for y in range(5)
            if x + y < 3
        ]
        ```

        We transform code into the following:

        ```
        def __v0(_acc, x):
            def __v1(_acc, y):
                return __acc + [(x, y)] if all([x + y < 3]) else __acc
            return __acc + recursive_reduce(__v1, range(5), []) if all([x % 2 == 0]) else __acc
        recursive_reduce(__v0, range(4), [])
        ```

        This function can also handle set and dictionary comprehensions.
        '''
        self.add_statement(ast.parse('from flippy import recursive_reduce').body[0])

        nested_cb = None
        if isinstance(node, ast.ListComp):
            op = '+'
            initial = '[]'
            nested_comp = ast.List(elts=[node.elt])
        elif isinstance(node, ast.SetComp):
            op = '|'
            initial = 'set()'
            nested_comp = ast.Set(elts=[node.elt])
        else:
            assert isinstance(node, ast.DictComp)
            op = '|'
            initial = '{}'
            nested_comp = ast.Dict(keys=[node.key], values=[node.value])

        for g in node.generators[::-1]:
            # We make sure to specialize the function body based on whether tests were provided.
            base_expr = f'''__acc {op} {Placeholder.new('nested_comp')}'''
            if len(g.ifs) == 0:
                test = None
                expr = base_expr
            else:
                # We separately handle the common case (one test) from the less common case of multiple tests.
                if len(g.ifs) == 1:
                    test = g.ifs[0]
                else:
                    test = ast.BoolOp(op=ast.And(), values=g.ifs)
                expr = f'''{base_expr} if {Placeholder.new('test')} else __acc'''

            # We add a nested callback so that we can destructure more complex targets. This is particularly useful
            # for dictionary comprehensions. Previously, we only had nested lambdas, which prevented supporting complex targets.
            cb_name = self.generate_name()
            cb = Placeholder.fill(ast.parse(textwrap.dedent(f'''
            def {cb_name}(__acc, __target):
                {Placeholder.new('target')} = __target
                {Placeholder.new('nested_cb')}
                return {expr}
            ''')), dict(
                # In the innermost loop, this CB is None, so the expression is just deleted.
                nested_cb=nested_cb,
                nested_comp=nested_comp,
                test=test,
                target=g.target
            )).body[0]

            new_node = Placeholder.fill(ast.parse(textwrap.dedent(f'''
            recursive_reduce({cb_name}, {Placeholder.new('iter')}, {initial})
            ''')), dict(
                iter=g.iter,
            )).body[0].value

            nested_cb = cb
            nested_comp = new_node

        # We need to add the outermost CB as a statement to the program.
        self.add_statement(self.visit(cb))

        return self.visit(new_node)

    visit_ListComp = visit_comprehension
    visit_SetComp = visit_comprehension
    visit_DictComp = visit_comprehension

    def visit_AugAssign(self, node):
        loadable_target = copy.copy(node.target)
        loadable_target.ctx = ast.Load()

        return self.visit(ast.Assign(
            targets=[node.target],
            value=ast.BinOp(
                left=loadable_target,
                op=node.op,
                right=node.value,
            ),
            lineno=node.lineno,
        ))

    def visit_AnnAssign(self, node):
        if node.value is None:
            # In this case, it's just a type annotation.
            return None
        return self.visit(ast.Assign(
            targets=[node.target],
            value=node.value,
            lineno=node.lineno,
        ))

    def visit_While(self, node):
        '''
        For simplicity, we desugar while loops to make tests be explicit if statements, mirroring how we handle
        other conditional control flow (like if expressions, logical and, and logical or).
        '''
        loop, = Placeholder.fill_from_source(f'''
        while True:
            if not ({Placeholder.new('test')}):
                break
            {Placeholder.new('body')}
        ''', dict(
            test=node.test,
            body=node.body,
        )).body
        # Doing this out of completeness, but it's a constant so nothing will change.
        self.generic_visit_field(loop, 'test')
        self.generic_visit_field(loop, 'body')
        return loop

    def visit_For(self, node):
        '''
        We implement a CPS-compatible `for` by assuming a loop argument that can be indexed and has a length.
        '''

        assert not node.orelse

        iter_sym = f'_for_iter_{node.lineno}'
        idx_sym = f'_for_idx_{node.lineno}'
        code = Placeholder.fill(ast.parse(textwrap.dedent(f'''
        {iter_sym} = {Placeholder.new('iter')}
        {idx_sym} = 0
        while {idx_sym} < len({iter_sym}):
            {Placeholder.new('target')} = {iter_sym}[{idx_sym}]
            {idx_sym} += 1
            {Placeholder.new('body')}
        ''')), dict(
            iter=node.iter,
            target=node.target,
            body=node.body,
        )).body

        return [
            self.visit(s)
            for s in code
        ]

    def desugar_to_IfElse_block(
        self,
        test_expr,
        if_expr,
        else_expr,
        test_name,
        return_name
    ):
        test_node, if_node = ast.parse(textwrap.dedent(f"""
        {test_name} = test
        if {test_name}:
            {return_name} = if_body
        else:
            {return_name} = else_body
        """)).body
        test_node.value = test_expr
        test_node = self.generic_visit(test_node)
        if not isinstance(test_node, list):
            test_node = [test_node]
        if_node.body[0].value = if_expr
        if_node.orelse[0].value = else_expr
        self.generic_visit(if_node)
        for stmt in [*test_node, if_node]:
            self.add_statement(stmt)
        return ast.Name(id=return_name, ctx=ast.Load())

    def add_statement(self, node):
        self.new_stmt_stack[-1].append(node)

class SetLineNumbers(ast.NodeTransformer):
    """
    Reset line numbers by statement.
    """
    def __call__(self, node, starting_line=0):
        self.cur_line = starting_line
        self.visit(node)
        ast.fix_missing_locations(node)
        return node
    def set_line(self, node):
        node.lineno = self.cur_line
        self.cur_line += 1
    def generic_visit(self, node):
        if isinstance(node, (ast.Expr, ast.stmt)):
            self.set_line(node)
        elif hasattr(node, 'lineno'):
            del node.lineno
        ast.NodeTransformer.generic_visit(self, node)
        return node

class GetLineNumber(ast.NodeVisitor):
    """
    Reset line numbers by statement.
    """
    def __call__(self, node: ast.AST, lineno: int):
        self.cur_line = 0
        self.lineno = lineno
        try:
            self.visit(node)
        except StopIteration as err:
            return err.value
        raise ValueError(f"Could not find line number {lineno} in {ast.dump(node)}")
    def set_line(self, node):
        node.lineno = self.cur_line
        self.cur_line += 1
    def generic_visit(self, node):
        if isinstance(node, (ast.Expr, ast.stmt)):
            if self.cur_line == self.lineno:
                raise StopIteration(node)
            self.cur_line += 1
        ast.NodeVisitor.generic_visit(self, node)
        return node

class CPSFunction:
    """
    This class is used to represent a function that has been transformed into
    continuation-passing style (CPS). It wraps a callable and stores the source code.
    It also provides a closure property that contains the variables captured by the function.
    The identity of a CPSFunction is determined by its source code and closure variables
    (see `__hash__` and `__eq__`).

    Note that the behavior of a function could depend on the contents of `global_store`,
    in which case the function is not a pure function. Inference algorithms are
    expected to be aware of this and handle the combination of the function and
    `global_store` appropriately.

    Separately, the behavior of a function could depend on variables defined in
    the global scope of the Python interpreter. In this case, it is assumed
    that these variables are not changed during the execution of the function.
    """
    def __init__(
        self,
        cps_func: Callable,
        func_src: str,
    ):
        self.cps_func = cps_func
        self.func_src = func_src
        setattr(self, CPSTransform.is_transformed_property, True)

    def __call__(self, *args, **kwargs):
        return self.cps_func(*args, **kwargs)

    @property
    def closure(self) -> hashabledict:
        if getattr(self.cps_func, "__closure__", None) is not None:
            closure_keys = self.cps_func.__code__.co_freevars
            closure_values = [cell.cell_contents for cell in self.cps_func.__closure__]
            return hashabledict(zip(closure_keys, closure_values))
        else:
            return hashabledict()

    @cached_property
    def _hash(self):
        # we sort the closure and use repr to avoid circular references
        return hash((self.func_src, repr(sorted(self.closure.items()))))

    def __hash__(self):
        return self._hash

    def __eq__(self, other):
        if not isinstance(other, CPSFunction):
            return False
        return \
            self.func_src == other.func_src and \
            self.closure == other.closure

    @property
    def __name__(self):
        return self.cps_func.__name__

class CPSTransform(NodeTransformer):
    """
    Convert python to a form of continuation passing style.
    """
    is_transformed_property = "_cps_transformed"
    func_src_name = "__func_src"
    cps_interpreter_name = "_cps"
    final_continuation_name = "_cont"
    stack_name = "_stack"

    def __init__(self):
        self.current_continuation_stmts = None
        self.parent_function_lineno = 0
        self.continue_break_replacement = None
        self.call_count = 0

    def __call__(self, node):
        self.scope_map = ClosureScopeAnalysis().analyze_scope(node)
        self.scope = None
        self.call_count = 0
        with self.in_scope(node):
            return self.visit(node)

    @classmethod
    def is_transformed(cls, func):
        return getattr(func, CPSTransform.is_transformed_property, False)

    def visit_Lambda(self, node):
        # Most lambdas are desugared, so this only has to handle the ones we generate,
        # like the default value for _cont of lambda val: val.
        with self.in_scope(node):
            return self.generic_visit(node)

    def visit_Module(self, node):
        # If we're in the outermost module scope, we only transform function definitions
        # Everything else is executed as normal (deterministic) python
        # This is primarily to handle cases where non-function definition statements
        # are created in the course of desugaring but still need to be executed.
        for stmt in node.body:
            if isinstance(stmt, ast.FunctionDef):
                self.visit(stmt)
        return node

    def visit_FunctionDef(self, node):
        node = self.add_function_src_to_FunctionDef_body(node)
        with self.in_scope(node):
            self.visit(node.args)
            # we track this so that line numbers passed to interpreter are
            # with respect to the start of the function definition
            old_parent_lineno = self.parent_function_lineno
            self.parent_function_lineno = node.lineno
            node.body, _ = self.transform_block(node.body)
            self.parent_function_lineno = old_parent_lineno

        if hasattr(node, 'body_imports'):
            for i in node.body_imports:
                node.body.insert(1, i) # 1 to make them after function source
            del node.body_imports

        node = self.add_keyword_to_FunctionDef(node, self.stack_name, "()")
        node = self.add_keyword_to_FunctionDef(node, self.cps_interpreter_name, self.cps_interpreter_name)
        node = self.add_keyword_to_FunctionDef(node, self.final_continuation_name, "lambda val: val")

        # decorator = ast.parse(f"lambda fn: (fn, setattr(fn, '{self.is_transformed_property}', True))[0]").body[0].value

        # This decorator position makes it the last decorator called.
        # This ensures that the outermost function has `is_transformed_property` set.
        # node.decorator_list.insert(0, decorator)
        # This decorator position makes it the first decorator called.
        # This is important so that intermediate decorators can change their behavior
        # the second time that this function is executed.
        # Functions are executed a second time when being evaluted after being CPS-transformed
        # because decorators aren't removed by the transform.
        # node.decorator_list.append(decorator)
        return node

    def add_keyword_to_FunctionDef(self, node, key, value):
        assert isinstance(node, ast.FunctionDef)
        cur_keywords = [kw.arg for kw in node.args.kwonlyargs + node.args.args]
        assert key not in cur_keywords
        node.args.kwonlyargs.append(ast.arg(arg=key))
        node.args.kw_defaults.append(ast.parse(value).body[0].value)
        return node

    def add_function_src_to_FunctionDef_body(self, node):
        assert isinstance(node, ast.FunctionDef)
        ctx_id = repr(ast.unparse(node))
        ctx_id_assn = ast.parse(f"{self.func_src_name} = {ctx_id}").body[0]
        node.body.insert(0, ctx_id_assn)
        node.decorator_list.append(
            ast.parse(f"lambda fn: CPSFunction(fn, {ctx_id})").body[0].value
        )
        return node

    @contextlib.contextmanager
    def in_scope(self, node):
        prev_scope = self.scope
        self.scope = self.scope_map[node]
        try:
            yield
        finally:
            self.scope = prev_scope

    def scope_packing_for_current_scope(self, locals_name, scope_name, *, hashabledict=False):
        '''
        Our general scheme for scope packing is to pack the possibly-present names
        into a dictionary, if they are defined. Variables are unpacked in a way that
        handles the issue of definition, by only conditionally defining variables.

        Note that this formerly tracked names as they were defined in the local scope,
        but this was removed in favor of simply adding all names that could possibly be in
        the present scope. We did this because loops make it more difficult to reason about
        variables, requiring some care to ensure that nested continuation-creating nodes
        would propagate variables that would be defined later in the loop body.
        '''
        names = sorted(self.scope.ns.keys())
        import_ = 'from flippy.hashable import hashabledict' if hashabledict else ''
        scope_type = 'hashabledict' if hashabledict else ''
        pack = f'''
        {import_}
        # We capture locals before we enter any other scope (like a lambda or list comprehension)
        {locals_name} = locals()
        {scope_name} = {scope_type}({{name: {locals_name}[name] for name in {names} if name in {locals_name}}})
        '''
        unpack = '\n'.join(f'''
        if "{name}" in {scope_name}:
            {name} = {scope_name}["{name}"]
        ''' for name in names) or 'pass'
        return (
            ast.parse(textwrap.dedent(pack)).body,
            ast.parse(textwrap.dedent(unpack)).body,
        )

    def visit_Assign(self, node):
        # We make sure to visit values first.
        self.generic_visit_field(node, 'value')
        # Then we visit targets.
        self.generic_visit_field(node, 'targets')
        return node

    def visit_Call(self, node):
        '''
        Replace the Call node with a continuation and call to the interpreter
        that picks up at the continuation.
        '''
        call_id = self.call_count
        self.call_count += 1
        continuation_name = f"_cont_{call_id}"
        result_name = f"_res_{call_id}"
        locals_name = f"_locals_{call_id}" #NOTE: could optimize with _locals_{node.lineno}
        pack, unpack = self.scope_packing_for_current_scope(locals_name, f"_scope_{call_id}")
        code = Placeholder.fill(ast.parse(textwrap.dedent(f'''
            {Placeholder.new('pack')}
            def {continuation_name}({result_name}):
                {Placeholder.new('unpack')}
            return lambda : {self.cps_interpreter_name}.interpret(
                {Placeholder.new('func')},
                cont={continuation_name},
                stack={self.stack_name},
                func_src={self.func_src_name},
                locals_={locals_name},
                call_id={call_id},
                lineno={node.lineno - self.parent_function_lineno}
            )()
        ''')), dict(
            func=self.visit(node.func),
            pack=pack,
            unpack=unpack,
        )).body
        continuation_node, thunk_node = code[-2:]
        thunk_node.value.body.args = [self.visit(arg) for arg in node.args]
        thunk_node.value.body.keywords = [self.visit(kw) for kw in node.keywords]
        self.current_continuation_stmts.extend(code)
        # Update CCS to the continuation of our function.
        self.current_continuation_stmts = continuation_node.body
        return ast.Name(id=result_name, ctx=ast.Load())

    def visit_While(self, node):
        '''
        We transform While nodes into a CPS style, by having
        1) a conditionally-recursive (dependent on loop test) continuation that executes the loop body and
        2) a final continuation, visited on loop break or False value from loop test.
        We have to carefully pack and unpack scope through all these functions. The current implementation
        handles temporary variables and changes to scope before a continue/break.
        '''
        # Store previous continue/break replacements, for recursive processing
        previous = self.continue_break_replacement

        call_id = self.call_count
        self.call_count += 1

        loop_name = f"_loop_fn_{call_id}"
        loop_exit_name = f"_loop_exit_fn_{call_id}"
        scope_name = f"_scope_{call_id}"
        locals_name = f"_locals_{call_id}"

        # We process body, replacing continue/break with placeholders.
        # We initially replace with placeholders because continue/break requires packing any
        # variables that could be defined in loop body. Once we know the variables that
        # need to be packed, we replace with the appropriate calls (below).
        self.continue_break_replacement = (
            ast.Expr(ast.Name(Placeholder.new('continue_'))),
            ast.Expr(ast.Name(Placeholder.new('break_'))))
        node_body, continuation_block = self.transform_block(node.body)

        # Now that we've transformed the loop body, we know all variables that can be defined.
        pack, unpack = self.scope_packing_for_current_scope(locals_name, scope_name, hashabledict=True)
        def loop_call_template(cont):
            return textwrap.dedent(f'''
                return lambda: {self.cps_interpreter_name}.interpret(
                    {loop_name},
                    cont={cont},
                    stack={self.stack_name},
                    func_src={self.func_src_name},
                    locals_={locals_name},
                    lineno={node.lineno - self.parent_function_lineno}
                )({scope_name})''')
        # For the loop's initial call, we supply the post-loop continuation. However, for the
        # recursive calls (at end of loop body or for continue statements), we use final_continuation_name
        loop_init_call = pack + ast.parse(loop_call_template(loop_exit_name)).body
        loop_recur_call = pack + ast.parse(loop_call_template(self.final_continuation_name)).body
        loop_exit_call = pack + ast.parse(f'return lambda: {self.final_continuation_name}(_ValueWrapper({scope_name}))').body
        # At the end of body execution, we recurse. Must happen before filling placeholders in node_body.
        continuation_block.extend(loop_recur_call)
        # Replace continue/break with call that has appropriate variable packing.
        node_body = [Placeholder.fill(s, dict(continue_=loop_recur_call, break_=loop_exit_call)) for s in node_body]

        # We import _ValueWrapper at the beginning of our enclosing scope.
        assert isinstance(self.scope.node, ast.FunctionDef)
        self.scope.node.body_imports = ast.parse('from flippy.transforms import _ValueWrapper').body

        code = Placeholder.fill(ast.parse(textwrap.dedent(f'''
            def {loop_name}({scope_name}):
                {Placeholder.new('unpack')}
                if {Placeholder.new('test')}:
                    {Placeholder.new('body')}
                else:
                    # When exiting, we have no further assignments, so we could reuse initial scope, but don't.
                    {Placeholder.new('loop_exit_call')}
            def {loop_exit_name}({scope_name}):
                if not isinstance({scope_name}, _ValueWrapper):
                    return lambda: {self.final_continuation_name}({scope_name})
                {scope_name} = {scope_name}.value
                {Placeholder.new('unpack')}
            {Placeholder.new('loop_init_call')}
        ''')), dict(
            # NOTE: We assume the above desugaring, so that the loop test is trivial (always True), and our loop predicate
            # is instead being tested with an if statement in the loop's body.
            test=node.test,
            body=node_body,
            unpack=unpack,
            loop_init_call=loop_init_call,
            loop_recur_call=loop_recur_call,
            loop_exit_call=loop_exit_call,
        )).body
        loop_fn, loop_exit_fn = code[:2]

        # Next few lines are copied from visit_FunctionDef
        loop_fn = self.add_function_src_to_FunctionDef_body(loop_fn)
        loop_fn = self.add_keyword_to_FunctionDef(loop_fn, self.stack_name, "()")
        loop_fn = self.add_keyword_to_FunctionDef(loop_fn, self.cps_interpreter_name, self.cps_interpreter_name)
        loop_fn = self.add_keyword_to_FunctionDef(loop_fn, self.final_continuation_name, "lambda val: val")
        assert loop_fn is code[0], 'Making sure the object was mutated'

        self.current_continuation_stmts.extend(code)
        # Execution continues in the loop exit continuation.
        self.current_continuation_stmts = loop_exit_fn.body
        # Restore, for recursive processing
        self.continue_break_replacement = previous
        return None

    def visit_Continue(self, node):
        c, b = self.continue_break_replacement
        return c
    def visit_Break(self, node):
        c, b = self.continue_break_replacement
        return b

    def transform_block(self, untransformed_stmts):
        '''
        This function transforms each statement in the argument `untransformed_stmts`, in order,
        into CPS. It maintains a list to hold the transformed statements. However,
        in the course of the CPS transform, Call/If/While will produce new continuations.
        So, this algorithm maintains a pointer to the current continuation, updated by
        continuation-generating visitor methods.
        '''

        # Store current state, so we can recursively transform.
        previous = self.current_continuation_stmts

        # We create a block to hold CPS-transformed statements.
        # Initially, these refer to the same list, but self.current_continuation_stmts
        # will change after we process a statement that needs a CPS-transform
        # (i.e. after a Call, While, or If).
        transformed_stmts = self.current_continuation_stmts = []

        for us in untransformed_stmts:
            rv = self.visit(us)
            # We drop return values that are None. This happens in visit_If and visit_While.
            if rv is not None:
                # Note: self.current_continuation_stmts can be modified in other functions,
                # which means this doesn't always append into transformed_stmts.
                if isinstance(rv, list):
                    self.current_continuation_stmts.extend(rv)
                else:
                    self.current_continuation_stmts.append(rv)

        # Hold onto a reference before we restore state.
        ccs = self.current_continuation_stmts
        # Restore previous state.
        self.current_continuation_stmts = previous

        return transformed_stmts, ccs

    def visit_If(self, node):
        """
        We transform each branch of the If node individually, making sure they each wind up
        executing the same continuation, which contains future statements.
        """
        node.test = self.visit(node.test)
        node.body, body_continuation = self.transform_block(node.body)
        node.orelse, orelse_continuation = self.transform_block(node.orelse)
        call_id = self.call_count
        self.call_count += 1

        continuation_name = f"_cont_{call_id}"
        scope_name = f"_scope_{call_id}"
        # Should only do scope packing after variables from both branches have been analyzed.
        pack, unpack = self.scope_packing_for_current_scope(f"_locals_{call_id}", scope_name)
        continuation_node, = Placeholder.fill(ast.parse(textwrap.dedent(f'''
            def {continuation_name}({scope_name}):
                {Placeholder.new('unpack')}
        ''')), dict(unpack=unpack)).body
        conditional_done_call = pack + ast.parse(f'return lambda: {continuation_name}({scope_name})').body
        # Add the continuation
        self.current_continuation_stmts.extend([continuation_node, node])
        # Make branches go to continuation
        body_continuation += conditional_done_call
        orelse_continuation += conditional_done_call
        # Future statements go in continuation body
        self.current_continuation_stmts = continuation_node.body
        return None

    def visit_Return(self, node):
        new_node = ast.parse(f"return lambda : {self.final_continuation_name}(_res)").body[0]
        new_node.value.body.args[0] = self.visit(node.value)
        return new_node

class HashableCollectionTransform(ast.NodeTransformer):
    """
    Wrap all lists and dicts in hashable versions.
    """
    def __call__(self, rootnode):
        self.visit(rootnode)
        return rootnode

    def visit_List(self, node):
        return ast.Call(
            func=ast.Name(id='hashablelist', ctx=ast.Load()),
            args=[node],
            keywords=[]
        )

    def visit_Dict(self, node):
        return ast.Call(
            func=ast.Name(id='hashabledict', ctx=ast.Load()),
            args=[node],
            keywords=[]
        )
