from __future__ import annotations

import ast
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

import libcst as cst

import codeflash.code_utils._libcst_cache  # noqa: F401
from codeflash.cli_cmds.console import logger
from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path
from codeflash.code_utils.formatter import sort_imports
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from codeflash.models.models import FunctionParent, TestingMode, VerificationType

if TYPE_CHECKING:
    from collections.abc import Iterable

    from codeflash.models.models import CodePosition


@dataclass(frozen=True)
class FunctionCallNodeArguments:
    args: list[ast.expr]
    keywords: list[ast.keyword]


def get_call_arguments(call_node: ast.Call) -> FunctionCallNodeArguments:
    return FunctionCallNodeArguments(call_node.args, call_node.keywords)


def node_in_call_position(node: ast.AST, call_positions: list[CodePosition]) -> bool:
    # Profile: The most meaningful speedup here is to reduce attribute lookup and to localize call_positions if not empty.
    # Small optimizations for tight loop:
    if isinstance(node, ast.Call):
        node_lineno = getattr(node, "lineno", None)
        node_col_offset = getattr(node, "col_offset", None)
        node_end_lineno = getattr(node, "end_lineno", None)
        node_end_col_offset = getattr(node, "end_col_offset", None)
        if node_lineno is not None and node_col_offset is not None and node_end_lineno is not None:
            # Faster loop: reduce attribute lookups, use local variables for conditionals.
            for pos in call_positions:
                pos_line = pos.line_no
                if pos_line is not None and node_lineno <= pos_line <= node_end_lineno:
                    if pos_line == node_lineno and node_col_offset <= pos.col_no:
                        return True
                    if (
                        pos_line == node_end_lineno
                        and node_end_col_offset is not None
                        and node_end_col_offset >= pos.col_no
                    ):
                        return True
                    if node_lineno < pos_line < node_end_lineno:
                        return True
    return False


def is_argument_name(name: str, arguments_node: ast.arguments) -> bool:
    return any(
        element.arg == name
        for attribute_name in dir(arguments_node)
        if isinstance(attribute := getattr(arguments_node, attribute_name), list)
        for element in attribute
        if isinstance(element, ast.arg)
    )


