# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.

"""
This module defines the `DSLPreprocessor` class, which acts as a Python preprocessor.
It uses Python's AST and rewrites specific Python statements such as `for` and `if-else`.

The preprocessor operates on the following constructs:
    - `for` loops:
        - Rewrites `for` loops with the `@loop_selector` decorator.
        - Supports `range`, `range_dynamic`, and `range_constexpr` for loop iteration.
    - `if-elif-else` statements:
        - Rewrites conditional statements with the `@if_selector` decorator.
        - Supports `dynamic_expr` and `const_expr` in the condition expressions.

Additionally, both `for` loops and `if-else` statements require `yield`
operation generation. The preprocessor handles this by:
    - Using a `ScopeManager` to track symbols across different scopes during AST traversal.
    - Identifying read-only, read-write, and active variables for DSL constructs.
    - Generating `yield` operations for symbols that are classified as read-write or write.

It is designed to be generic and can handle `for` and `if` constructs from other dialects.
In such cases, the user's DSL should implement `@loop_selector` and `@if_selector`
to generate dialect-specific operations for `for` and `if` statements.
"""

import ast
import importlib
import inspect
import textwrap
from dataclasses import dataclass
from typing import List, Set, Dict, Any, Callable, Optional
from types import ModuleType

from .common import *
from .utils.logger import log


class OrderedSet:
    """
    A deterministic set implementation for ordered operations.
    """

    def __init__(self, iterable=None):
        self._dict = dict.fromkeys(iterable or [])

    def add(self, item):
        self._dict[item] = None

    def __iter__(self):
        return iter(self._dict)

    def __and__(self, other):
        return OrderedSet(key for key in self._dict if key in other)

    def __or__(self, other):
        new_dict = self._dict.copy()
        new_dict.update(dict.fromkeys(other))
        return OrderedSet(new_dict)

    def __sub__(self, other):
        return OrderedSet(key for key in self._dict if key not in other)


@dataclass
class ScopeManager:
    """
    Manages symbol scopes during AST traversal.
    Manage nested scopes during transformations.
    """

    scopes: List[Set[str]]
    current_scope: Set[str]

    @classmethod
    def create(cls) -> "ScopeManager":
        return cls([], set())

    def enter_scope(self) -> None:
        self.scopes.append(self.current_scope.copy())

    def exit_scope(self) -> None:
        self.current_scope = self.scopes.pop()

    def add_to_scope(self, name: str) -> None:
        self.current_scope.add(name)

    def get_active_symbols(self) -> Set[str]:
        return set(self.current_scope)


