"""Shared AST utilities, device-sync helpers, and the wrapper function builder.

Provides low-level helpers used by both sync and async instrumentation
paths: call-position matching, argument extraction, framework detection,
GPU device synchronization AST generation, and the ``codeflash_wrap``
wrapper function builder.
"""

from __future__ import annotations

import ast
import logging
from typing import TYPE_CHECKING

import attrs

from .._model import (
    FunctionParent,
    FunctionToOptimize,
    TestingMode,
    VerificationType,
)

if TYPE_CHECKING:
    from ..test_discovery.models import CodePosition

log = logging.getLogger(__name__)


@attrs.frozen
class FunctionCallNodeArguments:
    """Arguments extracted from an AST Call node."""

    args: list[ast.expr]
    keywords: list[ast.keyword]


def get_call_arguments(
    call_node: ast.Call,
) -> FunctionCallNodeArguments:
    """Extract args and keywords from an AST Call node."""
    return FunctionCallNodeArguments(call_node.args, call_node.keywords)


def node_in_call_position(
    node: ast.AST, call_positions: list[CodePosition]
) -> bool:
    """Return True if the AST node overlaps any of the given call positions.

    When *call_positions* is empty, every ``ast.Call`` node matches.
    """
    if not call_positions:
        return isinstance(node, ast.Call)
    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:
    """Check if *name* is an argument in the given arguments node."""
    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 FunctionImportedAsVisitor(ast.NodeVisitor):
    """Check if a function was imported as an alias.

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

    def __init__(self, function: FunctionToOptimize) -> None:
        """Initialize with the target function to look for import aliases."""
        assert len(function.parents) <= 1, (  # noqa: S101
            "Only support functions with one or less parent"
        )
        self.imported_as: FunctionToOptimize = function
        self.function = function
        if function.parents:
            self.to_match = function.parents[0].name
        else:
            self.to_match = function.function_name

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        """Detect import aliases for the target function."""
        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 detect_frameworks_from_code(
    code: str,
) -> dict[str, str]:
    """Detect GPU/device frameworks used in code.

    Analyzes imports for torch, tensorflow, and jax.

    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 or module_name
                elif module_name == "tensorflow":
                    frameworks["tensorflow"] = alias.asname or module_name
                elif module_name == "jax":
                    frameworks["jax"] = alias.asname or 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 create_device_sync_precompute_statements(
    used_frameworks: dict[str, str] | None,
) -> list[ast.stmt]:
    """Pre-compute device sync conditions.

    Moves conditional checks (is_available,
    hasattr, etc.) outside the timing block to
    avoid overhead affecting measurements.

    Args:
        used_frameworks: Framework-to-alias map

    Returns:
        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"]
        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,
            )
        )
        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"]
        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"]
        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,  # noqa: FBT001, FBT002
) -> list[ast.stmt]:
    """Create AST device sync statements.

    Uses pre-computed boolean conditions.

    Args:
        used_frameworks: Framework-to-alias map
        for_return_value: If True, sync after
            function call (includes JAX).

    Returns:
        AST statements for device sync.

    """
    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"]
        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 sync (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"]
        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"]
        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:
    """Build an AST FunctionDef for the codeflash_wrap instrumentation wrapper."""
    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
        # 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
                *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,
        type_params=[],
    )