class InjectPerfOnly(ast.NodeTransformer):
    def __init__(
        self,
        function: FunctionToOptimize,
        module_path: str,
        call_positions: list[CodePosition],
        mode: TestingMode = TestingMode.BEHAVIOR,
    ) -> None:
        self.mode: TestingMode = mode
        self.function_object = function
        self.class_name = None
        self.only_function_name = function.function_name
        self.module_path = module_path
        self.call_positions = call_positions
        if len(function.parents) == 1 and function.parents[0].type == "ClassDef":
            self.class_name = function.top_level_parent_name

    def find_and_update_line_node(
        self, test_node: ast.stmt, node_name: str, index: str, test_class_name: str | None = None
    ) -> Iterable[ast.stmt] | None:
        # Major optimization: since ast.walk is *very* expensive for big trees and only checks for ast.Call,
        # it's much more efficient to visit nodes manually. We'll only descend into expressions/statements.

        # Helper for manual walk
        def iter_ast_calls(node):
            # Generator to yield each ast.Call in test_node, preserves node identity
            stack = [node]
            while stack:
                n = stack.pop()
                if isinstance(n, ast.Call):
                    yield n
                # Instead of using ast.walk (which calls iter_child_nodes under the hood in Python, which copy lists and stack-frames for EVERY node),
                # do a specialized BFS with only the necessary attributes
                for _field, value in ast.iter_fields(n):
                    if isinstance(value, list):
                        for item in reversed(value):
                            if isinstance(item, ast.AST):
                                stack.append(item)
                    elif isinstance(value, ast.AST):
                        stack.append(value)

        # This change improves from O(N) stack-frames per child-node to a single stack, less python call overhead
        return_statement = [test_node]
        call_node = None

        # Minor optimization: Convert mode, function_name, test_class_name, qualified_name, etc to locals
        fn_obj = self.function_object
        module_path = self.module_path
        mode = self.mode
        qualified_name = fn_obj.qualified_name

        # Use locals for all 'current' values, only look up class/function/constant AST object once.
        codeflash_loop_index = ast.Name(id="codeflash_loop_index", ctx=ast.Load())
        codeflash_cur = ast.Name(id="codeflash_cur", ctx=ast.Load())
        codeflash_con = ast.Name(id="codeflash_con", ctx=ast.Load())

        for node in iter_ast_calls(test_node):
            if not node_in_call_position(node, self.call_positions):
                continue

            call_node = node
            all_args = get_call_arguments(call_node)
            # Two possible call types: Name and Attribute
            node_func = node.func

            if isinstance(node_func, ast.Name):
                function_name = node_func.id

                # Check if this is the function we want to instrument
                if function_name != fn_obj.function_name:
                    continue

                if fn_obj.is_async:
                    return [test_node]

                # Build once, reuse objects.
                inspect_name = ast.Name(id="inspect", ctx=ast.Load())
                bind_call = ast.Assign(
                    targets=[ast.Name(id="_call__bound__arguments", ctx=ast.Store())],
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Call(
                                func=ast.Attribute(value=inspect_name, attr="signature", ctx=ast.Load()),
                                args=[ast.Name(id=function_name, ctx=ast.Load())],
                                keywords=[],
                            ),
                            attr="bind",
                            ctx=ast.Load(),
                        ),
                        args=all_args.args,
                        keywords=all_args.keywords,
                    ),
                    lineno=test_node.lineno,
                    col_offset=test_node.col_offset,
                )

                apply_defaults = ast.Expr(
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                            attr="apply_defaults",
                            ctx=ast.Load(),
                        ),
                        args=[],
                        keywords=[],
                    ),
                    lineno=test_node.lineno + 1,
                    col_offset=test_node.col_offset,
                )

                node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
                base_args = [
                    ast.Name(id=function_name, ctx=ast.Load()),
                    ast.Constant(value=module_path),
                    ast.Constant(value=test_class_name or None),
                    ast.Constant(value=node_name),
                    ast.Constant(value=qualified_name),
                    ast.Constant(value=index),
                    codeflash_loop_index,
                ]
                # Extend with BEHAVIOR extras if needed
                if mode == TestingMode.BEHAVIOR:
                    base_args += [codeflash_cur, codeflash_con]
                # Extend with call args (performance) or starred bound args (behavior)
                if mode == TestingMode.PERFORMANCE:
                    base_args += call_node.args
                else:
                    base_args.append(
                        ast.Starred(
                            value=ast.Attribute(
                                value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                attr="args",
                                ctx=ast.Load(),
                            ),
                            ctx=ast.Load(),
                        )
                    )
                node.args = base_args
                # Prepare keywords
                if mode == TestingMode.BEHAVIOR:
                    node.keywords = [
                        ast.keyword(
                            value=ast.Attribute(
                                value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                attr="kwargs",
                                ctx=ast.Load(),
                            )
                        )
                    ]
                else:
                    node.keywords = call_node.keywords

                return_statement = (
                    [bind_call, apply_defaults, test_node] if mode == TestingMode.BEHAVIOR else [test_node]
                )
                break
            if isinstance(node_func, ast.Attribute):
                function_to_test = node_func.attr
                if function_to_test == fn_obj.function_name:
                    if fn_obj.is_async:
                        return [test_node]

                    # Create the signature binding statements

                    # Unparse only once
                    function_name_expr = ast.parse(ast.unparse(node_func), mode="eval").body

                    inspect_name = ast.Name(id="inspect", ctx=ast.Load())
                    bind_call = ast.Assign(
                        targets=[ast.Name(id="_call__bound__arguments", ctx=ast.Store())],
                        value=ast.Call(
                            func=ast.Attribute(
                                value=ast.Call(
                                    func=ast.Attribute(value=inspect_name, attr="signature", ctx=ast.Load()),
                                    args=[function_name_expr],
                                    keywords=[],
                                ),
                                attr="bind",
                                ctx=ast.Load(),
                            ),
                            args=all_args.args,
                            keywords=all_args.keywords,
                        ),
                        lineno=test_node.lineno,
                        col_offset=test_node.col_offset,
                    )

                    apply_defaults = ast.Expr(
                        value=ast.Call(
                            func=ast.Attribute(
                                value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                attr="apply_defaults",
                                ctx=ast.Load(),
                            ),
                            args=[],
                            keywords=[],
                        ),
                        lineno=test_node.lineno + 1,
                        col_offset=test_node.col_offset,
                    )

                    node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
                    base_args = [
                        function_name_expr,
                        ast.Constant(value=module_path),
                        ast.Constant(value=test_class_name or None),
                        ast.Constant(value=node_name),
                        ast.Constant(value=qualified_name),
                        ast.Constant(value=index),
                        codeflash_loop_index,
                    ]
                    if mode == TestingMode.BEHAVIOR:
                        base_args += [codeflash_cur, codeflash_con]
                    if mode == TestingMode.PERFORMANCE:
                        base_args += call_node.args
                    else:
                        base_args.append(
                            ast.Starred(
                                value=ast.Attribute(
                                    value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                    attr="args",
                                    ctx=ast.Load(),
                                ),
                                ctx=ast.Load(),
                            )
                        )
                    node.args = base_args
                    if mode == TestingMode.BEHAVIOR:
                        node.keywords = [
                            ast.keyword(
                                value=ast.Attribute(
                                    value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                    attr="kwargs",
                                    ctx=ast.Load(),
                                )
                            )
                        ]
                    else:
                        node.keywords = call_node.keywords

                    # Return the signature binding statements along with the test_node
                    return_statement = (
                        [bind_call, apply_defaults, test_node] if mode == TestingMode.BEHAVIOR else [test_node]
                    )
                    break

        if call_node is None:
            return None
        return return_statement

    def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
        # TODO: Ensure that this class inherits from unittest.TestCase. Don't modify non unittest.TestCase classes.
        for inner_node in ast.walk(node):
            if isinstance(inner_node, ast.FunctionDef):
                self.visit_FunctionDef(inner_node, node.name)

        return node

    def visit_FunctionDef(self, node: ast.FunctionDef, test_class_name: str | None = None) -> ast.FunctionDef:
        if node.name.startswith("test_"):
            did_update = False
            i = len(node.body) - 1
            while i >= 0:
                line_node = node.body[i]
                # TODO: Validate if the functional call actually did not raise any exceptions

                if isinstance(line_node, (ast.With, ast.For, ast.While, ast.If)):
                    j = len(line_node.body) - 1
                    while j >= 0:
                        compound_line_node: ast.stmt = line_node.body[j]
                        internal_node: ast.AST
                        for internal_node in ast.walk(compound_line_node):
                            if isinstance(internal_node, (ast.stmt, ast.Assign)):
                                updated_node = self.find_and_update_line_node(
                                    internal_node, node.name, str(i) + "_" + str(j), test_class_name
                                )
                                if updated_node is not None:
                                    line_node.body[j : j + 1] = updated_node
                                    did_update = True
                                    break
                        j -= 1
                else:
                    updated_node = self.find_and_update_line_node(line_node, node.name, str(i), test_class_name)
                    if updated_node is not None:
                        node.body[i : i + 1] = updated_node
                        did_update = True
                i -= 1
            if did_update:
                node.body = [
                    ast.Assign(
                        targets=[ast.Name(id="codeflash_loop_index", ctx=ast.Store())],
                        value=ast.Call(
                            func=ast.Name(id="int", ctx=ast.Load()),
                            args=[
                                ast.Subscript(
                                    value=ast.Attribute(
                                        value=ast.Name(id="os", ctx=ast.Load()), attr="environ", ctx=ast.Load()
                                    ),
                                    slice=ast.Constant(value="CODEFLASH_LOOP_INDEX"),
                                    ctx=ast.Load(),
                                )
                            ],
                            keywords=[],
                        ),
                        lineno=node.lineno + 2,
                        col_offset=node.col_offset,
                    ),
                    *(
                        [
                            ast.Assign(
                                targets=[ast.Name(id="codeflash_iteration", ctx=ast.Store())],
                                value=ast.Subscript(
                                    value=ast.Attribute(
                                        value=ast.Name(id="os", ctx=ast.Load()), attr="environ", ctx=ast.Load()
                                    ),
                                    slice=ast.Constant(value="CODEFLASH_TEST_ITERATION"),
                                    ctx=ast.Load(),
                                ),
                                lineno=node.lineno + 1,
                                col_offset=node.col_offset,
                            ),
                            ast.Assign(
                                targets=[ast.Name(id="codeflash_con", ctx=ast.Store())],
                                value=ast.Call(
                                    func=ast.Attribute(
                                        value=ast.Name(id="sqlite3", ctx=ast.Load()), attr="connect", ctx=ast.Load()
                                    ),
                                    args=[
                                        ast.JoinedStr(
                                            values=[
                                                ast.Constant(
                                                    value=f"{get_run_tmp_file(Path('test_return_values_')).as_posix()}"
                                                ),
                                                ast.FormattedValue(
                                                    value=ast.Name(id="codeflash_iteration", ctx=ast.Load()),
                                                    conversion=-1,
                                                ),
                                                ast.Constant(value=".sqlite"),
                                            ]
                                        )
                                    ],
                                    keywords=[],
                                ),
                                lineno=node.lineno + 3,
                                col_offset=node.col_offset,
                            ),
                            ast.Assign(
                                targets=[ast.Name(id="codeflash_cur", ctx=ast.Store())],
                                value=ast.Call(
                                    func=ast.Attribute(
                                        value=ast.Name(id="codeflash_con", ctx=ast.Load()),
                                        attr="cursor",
                                        ctx=ast.Load(),
                                    ),
                                    args=[],
                                    keywords=[],
                                ),
                                lineno=node.lineno + 4,
                                col_offset=node.col_offset,
                            ),
                            ast.Expr(
                                value=ast.Call(
                                    func=ast.Attribute(
                                        value=ast.Name(id="codeflash_cur", ctx=ast.Load()),
                                        attr="execute",
                                        ctx=ast.Load(),
                                    ),
                                    args=[
                                        ast.Constant(
                                            value="CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT,"
                                            " test_class_name TEXT, test_function_name TEXT, function_getting_tested TEXT,"
                                            " loop_index INTEGER, iteration_id TEXT, runtime INTEGER, return_value BLOB, verification_type TEXT)"
                                        )
                                    ],
                                    keywords=[],
                                ),
                                lineno=node.lineno + 5,
                                col_offset=node.col_offset,
                            ),
                        ]
                        if self.mode == TestingMode.BEHAVIOR
                        else []
                    ),
                    *node.body,
                    *(
                        [
                            ast.Expr(
                                value=ast.Call(
                                    func=ast.Attribute(
                                        value=ast.Name(id="codeflash_con", ctx=ast.Load()), attr="close", ctx=ast.Load()
                                    ),
                                    args=[],
                                    keywords=[],
                                )
                            )
                        ]
                        if self.mode == TestingMode.BEHAVIOR
                        else []
                    ),
                ]
        return node