class DSLPreprocessor(ast.NodeTransformer):
    """
    A preprocessor for transforming Python ASTs. It supports:

    - Rewriting `for` loops with the `@loop_selector` decorator.
    - Rewriting `if-elif-else` statements with the `@if_selector` decorator.
    - Generating `yield` operations for read-write or write symbols.
    """

    DECORATOR_FOR_STATEMENT = "loop_selector"
    DECORATOR_IF_STATEMENT = "if_selector"
    DECORATOR_WHILE_STATEMENT = "while_selector"
    IF_EXECUTOR = "if_executor"
    WHILE_EXECUTOR = "while_executor"
    ASSERT_EXECUTOR = "assert_executor"
    BOOL_CAST = "bool_cast"
    IMPLICIT_DOWNCAST_NUMERIC_TYPE = "implicitDowncastNumericType"
    SUPPORTED_FOR_RANGE_STATEMENTS = {"range", "range_dynamic", "range_constexpr"}

    def __init__(self):
        super().__init__()
        self.counter = 0  # Unique function names for multiple loops
        self.scope_manager = ScopeManager.create()
        self.processed_functions = set()
        self.function_counter = 0
        self.function_name = "<unknown function>"
        self.class_name = None
        self.file_name = "<unknown filename>"
        self.function_depth = 0
        self.local_closures = set()
        self.function_globals = None

    def _get_module_imports(self, decorated_func):
        """Extract imports from the module containing the decorated function"""
        # Get the module containing the decorated function
        module = inspect.getmodule(decorated_func)
        if module is None:
            return {}

        # Get the module source code
        try:
            source = inspect.getsource(module)
            module_ast = ast.parse(source)

            # Extract imports from the full module
            imports = {}
            for node in ast.walk(module_ast):
                if isinstance(node, ast.Import):
                    for name in node.names:
                        imports[name.name] = name.asname if name.asname else name.name
                elif isinstance(node, ast.ImportFrom):
                    module_name = node.module
                    for name in node.names:
                        if name.name == "*":
                            # Handle wildcard imports
                            try:
                                imported_module = importlib.import_module(module_name)
                                imports[module_name] = imported_module
                            except ImportError:
                                pass
                        else:
                            full_name = f"{module_name}.{name.name}"
                            imports[full_name] = (
                                name.asname if name.asname else name.name
                            )
            return imports
        except (IOError, TypeError):
            return {}

    def exec(self, function_name, original_function, code_object, exec_globals):
        # Get imports from the original module
        module_imports = self._get_module_imports(original_function)

        # Import all required modules
        for module_path, alias in module_imports.items():
            try:
                if "." in module_path:
                    base_module, attribute = module_path.rsplit(".", 1)
                    module = importlib.import_module(base_module)
                    if hasattr(module, attribute):
                        attr = getattr(module, attribute)
                        exec_globals[alias] = attr
                else:
                    path = importlib.import_module(module_path)
                    exec_globals[alias] = path
            except (ImportError, AttributeError) as e:
                raise ImportError(f"Failed to import {module_path}: {str(e)}")

        # Execute the transformed code
        log().info(
            "ASTPreprocessor Executing transformed code for function [%s]",
            function_name,
        )
        exec(code_object, exec_globals)
        return exec_globals.get(function_name)

    @staticmethod
    def print_ast(transformed_tree=None):
        print("#", "-" * 40, "Transformed AST", "-" * 40)
        unparsed_code = ast.unparse(transformed_tree)
        print(unparsed_code)
        print("#", "-" * 40, "End Transformed AST", "-" * 40)

    def make_func_param_name(self, base_name, used_names):
        """Generate a unique parameter name that doesn't collide with existing names."""
        if base_name not in used_names:
            return base_name

        i = 0
        while f"{base_name}_{i}" in used_names:
            i += 1
        return f"{base_name}_{i}"

    def transform_function(self, func_name, function_pointer):
        """
        Transforms a function.
        """
        # Skip if the function has already been processed
        if function_pointer in self.processed_functions:
            log().info(
                "ASTPreprocessor Skipping already processed function [%s]", func_name
            )
            return []

        # Step 1. Parse the given function
        file_name = inspect.getsourcefile(function_pointer)
        lines, start_line = inspect.getsourcelines(function_pointer)
        dedented_source = textwrap.dedent("".join(lines))
        tree = ast.parse(dedented_source, filename=file_name)
        # Bump the line numbers so they match the real source file
        ast.increment_lineno(tree, start_line - 1)

        # Step 1.2 Check the decorator
        if not self.check_decorator(tree.body[0]):
            log().info(
                "[%s] - Skipping function due to missing decorator",
                func_name,
            )
            return []

        self.processed_functions.add(function_pointer)
        log().info("ASTPreprocessor Transforming function [%s]", func_name)

        # Step 2. Transform the function
        transformed_tree = self.visit(tree)
        ast.fix_missing_locations(transformed_tree)
        combined_body = transformed_tree.body

        # Step 3. Return the transformed tree
        return combined_body

    def check_early_exit(self, tree, kind):
        """
        Checks if a given region or scope in the provided Python code has early exits.
        """

        class EarlyExitChecker(ast.NodeVisitor):
            def __init__(self, kind):
                self.has_early_exit = False
                self.early_exit_node = None
                self.early_exit_type = None
                self.kind = kind
                self.loop_nest_level = 0

            # Early exit is not allowed in any level of dynamic control flow
            def visit_Return(self, node):
                self.has_early_exit = True
                self.early_exit_node = node
                self.early_exit_type = "return"

            def visit_Raise(self, node):
                self.has_early_exit = True
                self.early_exit_node = node
                self.early_exit_type = "raise"

            def visit_Break(self, node):
                # For break/continue in inner loops, we don't consider it as early exit
                if self.loop_nest_level == 0 and self.kind != "if":
                    self.has_early_exit = True
                    self.early_exit_node = node
                    self.early_exit_type = "break"

            def visit_Continue(self, node):
                if self.loop_nest_level == 0 and self.kind != "if":
                    self.has_early_exit = True
                    self.early_exit_node = node
                    self.early_exit_type = "continue"

            def visit_For(self, node):
                self.loop_nest_level += 1
                self.generic_visit(node)
                self.loop_nest_level -= 1

            def visit_While(self, node):
                self.loop_nest_level += 1
                self.generic_visit(node)
                self.loop_nest_level -= 1

        checker = EarlyExitChecker(kind)
        checker.generic_visit(tree)
        if not checker.has_early_exit:
            return
        raise DSLAstPreprocessorError(
            message=f"Early exit ({checker.early_exit_type}) is not allowed in `{self.function_name}`"
            + (f" in `{self.class_name}`" if self.class_name else ""),
            filename=self.file_name,
            snippet=ast.unparse(tree),
            suggestion=(
                "If predicates are constant expression, write like "
                "`if const_expr(...)` or `for ... in range_constexpr`. "
                "In that case, early exit will be executed by Python "
                "interpreter, so it's supported."
            ),
        )

    def is_node_constexpr(self, node) -> bool:
        """
        Determines if the node is a constexpr.
        Supported nodes are if, for, while statements.
        """
        if isinstance(node, ast.If) or isinstance(node, ast.While):
            if isinstance(node.test, ast.Call):
                func = node.test.func

                if isinstance(func, ast.Attribute) and func.attr == "const_expr":
                    return True

                elif isinstance(func, ast.Name) and func.id == "const_expr":
                    return True
        elif isinstance(node, ast.For):
            if isinstance(node.iter, ast.Call):
                func = node.iter.func
                if isinstance(func, ast.Attribute) and func.attr == "range_constexpr":
                    return True

                elif isinstance(func, ast.Name) and func.id == "range_constexpr":
                    return True
        return False

    def transform(self, original_function, exec_globals):
        """
        Transforms the provided function using the preprocessor.
        """
        self.file_name = inspect.getsourcefile(original_function)
        self.function_globals = exec_globals
        transformed_tree = self.transform_function(
            original_function.__name__, original_function
        )
        unified_tree = ast.Module(body=transformed_tree, type_ignores=[])
        unified_tree = ast.fix_missing_locations(unified_tree)

        return unified_tree

    def analyze_region_variables(self, node: Union[ast.For, ast.If], active_symbols):
        """
        Analyze variables in different code regions to identify read-only, write-only,
        and active variables for DSL constructs.
        """

        # we need orderedset to keep the insertion order the same. otherwise generated IR is different each time
        read_args = OrderedSet()
        write_args = OrderedSet()
        local_closure = self.local_closures
        file_name = self.file_name
        region_node = node

        class RegionAnalyzer(ast.NodeVisitor):

            def visit_Name(self, node):
                """
                Mark every load as read, and every store as write.
                """
                if isinstance(node.ctx, ast.Load):
                    read_args.add(node.id)
                elif isinstance(node.ctx, ast.Store):
                    write_args.add(node.id)

            @staticmethod
            def get_call_base(func_node):
                if isinstance(func_node, ast.Attribute):
                    # If the .value is another Attribute, keep digging
                    if isinstance(func_node.value, ast.Attribute):
                        return RegionAnalyzer.get_call_base(func_node.value)
                    # If the .value is a Name, that's our base
                    elif isinstance(func_node.value, ast.Name):
                        return func_node.value.id
                    else:
                        # Could be something else (lambda, call, etc.)
                        return None
                elif isinstance(func_node, ast.Name):
                    return None
                return None

            @staticmethod
            def get_function_name(func_node: ast.Call):
                if isinstance(func_node.func, ast.Name):
                    function_name = func_node.func.id
                # Check if it's a method or attribute call
                elif isinstance(func_node.func, ast.Attribute):
                    function_name = func_node.func.attr
                else:
                    function_name = None
                return function_name

            def visit_Call(self, node):
                base_name = RegionAnalyzer.get_call_base(node.func)

                if isinstance(node.func, ast.Name):
                    func_name = node.func.id
                    if func_name in local_closure:
                        raise DSLAstPreprocessorError(
                            f"Function `{func_name}` is a closure and is not supported in for/if statements",
                            filename=file_name,
                            snippet=ast.unparse(region_node),
                        )

                # Classes are mutable by default. Mark them as write. If they are
                # dataclass(frozen=True), treat them as read in runtime.
                if base_name is not None and base_name not in ("self"):
                    write_args.add(base_name)

                self.generic_visit(node)

        analyzer = RegionAnalyzer()
        analyzer.visit(ast.Module(body=node))

        # Argument can be Load and Store. We should just mark it as Store.
        read_args = read_args - write_args

        used_args = read_args & active_symbols
        iter_args = write_args & active_symbols
        flattend_args = used_args | iter_args

        return list(used_args), list(iter_args), list(flattend_args)

    def extract_range_args(self, iter_node):
        args = iter_node.args
        if len(args) == 1:
            return ast.Constant(value=0), self.visit(args[0]), ast.Constant(value=1)
        elif len(args) == 2:
            return self.visit(args[0]), self.visit(args[1]), ast.Constant(value=1)
        elif len(args) == 3:
            return self.visit(args[0]), self.visit(args[1]), self.visit(args[2])
        else:
            raise DSLAstPreprocessorError(
                "Unsupported number of arguments in range", filename=self.file_name
            )

    def extract_unroll_args(self, iter_node):
        keywords = {kw.arg: kw.value for kw in iter_node.keywords}
        return (
            keywords.get("unroll", ast.Constant(value=-1)),
            keywords.get("unroll_full", ast.Constant(value=False)),
        )

    def create_loop_function(
        self,
        func_name,
        node,
        start,
        stop,
        step,
        unroll,
        unroll_full,
        used_args,
        iter_args,
        flattened_args,
        is_loop_constexpr,
    ):
        """
        Creates a loop body function with the `loop_selector` decorator.
        """

        func_args = [ast.arg(arg=node.target.id, annotation=None)]
        func_args += [ast.arg(arg=var, annotation=None) for var in flattened_args]

        # Create the loop body
        transformed_body = []
        for stmt in node.body:
            transformed_stmt = self.visit(stmt)  # Recursively visit inner statements
            if isinstance(transformed_stmt, list):
                transformed_body.extend(transformed_stmt)
            else:
                transformed_body.append(transformed_stmt)

        # Handle the return for a single iterated argument correctly
        if len(iter_args) == 0:
            transformed_body.append(ast.Return())
        else:
            transformed_body.append(
                ast.Return(
                    value=ast.List(
                        elts=[ast.Name(id=var, ctx=ast.Load()) for var in iter_args],
                        ctx=ast.Load(),
                    )
                )
            )

        # Define the decorator with parameters
        decorator = ast.copy_location(
            ast.Call(
                func=ast.Name(id=self.DECORATOR_FOR_STATEMENT, ctx=ast.Load()),
                args=[start, stop, step],
                keywords=[
                ast.keyword(arg="unroll", value=unroll),
                ast.keyword(arg="unroll_full", value=unroll_full),
                ast.keyword(arg="constexpr", value=is_loop_constexpr),
                ast.keyword(
                    arg="used_args",
                    value=ast.List(
                        elts=[ast.Name(id=arg, ctx=ast.Load()) for arg in used_args],
                        ctx=ast.Load(),
                    ),
                ),
                ast.keyword(
                    arg="iter_args",
                    value=ast.List(
                        elts=[ast.Name(id=arg, ctx=ast.Load()) for arg in iter_args],
                        ctx=ast.Load(),
                        ),
                    ),
                    ast.keyword(
                        arg="iter_arg_names",
                        value=ast.List(
                            elts=[ast.Constant(value=arg) for arg in iter_args],
                            ctx=ast.Load(),
                        ),
                    ),
                ],
            ),
            node,
        )

        return ast.copy_location(
            ast.FunctionDef(
                name=func_name,
                args=ast.arguments(
                    posonlyargs=[],
                    args=func_args,
                    kwonlyargs=[],
                    kw_defaults=[],
                    defaults=[],
                ),
                body=transformed_body,
                decorator_list=[decorator],
            ),
            node,
        )

    def create_loop_call(self, func_name, iter_args):
        """
        Assigns the returned value from the loop function directly (without a tuple unpacking).
        """
        if len(iter_args) == 0:
            return ast.Expr(value=ast.Name(id=func_name, ctx=ast.Load()))
        elif len(iter_args) == 1:
            return ast.Assign(
                targets=[ast.Name(id=iter_args[0], ctx=ast.Store())],
                value=ast.Name(id=func_name, ctx=ast.Load()),
            )
        else:
            return ast.Assign(
                targets=[
                    ast.Tuple(
                        elts=[ast.Name(id=var, ctx=ast.Store()) for var in iter_args],
                        ctx=ast.Store(),
                    )
                ],
                value=ast.Name(id=func_name, ctx=ast.Load()),
            )

    def is_supported_range_call(self, node):
        return (
            isinstance(node, ast.For)
            and isinstance(node.iter, ast.Call)
            and (
                (
                    isinstance(node.iter.func, ast.Name)
                    and node.iter.func.id in self.SUPPORTED_FOR_RANGE_STATEMENTS
                )
                or (
                    isinstance(node.iter.func, ast.Attribute)
                    and node.iter.func.attr in self.SUPPORTED_FOR_RANGE_STATEMENTS
                )
            )
        )

    def get_loop_constexpr(self, node):
        if not self.is_supported_range_call(node):
            return None

        # Map function names to their constexpr values
        constexpr_map = {"range": None, "range_dynamic": False, "range_constexpr": True}
        range_name = (
            node.iter.func.id
            if isinstance(node.iter.func, ast.Name)
            else node.iter.func.attr
        )
        return ast.Constant(value=constexpr_map[range_name])

    def transform_for_loop(self, node, active_symbols):
        # Constexpr doesn't get preprocessed
        if self.is_node_constexpr(node):
            self.generic_visit(node)
            return node

        # We only support range, range_constexpr, range_dynamic
        if self.is_supported_range_call(node):
            constexpr_val = self.get_loop_constexpr(node)
            # Check for early exit and raise exception
            self.check_early_exit(node, "for")
            start, stop, step = self.extract_range_args(node.iter)
            unroll, unroll_full = self.extract_unroll_args(node.iter)
            used_args, iter_args, flat_args = self.analyze_region_variables(
                node, active_symbols
            )

            func_name = f"loop_body_{self.counter}"
            self.counter += 1

            func_def = self.create_loop_function(
                func_name,
                node,
                start,
                stop,
                step,
                unroll,
                unroll_full,
                used_args,
                iter_args,
                flat_args,
                constexpr_val,
            )

            assign = ast.copy_location(
                self.create_loop_call(func_name, iter_args), node
            )

            # This should work fine as it modifies the AST structure
            return [func_def, assign]

        self.generic_visit(node)

        return node

    def visit_BoolOp(self, node):
        # Visit child nodes first
        self.generic_visit(node)

        # It is necessary to expand short circuit evaluation explicit here
        # Although we do not support inline if-else for IR generation, this is actually evaluated in Python
        # So it's fine here
        # Transform "and" to "and_"
        if isinstance(node.op, ast.And):
            # Create an if-else statement in AST form
            # if type(lhs) == bool and lhs == False:
            #     return lhs
            # else
            #     return and_(lhs, rhs)
            short_circuit_value = ast.Constant(value=False)
            helper_func = ast.Name(id="and_", ctx=ast.Load())
        # Transform "or" to "or_"
        elif isinstance(node.op, ast.Or):
            # Create an if-else statement in AST form
            # if type(lhs) == bool and lhs == True:
            #     return lhs
            # else
            #     return or_(lhs, rhs)
            short_circuit_value = ast.Constant(value=True)
            helper_func = ast.Name(id="or_", ctx=ast.Load())
        else:
            # BoolOp should be either And or Or
            raise DSLAstPreprocessorError(
                f"Unsupported boolean operation: {node.op}",
                filename=self.file_name,
                snippet=ast.unparse(node),
            )

        def short_circuit_eval(value, short_circuit_value):
            return ast.BoolOp(
                op=ast.And(),
                values=[
                    ast.Compare(
                        left=ast.Call(
                            func=ast.Name(id="type", ctx=ast.Load()),
                            args=[value],
                            keywords=[],
                        ),
                        ops=[ast.Eq()],
                        comparators=[ast.Name(id="bool", ctx=ast.Load())],
                    ),
                    ast.Compare(
                        left=value,
                        ops=[ast.Eq()],
                        comparators=[short_circuit_value],
                    ),
                ],
            )

        lhs = node.values[0]

        for i in range(1, len(node.values)):
            test = short_circuit_eval(lhs, short_circuit_value)
            lhs = ast.IfExp(
                test=test,
                body=lhs,
                orelse=ast.Call(
                    func=helper_func,
                    args=[lhs, node.values[i]],
                    keywords=[],
                ),
            )

        return ast.copy_location(lhs, node)

    def visit_UnaryOp(self, node):
        # Visit child nodes first
        self.generic_visit(node)

        # Transform "not" to "~" as we overload __invert__
        if isinstance(node.op, ast.Not):
            func_name = ast.Name(id="not_", ctx=ast.Load())
            return ast.copy_location(
                ast.Call(func=func_name, args=[node.operand], keywords=[]), node
            )

        return node

    def visit_For(self, node):
        active_symbols = self.scope_manager.get_active_symbols()
        self.scope_manager.enter_scope()

        if isinstance(node.target, ast.Name):
            self.scope_manager.add_to_scope(node.target.id)

        new_for_node = self.transform_for_loop(node, active_symbols)
        self.scope_manager.exit_scope()
        return new_for_node

    def visit_Name(self, node):
        self.generic_visit(node)
        return node

    def visit_Assert(self, node):
        test = self.visit(node.test)

        args = [ast.keyword(arg="test", value=test)]
        if node.msg:
            msg = self.visit(node.msg)
            args.append(ast.keyword(arg="msg", value=msg))

        # Rewrite to assert_executor(test, msg)
        new_node = ast.Expr(
            ast.Call(
                func=ast.Name(id=self.ASSERT_EXECUTOR, ctx=ast.Load()),
                args=[],
                keywords=args,
            )
        )

        # Propagate line number from original node to new node
        ast.copy_location(new_node, node)
        return new_node

    def visit_Call(self, node):
        func = node.func
        self.generic_visit(node)

        # Check if the function is 'bool'
        if isinstance(func, ast.Name) and func.id == "bool":
            return ast.copy_location(
                ast.Call(
                    func=ast.Name(id=self.BOOL_CAST, ctx=ast.Load()),
                    args=[node.args[0]],
                    keywords=[],
                ),
                node,
            )
        elif isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name):
            def create_downcast_call(arg):
                return ast.copy_location(
                    ast.Call(
                        func=ast.Name(
                            id=self.IMPLICIT_DOWNCAST_NUMERIC_TYPE, ctx=ast.Load()
                        ),
                        args=[arg],
                        keywords=[],
                    ),
                    arg,
                )
            module = self.function_globals.get(func.value.id)
            if isinstance(module, ModuleType) and module.__package__.endswith(
                "._mlir.dialects"
            ):
                # Check if argument is Numeric, if so, call ir_value()
                args = []
                for arg in node.args:
                    args.append(create_downcast_call(arg))
                kwargs = []
                for kwarg in node.keywords:
                    kwargs.append(
                        ast.copy_location(
                            ast.keyword(
                                arg=kwarg.arg,
                                value=create_downcast_call(kwarg.value),
                            ),
                            kwarg,
                        )
                    )
                return ast.copy_location(
                    ast.Call(func=func, args=args, keywords=kwargs), node
                )

        return node

    def visit_ClassDef(self, node):
        self.class_name = node.name
        self.generic_visit(node)
        self.class_name = None
        return node

    def _visit_target(self, target):
        if isinstance(target, ast.Name):
            self.scope_manager.add_to_scope(target.id)
        elif isinstance(target, ast.Tuple):
            for t in target.elts:
                if isinstance(t, ast.Name):
                    self.scope_manager.add_to_scope(t.id)

    def visit_Assign(self, node):
        for target in node.targets:
            self._visit_target(target)
        self.generic_visit(node)
        return node

    def visit_AugAssign(self, node):
        self._visit_target(node.target)
        self.generic_visit(node)
        return node

    def check_decorator(self, node: ast.AST) -> bool:
        """
        Check if the function has the correct decorator for preprocessing.
        """
        if not isinstance(node, ast.FunctionDef):
            return False
        decorator_list = node.decorator_list
        if len(decorator_list) == 0:
            return False

        for d in decorator_list:
            if isinstance(d, ast.Call):
                if isinstance(d.func, ast.Attribute):
                    if d.func.attr in ["jit", "kernel"]:
                        if d.keywords == []:
                            return True
                        for keyword in d.keywords:
                            if keyword.arg == "preprocess":
                                try:
                                    if isinstance(keyword.value, ast.Constant):
                                        return keyword.value.value
                                    else:
                                        return ast.literal_eval(keyword.value)
                                except:
                                    pass

            elif isinstance(d, ast.Attribute):
                if d.attr in ["jit", "kernel"]:
                    return True

        return False

    def remove_dsl_decorator(self, decorator_list):
        """
        Remove .jit and .kernel decorators
        The decorator can be in two forms:
        - @jit(...)
        - @jit
        """
        new_decorator_list = []
        decorator_names = ["jit", "kernel"]
        for d in decorator_list:
            is_jit_or_kernel = False
            if isinstance(d, ast.Call):
                if isinstance(d.func, ast.Attribute):
                    if d.func.attr in decorator_names:
                        is_jit_or_kernel = True
            elif isinstance(d, ast.Attribute):
                if d.attr in decorator_names:
                    is_jit_or_kernel = True

            if not is_jit_or_kernel:
                new_decorator_list.append(d)
        return new_decorator_list

    def visit_FunctionDef(self, node):
        self.scope_manager.enter_scope()
        self.function_counter += 1
        self.function_name = node.name
        if self.function_depth > 0:
            self.local_closures.add(node.name)

        self.function_depth += 1

        # Add function name and arguments
        self.scope_manager.add_to_scope(node.name)
        for arg in node.args.args:
            self.scope_manager.add_to_scope(arg.arg)

        self.generic_visit(node)
        self.scope_manager.exit_scope()

        self.function_depth -= 1

        # Remove .jit and .kernel decorators
        node.decorator_list = self.remove_dsl_decorator(node.decorator_list)
        return node

    def visit_With(self, node):
        self.scope_manager.enter_scope()

        for item in node.items:
            if isinstance(item.optional_vars, ast.Name):
                self.scope_manager.add_to_scope(item.optional_vars.id)
        self.generic_visit(node)

        self.scope_manager.exit_scope()
        return node

    def visit_While(self, node):
        active_symbols = self.scope_manager.get_active_symbols()
        self.scope_manager.enter_scope()

        # Constexpr doesn't get preprocessed
        if self.is_node_constexpr(node):
            self.generic_visit(node)
            self.scope_manager.exit_scope()
            return node

        # Check for early exit and raise exception
        self.check_early_exit(node, "while")

        used_args, yield_args, flat_args = self.analyze_region_variables(
            node, active_symbols
        )
        func_name = f"while_region_{self.counter}"
        self.counter += 1

        func_def = self.create_while_function(
            func_name, node, used_args, yield_args, flat_args
        )
        assign = ast.copy_location(self.create_loop_call(func_name, yield_args), node)

        self.scope_manager.exit_scope()
        return [func_def, assign]

    def visit_Try(self, node):
        self.scope_manager.enter_scope()
        self.generic_visit(node)
        self.scope_manager.exit_scope()
        return node

    def visit_ExceptHandler(self, node):
        self.scope_manager.enter_scope()
        if node.name:  # Exception variable
            self.scope_manager.add_to_scope(node.name)
        self.generic_visit(node)
        self.scope_manager.exit_scope()
        return node

    def create_if_call(self, func_name, yield_args, flat_args):
        """Creates the assignment statement for the if function call"""
        if not yield_args:
            return ast.Expr(value=ast.Name(id=func_name, ctx=ast.Load()))
        elif len(yield_args) == 1:
            return ast.Assign(
                targets=[ast.Name(id=yield_args[0], ctx=ast.Store())],
                value=ast.Name(id=func_name, ctx=ast.Load()),
            )
        else:
            return ast.Assign(
                targets=[
                    ast.Tuple(
                        elts=[ast.Name(id=var, ctx=ast.Store()) for var in yield_args],
                        ctx=ast.Store(),
                    )
                ],
                value=ast.Name(id=func_name, ctx=ast.Load()),
            )

    def visit_IfExp(self, node):
        """
        Visits an inline if-else expression (ternary operator).
        This is the Python equivalent of `x if condition else y`.
        """
        # Check if the condition is constexpr
        constexpr_val, test = self.is_constexpr(node)

        node.test = test
        node.body = self.visit(node.body)
        node.orelse = self.visit(node.orelse)

        # If it's a constexpr node, we don't need to transform it
        if constexpr_val.value is True:
            return node

        # Emit
        # node if type(pred) == bool else select_(pred, body, orelse)
        # so if pred is a python bool, use python to short-circuit and avoid emit arith.select
        return ast.copy_location(
            ast.IfExp(
                test=ast.Compare(
                    left=ast.Call(
                        func=ast.Name(id="type", ctx=ast.Load()),
                        args=[node.test],
                        keywords=[],
                    ),
                    ops=[ast.Eq()],
                    comparators=[ast.Name(id="bool", ctx=ast.Load())],
                ),
                body=node,  # Original ternary expression
                orelse=ast.Call(
                    func=ast.Name(id="select_", ctx=ast.Load()),
                    args=[
                        node.test,
                        node.body,
                        node.orelse,
                    ],
                    keywords=[],
                ),
            ),
            node,
        )

    def visit_If(self, node):
        active_symbols = self.scope_manager.get_active_symbols()
        self.scope_manager.enter_scope()

        # Constexpr doesn't get preprocessed
        if self.is_node_constexpr(node):
            self.generic_visit(node)
            self.scope_manager.exit_scope()
            return node

        # Check for early exit and raise exception
        self.check_early_exit(node, "if")

        used_args, yield_args, flat_args = self.analyze_region_variables(
            node, active_symbols
        )
        func_name = f"if_region_{self.counter}"
        self.counter += 1

        func_def = self.create_if_function(
            func_name, node, used_args, yield_args, flat_args
        )
        assign = ast.copy_location(
            self.create_if_call(func_name, yield_args, flat_args), node
        )

        self.scope_manager.exit_scope()
        return [func_def, assign]

    def is_constexpr(self, node):
        """Determines if the if condition is wrapped in const_expr or dynamic_expr"""
        if isinstance(node.test, ast.Call):
            func = node.test.func

            # Check if the function is 'const_expr'
            if isinstance(func, ast.Name) and func.id == "const_expr":
                return ast.Constant(value=True), node.test.args[0]

            # Check if the function is 'dynamic_expr'
            elif isinstance(func, ast.Name) and func.id == "dynamic_expr":
                return ast.Constant(value=False), self.visit(node.test.args[0])

            # Check if it's an attribute access for 'const_expr' or 'dynamic_expr'
            elif isinstance(func, ast.Attribute):
                if func.attr == "const_expr":
                    return ast.Constant(value=True), node.test.args[0]
                elif func.attr == "dynamic_expr":
                    return ast.Constant(value=False), self.visit(node.test.args[0])

        return ast.Constant(value=None), self.visit(node.test)

    def create_if_function(
        self, func_name, node, used_args, yield_args, flattened_args
    ):
        is_constexpr, test_expr = self.is_constexpr(node)
        pred_name = self.make_func_param_name("pred", flattened_args)
        func_args = [ast.arg(arg=pred_name, annotation=None)]
        func_args += [ast.arg(arg=var, annotation=None) for var in flattened_args]
        func_args_then_else = [
            ast.arg(arg=var, annotation=None) for var in flattened_args
        ]

        then_body = []
        for stmt in node.body:
            transformed_stmt = self.visit(stmt)  # Recursively visit inner statements
            if isinstance(transformed_stmt, list):
                then_body.extend(transformed_stmt)
            else:
                then_body.append(transformed_stmt)

        # Create common return list for all blocks
        return_list = ast.List(
            elts=[ast.Name(id=var, ctx=ast.Load()) for var in yield_args],
            ctx=ast.Load(),
        )

        # Create common function arguments
        func_decorator_arguments = ast.arguments(
            posonlyargs=[], args=func_args, kwonlyargs=[], kw_defaults=[], defaults=[]
        )
        func_then_else_arguments = ast.arguments(
            posonlyargs=[],
            args=func_args_then_else,
            kwonlyargs=[],
            kw_defaults=[],
            defaults=[],
        )

        then_block_name = f"then_block_{self.counter}"
        else_block_name = f"else_block_{self.counter}"
        elif_region_name = f"elif_region_{self.counter}"
        self.counter += 1

        # Create then block
        then_block = ast.copy_location(
            ast.FunctionDef(
                name=then_block_name,
                args=func_then_else_arguments,
                body=then_body + [ast.Return(value=return_list)],
                decorator_list=[],
            ),
            node,
        )

        # Decorator keywords
        decorator_keywords = [
            ast.keyword(
                arg="pred", value=test_expr
            ),  # ast.Name(id="pred", ctx=ast.Load())
            ast.keyword(
                arg="used_args",
                value=ast.List(
                    elts=[ast.Name(id=arg, ctx=ast.Load()) for arg in used_args],
                    ctx=ast.Load(),
                ),
            ),
            ast.keyword(
                arg="yield_args",
                value=ast.List(
                    elts=[ast.Name(id=arg, ctx=ast.Load()) for arg in yield_args],
                    ctx=ast.Load(),
                ),
            ),
        ]

        # Create decorator
        decorator = ast.copy_location(
            ast.Call(
                func=ast.Name(id=self.DECORATOR_IF_STATEMENT, ctx=ast.Load()),
                args=[],
                keywords=decorator_keywords,
            ),
            node,
        )

        # Executor keywords
        execute_keywords = [
            ast.keyword(arg="pred", value=ast.Name(id=pred_name, ctx=ast.Load())),
            ast.keyword(
                arg="used_args",
                value=ast.List(
                    elts=[ast.Name(id=arg, ctx=ast.Load()) for arg in used_args],
                    ctx=ast.Load(),
                ),
            ),
            ast.keyword(
                arg="yield_args",
                value=ast.List(
                    elts=[ast.Name(id=arg, ctx=ast.Load()) for arg in yield_args],
                    ctx=ast.Load(),
                ),
            ),
            ast.keyword(
                arg="yield_arg_names",
                value=ast.List(
                    elts=[ast.Constant(value=arg) for arg in yield_args],
                    ctx=ast.Load(),
                ),
            ),
            ast.keyword(
                arg="then_block", value=ast.Name(id=then_block_name, ctx=ast.Load())
            ),
        ]

        # Handle different cases
        if not yield_args and node.orelse == []:
            # No yield_args case - only then_block needed
            execute_call = ast.copy_location(
                ast.Call(
                    func=ast.copy_location(
                        ast.Name(id=self.IF_EXECUTOR, ctx=ast.Load()), node
                    ),
                    args=[],
                    keywords=execute_keywords,
                ),
                node,
            )
            func_body = [then_block, ast.Return(value=execute_call)]
        else:
            # Create else block based on node.orelse
            if node.orelse:
                if len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If):
                    # Handle elif case
                    elif_node = node.orelse[0]
                    nested_if_name = elif_region_name
                    # Recursion for nested elif
                    nested_if = self.create_if_function(
                        nested_if_name, elif_node, used_args, yield_args, flattened_args
                    )
                    else_block = ast.FunctionDef(
                        name=else_block_name,
                        args=func_then_else_arguments,
                        body=[
                            nested_if,
                            ast.Return(
                                value=ast.Name(id=nested_if_name, ctx=ast.Load())
                            ),
                        ],
                        decorator_list=[],
                    )
                else:

                    else_body = []
                    for stmt in node.orelse:
                        transformed_stmt = self.visit(
                            stmt
                        )  # Recursively visit inner statements
                        if isinstance(transformed_stmt, list):
                            else_body.extend(transformed_stmt)
                        else:
                            else_body.append(transformed_stmt)

                    # Regular else block
                    else_block = ast.FunctionDef(
                        name=else_block_name,
                        args=func_then_else_arguments,
                        body=else_body + [ast.Return(value=return_list)],
                        decorator_list=[],
                    )
            else:
                # Default else block
                else_block = ast.FunctionDef(
                    name=else_block_name,
                    args=func_then_else_arguments,
                    body=[ast.Return(value=return_list)],
                    decorator_list=[],
                )

            # Add else_block to execute keywords
            execute_keywords.append(
                ast.keyword(
                    arg="else_block", value=ast.Name(id=else_block_name, ctx=ast.Load())
                )
            )
            # Add constexpr
            execute_keywords.append(ast.keyword(arg="constexpr", value=is_constexpr))

            execute_call = ast.copy_location(
                ast.Call(
                    func=ast.Name(id=self.IF_EXECUTOR, ctx=ast.Load()),
                    args=[],
                    keywords=execute_keywords,
                ),
                node,
            )
            func_body = [
                then_block,
                ast.copy_location(else_block, node),
                ast.Return(value=execute_call),
            ]

        return ast.copy_location(
            ast.FunctionDef(
                name=func_name,
                args=func_decorator_arguments,
                body=func_body,
                decorator_list=[decorator],
            ),
            node,
        )

    def create_while_function(
        self, func_name, node, used_args, yield_args, flattened_args
    ):
        """Create a while function that looks like:

        @while_selector(pred, used_args=[], yield_args=[])
        def while_region(pred, flattened_args):
            def while_before_block(*used_args, *yield_args):
                # Note that during eval of pred can possibly alter yield_args
                return *pred, yield_args
            def while_after_block(*used_args, yield_args):
                ...loop_body_transformed...
                return yield_args
            return self.while_executor(pred, used_args, yield_args,
                while_before_block, while_after_block, constexpr)
        yield_args = while_region(pred, flattened_args)

        Which will later be executed as psuedo-code:

        # Dynamic mode:
        scf.WhileOp(types(yield_args), yield_args)
        with InsertionPoint(before_block):
            cond, yield_args = while_before_block(*flattened_args)
            scf.ConditionOp(cond, yield_args)
        with InsertionPoint(after_block):
            yield_args = while_after_block(yield_args)
            scf.YieldOp(yield_args)
        return while_op.results_

        # Const mode:
        cond, yield_args = while_before_block(yield_args)
        while pred:
            yield_args = body_block(yield_args)
            cond, yield_args = while_before_block(yield_args)
        return yield_args
        """
        is_constexpr, test_expr = self.is_constexpr(node)
        pred_name = self.make_func_param_name("pred", flattened_args)

        # Section: decorator construction
        decorator_keywords = [
            ast.keyword(arg="pred", value=test_expr),
            ast.keyword(
                arg="used_args",
                value=ast.List(
                    elts=[ast.Name(id=arg, ctx=ast.Load()) for arg in used_args],
                    ctx=ast.Load(),
                ),
            ),
            ast.keyword(
                arg="yield_args",
                value=ast.List(
                    elts=[ast.Name(id=arg, ctx=ast.Load()) for arg in yield_args],
                    ctx=ast.Load(),
                ),
            ),
        ]
        decorator = ast.copy_location(
            ast.Call(
                func=ast.Name(id=self.DECORATOR_WHILE_STATEMENT, ctx=ast.Load()),
                args=[],
                keywords=decorator_keywords,
            ),
            node,
        )

        # Section: Shared initialization for before and after blocks
        while_before_block_name = f"while_before_block_{self.counter}"
        while_after_block_name = f"while_after_block_{self.counter}"
        self.counter += 1
        block_args_args = [ast.arg(arg=var, annotation=None) for var in used_args]
        block_args_args += [ast.arg(arg=var, annotation=None) for var in yield_args]
        block_args = ast.arguments(
            posonlyargs=[],
            args=block_args_args,
            kwonlyargs=[],
            kw_defaults=[],
            defaults=[],
        )

        yield_args_ast_name_list = ast.List(
            elts=[ast.Name(id=var, ctx=ast.Load()) for var in yield_args],
            ctx=ast.Load(),
        )

        # Section: while_before_block FunctionDef, which contains condition
        while_before_return_list = ast.List(
            elts=[test_expr, yield_args_ast_name_list],
            ctx=ast.Load(),
        )
        while_before_stmts = [ast.Return(value=while_before_return_list)]
        while_before_block = ast.copy_location(
            ast.FunctionDef(
                name=while_before_block_name,
                args=block_args,
                body=while_before_stmts,
                decorator_list=[],
            ),
            test_expr,
        )

        # Section: while_after_block FunctionDef, which contains loop body
        while_after_stmts = []
        for stmt in node.body:
            transformed_stmt = self.visit(stmt)  # Recursively visit inner statements
            if isinstance(transformed_stmt, list):
                while_after_stmts.extend(transformed_stmt)
            else:
                while_after_stmts.append(transformed_stmt)
        while_after_stmts.append(ast.Return(value=yield_args_ast_name_list))

        while_after_block = ast.copy_location(
            ast.FunctionDef(
                name=while_after_block_name,
                args=block_args,
                body=while_after_stmts,
                decorator_list=[],
            ),
            node,
        )

        # Section: Execute via executor
        execute_keywords = [
            ast.keyword(arg="pred", value=ast.Name(id=pred_name, ctx=ast.Load())),
            ast.keyword(
                arg="used_args",
                value=ast.List(
                    elts=[ast.Name(id=arg, ctx=ast.Load()) for arg in used_args],
                    ctx=ast.Load(),
                ),
            ),
            ast.keyword(
                arg="yield_args",
                value=ast.List(
                    elts=[ast.Name(id=arg, ctx=ast.Load()) for arg in yield_args],
                    ctx=ast.Load(),
                ),
            ),
            ast.keyword(
                arg="while_before_block",
                value=ast.Name(id=while_before_block_name, ctx=ast.Load()),
            ),
            ast.keyword(
                arg="while_after_block",
                value=ast.Name(id=while_after_block_name, ctx=ast.Load()),
            ),
            ast.keyword(arg="constexpr", value=is_constexpr),
            ast.keyword(
                arg="yield_arg_names",
                value=ast.List(
                    elts=[ast.Constant(value=arg) for arg in yield_args],
                    ctx=ast.Load(),
                ),
            ),
        ]

        execute_call = ast.Call(
            func=ast.Name(id=self.WHILE_EXECUTOR, ctx=ast.Load()),
            args=[],
            keywords=execute_keywords,
        )

        # Putting everything together, FunctionDef for while_region
        func_args_args = [ast.arg(arg=pred_name, annotation=None)]
        func_args_args += [ast.arg(arg=var, annotation=None) for var in flattened_args]
        func_args = ast.arguments(
            posonlyargs=[],
            args=func_args_args,
            kwonlyargs=[],
            kw_defaults=[],
            defaults=[],
        )

        return ast.copy_location(
            ast.FunctionDef(
                name=func_name,
                args=func_args,
                body=[
                    while_before_block,
                    while_after_block,
                    ast.Return(value=execute_call),
                ],
                decorator_list=[decorator],
            ),
            node,
        )