class AsyncCallInstrumenter(ast.NodeTransformer):
    def __init__(
        self,
        function: FunctionToOptimize,
        module_path: str,
        call_positions: list[CodePosition],
        mode: TestingMode = TestingMode.BEHAVIOR,
    ) -> None:
        self.mode = mode
        self.function_object = function
        self.class_name = None
        self.only_function_name = function.function_name
        self.module_path = module_path
        self.call_positions = call_positions
        self.did_instrument = False
        # Track function call count per test function
        self.async_call_counter: dict[str, int] = {}
        if len(function.parents) == 1 and function.parents[0].type == "ClassDef":
            self.class_name = function.top_level_parent_name

    def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
        return self.generic_visit(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef:
        if not node.name.startswith("test_"):
            return node

        return self._process_test_function(node)

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
        # Only process test functions
        if not node.name.startswith("test_"):
            return node

        return self._process_test_function(node)

    def _process_test_function(
        self, node: ast.AsyncFunctionDef | ast.FunctionDef
    ) -> ast.AsyncFunctionDef | ast.FunctionDef:
        # Initialize counter for this test function
        if node.name not in self.async_call_counter:
            self.async_call_counter[node.name] = 0

        new_body = []

        # Optimize ast.walk calls inside _instrument_statement, by scanning only relevant nodes
        for _i, stmt in enumerate(node.body):
            transformed_stmt, added_env_assignment = self._optimized_instrument_statement(stmt)

            if added_env_assignment:
                current_call_index = self.async_call_counter[node.name]
                self.async_call_counter[node.name] += 1

                env_assignment = ast.Assign(
                    targets=[
                        ast.Subscript(
                            value=ast.Attribute(
                                value=ast.Name(id="os", ctx=ast.Load()), attr="environ", ctx=ast.Load()
                            ),
                            slice=ast.Constant(value="CODEFLASH_CURRENT_LINE_ID"),
                            ctx=ast.Store(),
                        )
                    ],
                    value=ast.Constant(value=f"{current_call_index}"),
                    lineno=stmt.lineno if hasattr(stmt, "lineno") else 1,
                )
                new_body.append(env_assignment)
                self.did_instrument = True

            new_body.append(transformed_stmt)

        node.body = new_body
        return node

    def _instrument_statement(self, stmt: ast.stmt, _node_name: str) -> tuple[ast.stmt, bool]:
        for node in ast.walk(stmt):
            if (
                isinstance(node, ast.Await)
                and isinstance(node.value, ast.Call)
                and self._is_target_call(node.value)
                and self._call_in_positions(node.value)
            ):
                # Check if this call is in one of our target positions
                return stmt, True  # Return original statement but signal we added env var

        return stmt, False

    def _is_target_call(self, call_node: ast.Call) -> bool:
        """Check if this call node is calling our target async function."""
        if isinstance(call_node.func, ast.Name):
            return call_node.func.id == self.function_object.function_name
        if isinstance(call_node.func, ast.Attribute):
            return call_node.func.attr == self.function_object.function_name
        return False

    def _call_in_positions(self, call_node: ast.Call) -> bool:
        if not hasattr(call_node, "lineno") or not hasattr(call_node, "col_offset"):
            return False

        return node_in_call_position(call_node, self.call_positions)

    # Optimized version: only walk child nodes for Await
    def _optimized_instrument_statement(self, stmt: ast.stmt) -> tuple[ast.stmt, bool]:
        # Stack-based DFS, manual for relevant Await nodes
        stack = [stmt]
        while stack:
            node = stack.pop()
            # Favor direct ast.Await detection
            if isinstance(node, ast.Await):
                val = node.value
                if isinstance(val, ast.Call) and self._is_target_call(val) and self._call_in_positions(val):
                    return stmt, True
            # Use _fields instead of ast.walk for less allocations
            for fname in getattr(node, "_fields", ()):
                child = getattr(node, fname, None)
                if isinstance(child, list):
                    stack.extend(child)
                elif isinstance(child, ast.AST):
                    stack.append(child)
        return stmt, False


class FunctionImportedAsVisitor(ast.NodeVisitor):
    """Checks if a function has been imported as an alias. We only care about the alias then.

    from numpy import array as np_array
    np_array is what we want
    """

    def __init__(self, function: FunctionToOptimize) -> None:
        assert len(function.parents) <= 1, "Only support functions with one or less parent"
        self.imported_as = function
        self.function = function
        if function.parents:
            self.to_match = function.parents[0].name
        else:
            self.to_match = function.function_name

    # TODO: Validate if the function imported is actually from the right module
    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        for alias in node.names:
            if alias.name == self.to_match and hasattr(alias, "asname") and alias.asname is not None:
                if self.function.parents:
                    self.imported_as = FunctionToOptimize(
                        function_name=self.function.function_name,
                        parents=[FunctionParent(alias.asname, "ClassDef")],
                        file_path=self.function.file_path,
                        starting_line=self.function.starting_line,
                        ending_line=self.function.ending_line,
                        is_async=self.function.is_async,
                    )
                else:
                    self.imported_as = FunctionToOptimize(
                        function_name=alias.asname,
                        parents=[],
                        file_path=self.function.file_path,
                        starting_line=self.function.starting_line,
                        ending_line=self.function.ending_line,
                        is_async=self.function.is_async,
                    )


def inject_async_profiling_into_existing_test(
    test_path: Path,
    call_positions: list[CodePosition],
    function_to_optimize: FunctionToOptimize,
    tests_project_root: Path,
    mode: TestingMode = TestingMode.BEHAVIOR,
) -> tuple[bool, str | None]:
    """Inject profiling for async function calls by setting environment variables before each call."""
    with test_path.open(encoding="utf8") as f:
        test_code = f.read()

    try:
        tree = ast.parse(test_code)
    except SyntaxError:
        logger.exception(f"Syntax error in code in file - {test_path}")
        return False, None
    # TODO: Pass the full name of function here, otherwise we can run into namespace clashes
    test_module_path = module_name_from_file_path(test_path, tests_project_root)
    import_visitor = FunctionImportedAsVisitor(function_to_optimize)
    import_visitor.visit(tree)
    func = import_visitor.imported_as

    async_instrumenter = AsyncCallInstrumenter(func, test_module_path, call_positions, mode=mode)
    tree = async_instrumenter.visit(tree)

    if not async_instrumenter.did_instrument:
        return False, None

    # Add necessary imports
    new_imports = [ast.Import(names=[ast.alias(name="os")])]

    tree.body = [*new_imports, *tree.body]
    return True, sort_imports(ast.unparse(tree), float_to_top=True)


def detect_frameworks_from_code(code: str) -> dict[str, str]:
    """Detect GPU/device frameworks (torch, tensorflow, jax) used in the code by analyzing imports.

    Returns:
        A dictionary mapping framework names to their import aliases.
        For example: {"torch": "th", "tensorflow": "tf", "jax": "jax"}

    """
    frameworks: dict[str, str] = {}
    try:
        tree = ast.parse(code)
    except SyntaxError:
        return frameworks

    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                module_name = alias.name.split(".")[0]
                if module_name == "torch":
                    # Use asname if available, otherwise use the module name
                    frameworks["torch"] = alias.asname if alias.asname else module_name
                elif module_name == "tensorflow":
                    frameworks["tensorflow"] = alias.asname if alias.asname else module_name
                elif module_name == "jax":
                    frameworks["jax"] = alias.asname if alias.asname else module_name
        elif isinstance(node, ast.ImportFrom) and node.module:
            module_name = node.module.split(".")[0]
            if module_name == "torch" and "torch" not in frameworks:
                frameworks["torch"] = module_name
            elif module_name == "tensorflow" and "tensorflow" not in frameworks:
                frameworks["tensorflow"] = module_name
            elif module_name == "jax" and "jax" not in frameworks:
                frameworks["jax"] = module_name

    return frameworks


def inject_profiling_into_existing_test(
    test_path: Path,
    call_positions: list[CodePosition],
    function_to_optimize: FunctionToOptimize,
    tests_project_root: Path,
    mode: TestingMode = TestingMode.BEHAVIOR,
) -> tuple[bool, str | None]:
    tests_project_root = tests_project_root.resolve()
    if function_to_optimize.is_async:
        return inject_async_profiling_into_existing_test(
            test_path, call_positions, function_to_optimize, tests_project_root, mode
        )

    with test_path.open(encoding="utf8") as f:
        test_code = f.read()

    used_frameworks = detect_frameworks_from_code(test_code)
    try:
        tree = ast.parse(test_code)
    except SyntaxError:
        logger.exception(f"Syntax error in code in file - {test_path}")
        return False, None

    test_module_path = module_name_from_file_path(test_path, tests_project_root)
    import_visitor = FunctionImportedAsVisitor(function_to_optimize)
    import_visitor.visit(tree)
    func = import_visitor.imported_as

    tree = InjectPerfOnly(func, test_module_path, call_positions, mode=mode).visit(tree)
    new_imports = [
        ast.Import(names=[ast.alias(name="time")]),
        ast.Import(names=[ast.alias(name="gc")]),
        ast.Import(names=[ast.alias(name="os")]),
    ]
    if mode == TestingMode.BEHAVIOR:
        new_imports.extend(
            [
                ast.Import(names=[ast.alias(name="inspect")]),
                ast.Import(names=[ast.alias(name="sqlite3")]),
                ast.Import(names=[ast.alias(name="dill", asname="pickle")]),
            ]
        )
    # Add framework imports for GPU sync code (needed when framework is only imported via submodule)
    for framework_name, framework_alias in used_frameworks.items():
        if framework_alias == framework_name:
            # Only add import if we're using the framework name directly (not an alias)
            # This handles cases like "from torch.nn import Module" where torch needs to be imported
            new_imports.append(ast.Import(names=[ast.alias(name=framework_name)]))
        else:
            # If there's an alias, use it (e.g., "import torch as th")
            new_imports.append(ast.Import(names=[ast.alias(name=framework_name, asname=framework_alias)]))
    additional_functions = [create_wrapper_function(mode, used_frameworks)]

    tree.body = [*new_imports, *additional_functions, *tree.body]
    return True, sort_imports(ast.unparse(tree), float_to_top=True)


def _create_device_sync_precompute_statements(used_frameworks: dict[str, str] | None) -> list[ast.stmt]:
    """Create AST statements to pre-compute device sync conditions before profiling.

    This moves the conditional checks (like is_available(), hasattr(), etc.) outside
    the timing block to avoid their overhead affecting the measurements.

    Args:
        used_frameworks: Dict mapping framework names to their import aliases

    Returns:
        List of AST statements that pre-compute sync conditions into boolean variables

    """
    if not used_frameworks:
        return []

    precompute_statements: list[ast.stmt] = []

    # PyTorch: pre-compute whether to sync CUDA or MPS
    if "torch" in used_frameworks:
        torch_alias = used_frameworks["torch"]
        # _codeflash_should_sync_cuda = torch.cuda.is_available() and torch.cuda.is_initialized()
        precompute_statements.append(
            ast.Assign(
                targets=[ast.Name(id="_codeflash_should_sync_cuda", ctx=ast.Store())],
                value=ast.BoolOp(
                    op=ast.And(),
                    values=[
                        ast.Call(
                            func=ast.Attribute(
                                value=ast.Attribute(
                                    value=ast.Name(id=torch_alias, ctx=ast.Load()), attr="cuda", ctx=ast.Load()
                                ),
                                attr="is_available",
                                ctx=ast.Load(),
                            ),
                            args=[],
                            keywords=[],
                        ),
                        ast.Call(
                            func=ast.Attribute(
                                value=ast.Attribute(
                                    value=ast.Name(id=torch_alias, ctx=ast.Load()), attr="cuda", ctx=ast.Load()
                                ),
                                attr="is_initialized",
                                ctx=ast.Load(),
                            ),
                            args=[],
                            keywords=[],
                        ),
                    ],
                ),
                lineno=1,
            )
        )
        # _codeflash_should_sync_mps = (not _codeflash_should_sync_cuda and
        #     hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() and
        #     hasattr(torch.mps, 'synchronize'))
        precompute_statements.append(
            ast.Assign(
                targets=[ast.Name(id="_codeflash_should_sync_mps", ctx=ast.Store())],
                value=ast.BoolOp(
                    op=ast.And(),
                    values=[
                        ast.UnaryOp(op=ast.Not(), operand=ast.Name(id="_codeflash_should_sync_cuda", ctx=ast.Load())),
                        ast.Call(
                            func=ast.Name(id="hasattr", ctx=ast.Load()),
                            args=[
                                ast.Attribute(
                                    value=ast.Name(id=torch_alias, ctx=ast.Load()), attr="backends", ctx=ast.Load()
                                ),
                                ast.Constant(value="mps"),
                            ],
                            keywords=[],
                        ),
                        ast.Call(
                            func=ast.Attribute(
                                value=ast.Attribute(
                                    value=ast.Attribute(
                                        value=ast.Name(id=torch_alias, ctx=ast.Load()), attr="backends", ctx=ast.Load()
                                    ),
                                    attr="mps",
                                    ctx=ast.Load(),
                                ),
                                attr="is_available",
                                ctx=ast.Load(),
                            ),
                            args=[],
                            keywords=[],
                        ),
                        ast.Call(
                            func=ast.Name(id="hasattr", ctx=ast.Load()),
                            args=[
                                ast.Attribute(
                                    value=ast.Name(id=torch_alias, ctx=ast.Load()), attr="mps", ctx=ast.Load()
                                ),
                                ast.Constant(value="synchronize"),
                            ],
                            keywords=[],
                        ),
                    ],
                ),
                lineno=1,
            )
        )

    # JAX: pre-compute whether jax.block_until_ready exists
    if "jax" in used_frameworks:
        jax_alias = used_frameworks["jax"]
        # _codeflash_should_sync_jax = hasattr(jax, 'block_until_ready')
        precompute_statements.append(
            ast.Assign(
                targets=[ast.Name(id="_codeflash_should_sync_jax", ctx=ast.Store())],
                value=ast.Call(
                    func=ast.Name(id="hasattr", ctx=ast.Load()),
                    args=[ast.Name(id=jax_alias, ctx=ast.Load()), ast.Constant(value="block_until_ready")],
                    keywords=[],
                ),
                lineno=1,
            )
        )

    # TensorFlow: pre-compute whether tf.test.experimental.sync_devices exists
    if "tensorflow" in used_frameworks:
        tf_alias = used_frameworks["tensorflow"]
        # _codeflash_should_sync_tf = hasattr(tf.test.experimental, 'sync_devices')
        precompute_statements.append(
            ast.Assign(
                targets=[ast.Name(id="_codeflash_should_sync_tf", ctx=ast.Store())],
                value=ast.Call(
                    func=ast.Name(id="hasattr", ctx=ast.Load()),
                    args=[
                        ast.Attribute(
                            value=ast.Attribute(
                                value=ast.Name(id=tf_alias, ctx=ast.Load()), attr="test", ctx=ast.Load()
                            ),
                            attr="experimental",
                            ctx=ast.Load(),
                        ),
                        ast.Constant(value="sync_devices"),
                    ],
                    keywords=[],
                ),
                lineno=1,
            )
        )

    return precompute_statements


def _create_device_sync_statements(
    used_frameworks: dict[str, str] | None, for_return_value: bool = False
) -> list[ast.stmt]:
    """Create AST statements for device synchronization using pre-computed conditions.

    Args:
        used_frameworks: Dict mapping framework names to their import aliases
                        (e.g., {'torch': 'th', 'tensorflow': 'tf', 'jax': 'jax'})
        for_return_value: If True, creates sync for after function call (includes JAX block_until_ready)

    Returns:
        List of AST statements for device synchronization using pre-computed boolean variables

    """
    if not used_frameworks:
        return []

    sync_statements: list[ast.stmt] = []

    # PyTorch synchronization using pre-computed conditions
    if "torch" in used_frameworks:
        torch_alias = used_frameworks["torch"]
        # if _codeflash_should_sync_cuda:
        #     torch.cuda.synchronize()
        # elif _codeflash_should_sync_mps:
        #     torch.mps.synchronize()
        cuda_sync = ast.If(
            test=ast.Name(id="_codeflash_should_sync_cuda", ctx=ast.Load()),
            body=[
                ast.Expr(
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Attribute(
                                value=ast.Name(id=torch_alias, ctx=ast.Load()), attr="cuda", ctx=ast.Load()
                            ),
                            attr="synchronize",
                            ctx=ast.Load(),
                        ),
                        args=[],
                        keywords=[],
                    )
                )
            ],
            orelse=[
                ast.If(
                    test=ast.Name(id="_codeflash_should_sync_mps", ctx=ast.Load()),
                    body=[
                        ast.Expr(
                            value=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Attribute(
                                        value=ast.Name(id=torch_alias, ctx=ast.Load()), attr="mps", ctx=ast.Load()
                                    ),
                                    attr="synchronize",
                                    ctx=ast.Load(),
                                ),
                                args=[],
                                keywords=[],
                            )
                        )
                    ],
                    orelse=[],
                )
            ],
        )
        sync_statements.append(cuda_sync)

    # JAX synchronization (only after function call, using block_until_ready on return value)
    if "jax" in used_frameworks and for_return_value:
        jax_alias = used_frameworks["jax"]
        # if _codeflash_should_sync_jax:
        #     jax.block_until_ready(return_value)
        jax_sync = ast.If(
            test=ast.Name(id="_codeflash_should_sync_jax", ctx=ast.Load()),
            body=[
                ast.Expr(
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id=jax_alias, ctx=ast.Load()), attr="block_until_ready", ctx=ast.Load()
                        ),
                        args=[ast.Name(id="return_value", ctx=ast.Load())],
                        keywords=[],
                    )
                )
            ],
            orelse=[],
        )
        sync_statements.append(jax_sync)

    # TensorFlow synchronization using pre-computed condition
    if "tensorflow" in used_frameworks:
        tf_alias = used_frameworks["tensorflow"]
        # if _codeflash_should_sync_tf:
        #     tf.test.experimental.sync_devices()
        tf_sync = ast.If(
            test=ast.Name(id="_codeflash_should_sync_tf", ctx=ast.Load()),
            body=[
                ast.Expr(
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Attribute(
                                value=ast.Attribute(
                                    value=ast.Name(id=tf_alias, ctx=ast.Load()), attr="test", ctx=ast.Load()
                                ),
                                attr="experimental",
                                ctx=ast.Load(),
                            ),
                            attr="sync_devices",
                            ctx=ast.Load(),
                        ),
                        args=[],
                        keywords=[],
                    )
                )
            ],
            orelse=[],
        )
        sync_statements.append(tf_sync)

    return sync_statements


def create_wrapper_function(
    mode: TestingMode = TestingMode.BEHAVIOR, used_frameworks: dict[str, str] | None = None
) -> ast.FunctionDef:
    lineno = 1
    wrapper_body: list[ast.stmt] = [
        ast.Assign(
            targets=[ast.Name(id="test_id", ctx=ast.Store())],
            value=ast.JoinedStr(
                values=[
                    ast.FormattedValue(value=ast.Name(id="codeflash_test_module_name", ctx=ast.Load()), conversion=-1),
                    ast.Constant(value=":"),
                    ast.FormattedValue(value=ast.Name(id="codeflash_test_class_name", ctx=ast.Load()), conversion=-1),
                    ast.Constant(value=":"),
                    ast.FormattedValue(value=ast.Name(id="codeflash_test_name", ctx=ast.Load()), conversion=-1),
                    ast.Constant(value=":"),
                    ast.FormattedValue(value=ast.Name(id="codeflash_line_id", ctx=ast.Load()), conversion=-1),
                    ast.Constant(value=":"),
                    ast.FormattedValue(value=ast.Name(id="codeflash_loop_index", ctx=ast.Load()), conversion=-1),
                ]
            ),
            lineno=lineno + 1,
        ),
        ast.If(
            test=ast.UnaryOp(
                op=ast.Not(),
                operand=ast.Call(
                    func=ast.Name(id="hasattr", ctx=ast.Load()),
                    args=[ast.Name(id="codeflash_wrap", ctx=ast.Load()), ast.Constant(value="index")],
                    keywords=[],
                ),
            ),
            body=[
                ast.Assign(
                    targets=[
                        ast.Attribute(
                            value=ast.Name(id="codeflash_wrap", ctx=ast.Load()), attr="index", ctx=ast.Store()
                        )
                    ],
                    value=ast.Dict(keys=[], values=[]),
                    lineno=lineno + 3,
                )
            ],
            orelse=[],
            lineno=lineno + 2,
        ),
        ast.If(
            test=ast.Compare(
                left=ast.Name(id="test_id", ctx=ast.Load()),
                ops=[ast.In()],
                comparators=[
                    ast.Attribute(value=ast.Name(id="codeflash_wrap", ctx=ast.Load()), attr="index", ctx=ast.Load())
                ],
            ),
            body=[
                ast.AugAssign(
                    target=ast.Subscript(
                        value=ast.Attribute(
                            value=ast.Name(id="codeflash_wrap", ctx=ast.Load()), attr="index", ctx=ast.Load()
                        ),
                        slice=ast.Name(id="test_id", ctx=ast.Load()),
                        ctx=ast.Store(),
                    ),
                    op=ast.Add(),
                    value=ast.Constant(value=1),
                    lineno=lineno + 5,
                )
            ],
            orelse=[
                ast.Assign(
                    targets=[
                        ast.Subscript(
                            value=ast.Attribute(
                                value=ast.Name(id="codeflash_wrap", ctx=ast.Load()), attr="index", ctx=ast.Load()
                            ),
                            slice=ast.Name(id="test_id", ctx=ast.Load()),
                            ctx=ast.Store(),
                        )
                    ],
                    value=ast.Constant(value=0),
                    lineno=lineno + 6,
                )
            ],
            lineno=lineno + 4,
        ),
        ast.Assign(
            targets=[ast.Name(id="codeflash_test_index", ctx=ast.Store())],
            value=ast.Subscript(
                value=ast.Attribute(value=ast.Name(id="codeflash_wrap", ctx=ast.Load()), attr="index", ctx=ast.Load()),
                slice=ast.Name(id="test_id", ctx=ast.Load()),
                ctx=ast.Load(),
            ),
            lineno=lineno + 7,
        ),
        ast.Assign(
            targets=[ast.Name(id="invocation_id", ctx=ast.Store())],
            value=ast.JoinedStr(
                values=[
                    ast.FormattedValue(value=ast.Name(id="codeflash_line_id", ctx=ast.Load()), conversion=-1),
                    ast.Constant(value="_"),
                    ast.FormattedValue(value=ast.Name(id="codeflash_test_index", ctx=ast.Load()), conversion=-1),
                ]
            ),
            lineno=lineno + 8,
        ),
        *(
            [
                ast.Assign(
                    targets=[ast.Name(id="test_stdout_tag", ctx=ast.Store())],
                    value=ast.JoinedStr(
                        values=[
                            ast.FormattedValue(
                                value=ast.Name(id="codeflash_test_module_name", ctx=ast.Load()), conversion=-1
                            ),
                            ast.Constant(value=":"),
                            ast.FormattedValue(
                                value=ast.IfExp(
                                    test=ast.Name(id="codeflash_test_class_name", ctx=ast.Load()),
                                    body=ast.BinOp(
                                        left=ast.Name(id="codeflash_test_class_name", ctx=ast.Load()),
                                        op=ast.Add(),
                                        right=ast.Constant(value="."),
                                    ),
                                    orelse=ast.Constant(value=""),
                                ),
                                conversion=-1,
                            ),
                            ast.FormattedValue(value=ast.Name(id="codeflash_test_name", ctx=ast.Load()), conversion=-1),
                            ast.Constant(value=":"),
                            ast.FormattedValue(
                                value=ast.Name(id="codeflash_function_name", ctx=ast.Load()), conversion=-1
                            ),
                            ast.Constant(value=":"),
                            ast.FormattedValue(
                                value=ast.Name(id="codeflash_loop_index", ctx=ast.Load()), conversion=-1
                            ),
                            ast.Constant(value=":"),
                            ast.FormattedValue(value=ast.Name(id="invocation_id", ctx=ast.Load()), conversion=-1),
                        ]
                    ),
                    lineno=lineno + 9,
                ),
                ast.Expr(
                    value=ast.Call(
                        func=ast.Name(id="print", ctx=ast.Load()),
                        args=[
                            ast.JoinedStr(
                                values=[
                                    ast.Constant(value="!$######"),
                                    ast.FormattedValue(
                                        value=ast.Name(id="test_stdout_tag", ctx=ast.Load()), conversion=-1
                                    ),
                                    ast.Constant(value="######$!"),
                                ]
                            )
                        ],
                        keywords=[],
                    )
                ),
            ]
        ),
        ast.Assign(
            targets=[ast.Name(id="exception", ctx=ast.Store())], value=ast.Constant(value=None), lineno=lineno + 10
        ),
        # Pre-compute device sync conditions before profiling to avoid overhead during timing
        *_create_device_sync_precompute_statements(used_frameworks),
        ast.Expr(
            value=ast.Call(
                func=ast.Attribute(value=ast.Name(id="gc", ctx=ast.Load()), attr="disable", ctx=ast.Load()),
                args=[],
                keywords=[],
            ),
            lineno=lineno + 9,
        ),
        ast.Try(
            body=[
                # Pre-sync: synchronize device before starting timer
                *_create_device_sync_statements(used_frameworks, for_return_value=False),
                ast.Assign(
                    targets=[ast.Name(id="counter", ctx=ast.Store())],
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id="time", ctx=ast.Load()), attr="perf_counter_ns", ctx=ast.Load()
                        ),
                        args=[],
                        keywords=[],
                    ),
                    lineno=lineno + 11,
                ),
                ast.Assign(
                    targets=[ast.Name(id="return_value", ctx=ast.Store())],
                    value=ast.Call(
                        func=ast.Name(id="codeflash_wrapped", ctx=ast.Load()),
                        args=[ast.Starred(value=ast.Name(id="args", ctx=ast.Load()), ctx=ast.Load())],
                        keywords=[ast.keyword(arg=None, value=ast.Name(id="kwargs", ctx=ast.Load()))],
                    ),
                    lineno=lineno + 12,
                ),
                # Post-sync: synchronize device after function call to ensure all device work is complete
                *_create_device_sync_statements(used_frameworks, for_return_value=True),
                ast.Assign(
                    targets=[ast.Name(id="codeflash_duration", ctx=ast.Store())],
                    value=ast.BinOp(
                        left=ast.Call(
                            func=ast.Attribute(
                                value=ast.Name(id="time", ctx=ast.Load()), attr="perf_counter_ns", ctx=ast.Load()
                            ),
                            args=[],
                            keywords=[],
                        ),
                        op=ast.Sub(),
                        right=ast.Name(id="counter", ctx=ast.Load()),
                    ),
                    lineno=lineno + 13,
                ),
            ],
            handlers=[
                ast.ExceptHandler(
                    type=ast.Name(id="Exception", ctx=ast.Load()),
                    name="e",
                    body=[
                        ast.Assign(
                            targets=[ast.Name(id="codeflash_duration", ctx=ast.Store())],
                            value=ast.BinOp(
                                left=ast.Call(
                                    func=ast.Attribute(
                                        value=ast.Name(id="time", ctx=ast.Load()),
                                        attr="perf_counter_ns",
                                        ctx=ast.Load(),
                                    ),
                                    args=[],
                                    keywords=[],
                                ),
                                op=ast.Sub(),
                                right=ast.Name(id="counter", ctx=ast.Load()),
                            ),
                            lineno=lineno + 15,
                        ),
                        ast.Assign(
                            targets=[ast.Name(id="exception", ctx=ast.Store())],
                            value=ast.Name(id="e", ctx=ast.Load()),
                            lineno=lineno + 13,
                        ),
                    ],
                    lineno=lineno + 14,
                )
            ],
            orelse=[],
            finalbody=[],
            lineno=lineno + 11,
        ),
        ast.Expr(
            value=ast.Call(
                func=ast.Attribute(value=ast.Name(id="gc", ctx=ast.Load()), attr="enable", ctx=ast.Load()),
                args=[],
                keywords=[],
            )
        ),
        ast.Expr(
            value=ast.Call(
                func=ast.Name(id="print", ctx=ast.Load()),
                args=[
                    ast.JoinedStr(
                        values=[
                            ast.Constant(value="!######"),
                            ast.FormattedValue(value=ast.Name(id="test_stdout_tag", ctx=ast.Load()), conversion=-1),
                            *(
                                [
                                    ast.Constant(value=":"),
                                    ast.FormattedValue(
                                        value=ast.Name(id="codeflash_duration", ctx=ast.Load()), conversion=-1
                                    ),
                                ]
                                if mode == TestingMode.PERFORMANCE
                                else []
                            ),
                            ast.Constant(value="######!"),
                        ]
                    )
                ],
                keywords=[],
            )
        ),
        *(
            [
                ast.Assign(
                    targets=[ast.Name(id="pickled_return_value", ctx=ast.Store())],
                    value=ast.IfExp(
                        test=ast.Name(id="exception", ctx=ast.Load()),
                        body=ast.Call(
                            func=ast.Attribute(
                                value=ast.Name(id="pickle", ctx=ast.Load()), attr="dumps", ctx=ast.Load()
                            ),
                            args=[ast.Name(id="exception", ctx=ast.Load())],
                            keywords=[],
                        ),
                        orelse=ast.Call(
                            func=ast.Attribute(
                                value=ast.Name(id="pickle", ctx=ast.Load()), attr="dumps", ctx=ast.Load()
                            ),
                            args=[ast.Name(id="return_value", ctx=ast.Load())],
                            keywords=[],
                        ),
                    ),
                    lineno=lineno + 18,
                )
            ]
            if mode == TestingMode.BEHAVIOR
            else []
        ),
        *(
            [
                ast.Expr(
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id="codeflash_cur", ctx=ast.Load()), attr="execute", ctx=ast.Load()
                        ),
                        args=[
                            ast.Constant(value="INSERT INTO test_results VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)"),
                            ast.Tuple(
                                elts=[
                                    ast.Name(id="codeflash_test_module_name", ctx=ast.Load()),
                                    ast.Name(id="codeflash_test_class_name", ctx=ast.Load()),
                                    ast.Name(id="codeflash_test_name", ctx=ast.Load()),
                                    ast.Name(id="codeflash_function_name", ctx=ast.Load()),
                                    ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
                                    ast.Name(id="invocation_id", ctx=ast.Load()),
                                    ast.Name(id="codeflash_duration", ctx=ast.Load()),
                                    ast.Name(id="pickled_return_value", ctx=ast.Load()),
                                    ast.Constant(value=VerificationType.FUNCTION_CALL.value),
                                ],
                                ctx=ast.Load(),
                            ),
                        ],
                        keywords=[],
                    ),
                    lineno=lineno + 20,
                ),
                ast.Expr(
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id="codeflash_con", ctx=ast.Load()), attr="commit", ctx=ast.Load()
                        ),
                        args=[],
                        keywords=[],
                    ),
                    lineno=lineno + 21,
                ),
            ]
            if mode == TestingMode.BEHAVIOR
            else []
        ),
        ast.If(
            test=ast.Name(id="exception", ctx=ast.Load()),
            body=[ast.Raise(exc=ast.Name(id="exception", ctx=ast.Load()), cause=None, lineno=lineno + 22)],
            orelse=[],
            lineno=lineno + 22,
        ),
        ast.Return(value=ast.Name(id="return_value", ctx=ast.Load()), lineno=lineno + 19),
    ]
    return ast.FunctionDef(
        name="codeflash_wrap",
        args=ast.arguments(
            args=[
                ast.arg(arg="codeflash_wrapped", annotation=None),
                ast.arg(arg="codeflash_test_module_name", annotation=None),
                ast.arg(arg="codeflash_test_class_name", annotation=None),
                ast.arg(arg="codeflash_test_name", annotation=None),
                ast.arg(arg="codeflash_function_name", annotation=None),
                ast.arg(arg="codeflash_line_id", annotation=None),
                ast.arg(arg="codeflash_loop_index", annotation=None),
                *([ast.arg(arg="codeflash_cur", annotation=None)] if mode == TestingMode.BEHAVIOR else []),
                *([ast.arg(arg="codeflash_con", annotation=None)] if mode == TestingMode.BEHAVIOR else []),
            ],
            vararg=ast.arg(arg="args"),
            kwarg=ast.arg(arg="kwargs"),
            posonlyargs=[],
            kwonlyargs=[],
            kw_defaults=[],
            defaults=[],
        ),
        body=wrapper_body,
        lineno=lineno,
        decorator_list=[],
        returns=None,
    )


class AsyncDecoratorAdder(cst.CSTTransformer):
    """Transformer that adds async decorator to async function definitions."""

    def __init__(self, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR) -> None:
        """Initialize the transformer.

        Args:
        ----
            function: The FunctionToOptimize object representing the target async function.
            mode: The testing mode to determine which decorator to apply.

        """
        super().__init__()
        self.function = function
        self.mode = mode
        self.qualified_name_parts = function.qualified_name.split(".")
        self.context_stack = []
        self.added_decorator = False

        # Choose decorator based on mode
        if mode == TestingMode.BEHAVIOR:
            self.decorator_name = "codeflash_behavior_async"
        elif mode == TestingMode.CONCURRENCY:
            self.decorator_name = "codeflash_concurrency_async"
        else:
            self.decorator_name = "codeflash_performance_async"

    def visit_ClassDef(self, node: cst.ClassDef) -> None:
        # Track when we enter a class
        self.context_stack.append(node.name.value)

    def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef:
        # Pop the context when we leave a class
        self.context_stack.pop()
        return updated_node

    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
        # Track when we enter a function
        self.context_stack.append(node.name.value)

    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
        # Check if this is an async function and matches our target
        if original_node.asynchronous is not None and self.context_stack == self.qualified_name_parts:
            # Check if the decorator is already present
            has_decorator = any(
                self._is_target_decorator(decorator.decorator) for decorator in original_node.decorators
            )

            # Only add the decorator if it's not already there
            if not has_decorator:
                new_decorator = cst.Decorator(decorator=cst.Name(value=self.decorator_name))

                # Add our new decorator to the existing decorators
                updated_decorators = [new_decorator, *list(updated_node.decorators)]
                updated_node = updated_node.with_changes(decorators=tuple(updated_decorators))
                self.added_decorator = True

        # Pop the context when we leave a function
        self.context_stack.pop()
        return updated_node

    def _is_target_decorator(self, decorator_node: cst.Name | cst.Attribute | cst.Call) -> bool:
        """Check if a decorator matches our target decorator name."""
        if isinstance(decorator_node, cst.Name):
            return decorator_node.value in {
                "codeflash_trace_async",
                "codeflash_behavior_async",
                "codeflash_performance_async",
                "codeflash_concurrency_async",
            }
        if isinstance(decorator_node, cst.Call) and isinstance(decorator_node.func, cst.Name):
            return decorator_node.func.value in {
                "codeflash_trace_async",
                "codeflash_behavior_async",
                "codeflash_performance_async",
                "codeflash_concurrency_async",
            }
        return False


ASYNC_HELPER_INLINE_CODE = """import asyncio
import gc
import os
import sqlite3
import time
from functools import wraps
from pathlib import Path
from tempfile import TemporaryDirectory

import dill as pickle


def get_run_tmp_file(file_path):
    if not hasattr(get_run_tmp_file, "tmpdir"):
        get_run_tmp_file.tmpdir = TemporaryDirectory(prefix="codeflash_")
    return Path(get_run_tmp_file.tmpdir.name) / file_path


def extract_test_context_from_env():
    test_module = os.environ["CODEFLASH_TEST_MODULE"]
    test_class = os.environ.get("CODEFLASH_TEST_CLASS", None)
    test_function = os.environ["CODEFLASH_TEST_FUNCTION"]
    if test_module and test_function:
        return (test_module, test_class if test_class else None, test_function)
    raise RuntimeError(
        "Test context environment variables not set - ensure tests are run through codeflash test runner"
    )


def codeflash_behavior_async(func):
    @wraps(func)
    async def async_wrapper(*args, **kwargs):
        loop = asyncio.get_running_loop()
        function_name = func.__name__
        line_id = os.environ["CODEFLASH_CURRENT_LINE_ID"]
        loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
        test_module_name, test_class_name, test_name = extract_test_context_from_env()
        test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{loop_index}"
        if not hasattr(async_wrapper, "index"):
            async_wrapper.index = {}
        if test_id in async_wrapper.index:
            async_wrapper.index[test_id] += 1
        else:
            async_wrapper.index[test_id] = 0
        codeflash_test_index = async_wrapper.index[test_id]
        invocation_id = f"{line_id}_{codeflash_test_index}"
        class_prefix = (test_class_name + ".") if test_class_name else ""
        test_stdout_tag = f"{test_module_name}:{class_prefix}{test_name}:{function_name}:{loop_index}:{invocation_id}"
        print(f"!$######{test_stdout_tag}######$!")
        iteration = os.environ.get("CODEFLASH_TEST_ITERATION", "0")
        db_path = get_run_tmp_file(Path(f"test_return_values_{iteration}.sqlite"))
        codeflash_con = sqlite3.connect(db_path)
        codeflash_cur = codeflash_con.cursor()
        codeflash_cur.execute(
            "CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT, test_class_name TEXT, "
            "test_function_name TEXT, function_getting_tested TEXT, loop_index INTEGER, iteration_id TEXT, "
            "runtime INTEGER, return_value BLOB, verification_type TEXT)"
        )
        exception = None
        counter = loop.time()
        gc.disable()
        try:
            ret = func(*args, **kwargs)
            counter = loop.time()
            return_value = await ret
            codeflash_duration = int((loop.time() - counter) * 1_000_000_000)
        except Exception as e:
            codeflash_duration = int((loop.time() - counter) * 1_000_000_000)
            exception = e
        finally:
            gc.enable()
        print(f"!######{test_stdout_tag}######!")
        pickled_return_value = pickle.dumps(exception) if exception else pickle.dumps((args, kwargs, return_value))
        codeflash_cur.execute(
            "INSERT INTO test_results VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                test_module_name,
                test_class_name,
                test_name,
                function_name,
                loop_index,
                invocation_id,
                codeflash_duration,
                pickled_return_value,
                "function_call",
            ),
        )
        codeflash_con.commit()
        codeflash_con.close()
        if exception:
            raise exception
        return return_value
    return async_wrapper


def codeflash_performance_async(func):
    @wraps(func)
    async def async_wrapper(*args, **kwargs):
        loop = asyncio.get_running_loop()
        function_name = func.__name__
        line_id = os.environ["CODEFLASH_CURRENT_LINE_ID"]
        loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
        test_module_name, test_class_name, test_name = extract_test_context_from_env()
        test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{loop_index}"
        if not hasattr(async_wrapper, "index"):
            async_wrapper.index = {}
        if test_id in async_wrapper.index:
            async_wrapper.index[test_id] += 1
        else:
            async_wrapper.index[test_id] = 0
        codeflash_test_index = async_wrapper.index[test_id]
        invocation_id = f"{line_id}_{codeflash_test_index}"
        class_prefix = (test_class_name + ".") if test_class_name else ""
        test_stdout_tag = f"{test_module_name}:{class_prefix}{test_name}:{function_name}:{loop_index}:{invocation_id}"
        print(f"!$######{test_stdout_tag}######$!")
        exception = None
        counter = loop.time()
        gc.disable()
        try:
            ret = func(*args, **kwargs)
            counter = loop.time()
            return_value = await ret
            codeflash_duration = int((loop.time() - counter) * 1_000_000_000)
        except Exception as e:
            codeflash_duration = int((loop.time() - counter) * 1_000_000_000)
            exception = e
        finally:
            gc.enable()
        print(f"!######{test_stdout_tag}:{codeflash_duration}######!")
        if exception:
            raise exception
        return return_value
    return async_wrapper


def codeflash_concurrency_async(func):
    @wraps(func)
    async def async_wrapper(*args, **kwargs):
        function_name = func.__name__
        concurrency_factor = int(os.environ.get("CODEFLASH_CONCURRENCY_FACTOR", "10"))
        test_module_name = os.environ.get("CODEFLASH_TEST_MODULE", "")
        test_class_name = os.environ.get("CODEFLASH_TEST_CLASS", "")
        test_function = os.environ.get("CODEFLASH_TEST_FUNCTION", "")
        loop_index = os.environ.get("CODEFLASH_LOOP_INDEX", "0")
        gc.disable()
        try:
            seq_start = time.perf_counter_ns()
            for _ in range(concurrency_factor):
                result = await func(*args, **kwargs)
            sequential_time = time.perf_counter_ns() - seq_start
        finally:
            gc.enable()
        gc.disable()
        try:
            conc_start = time.perf_counter_ns()
            tasks = [func(*args, **kwargs) for _ in range(concurrency_factor)]
            await asyncio.gather(*tasks)
            concurrent_time = time.perf_counter_ns() - conc_start
        finally:
            gc.enable()
        tag = f"{test_module_name}:{test_class_name}:{test_function}:{function_name}:{loop_index}"
        print(f"!@######CONC:{tag}:{sequential_time}:{concurrent_time}:{concurrency_factor}######@!")
        return result
    return async_wrapper
"""

ASYNC_HELPER_FILENAME = "codeflash_async_wrapper.py"


def get_decorator_name_for_mode(mode: TestingMode) -> str:
    if mode == TestingMode.BEHAVIOR:
        return "codeflash_behavior_async"
    if mode == TestingMode.CONCURRENCY:
        return "codeflash_concurrency_async"
    return "codeflash_performance_async"


def write_async_helper_file(target_dir: Path) -> Path:
    """Write the async decorator helper file to the target directory."""
    helper_path = target_dir / ASYNC_HELPER_FILENAME
    if not helper_path.exists():
        helper_path.write_text(ASYNC_HELPER_INLINE_CODE, "utf-8")
    return helper_path


def add_async_decorator_to_function(
    source_path: Path,
    function: FunctionToOptimize,
    mode: TestingMode = TestingMode.BEHAVIOR,
    project_root: Path | None = None,
) -> bool:
    """Add async decorator to an async function definition and write back to file.

    Writes a helper file containing the decorator implementation to project_root (or source directory
    as fallback) and adds a standard import + decorator to the source file.

    """
    if not function.is_async:
        return False

    try:
        with source_path.open(encoding="utf8") as f:
            source_code = f.read()

        module = cst.parse_module(source_code)

        # Add the decorator to the function
        decorator_transformer = AsyncDecoratorAdder(function, mode)
        module = module.visit(decorator_transformer)

        if decorator_transformer.added_decorator:
            # Write the helper file to project_root (on sys.path) or source dir as fallback
            helper_dir = project_root if project_root is not None else source_path.parent
            write_async_helper_file(helper_dir)
            # Add the import via CST so sort_imports can place it correctly
            decorator_name = get_decorator_name_for_mode(mode)
            import_node = cst.parse_statement(f"from codeflash_async_wrapper import {decorator_name}")
            module = module.with_changes(body=[import_node, *list(module.body)])

        modified_code = sort_imports(code=module.code, float_to_top=True)
    except Exception as e:
        logger.exception(f"Error adding async decorator to function {function.qualified_name}: {e}")
        return False
    else:
        if decorator_transformer.added_decorator:
            with source_path.open("w", encoding="utf8") as f:
                f.write(modified_code)
            logger.debug(f"Applied async {mode.value} instrumentation to {source_path}")
            return True
        return False


def create_instrumented_source_module_path(source_path: Path, temp_dir: Path) -> Path:
    instrumented_filename = f"instrumented_{source_path.name}"
    return temp_dir / instrumented_filename
