from __future__ import annotations

import ast

import isort

from aiservice.common_utils import safe_isort
from core.languages.python.testgen.instrumentation.device_sync import (
    create_device_sync_precompute_statements,
    create_device_sync_statements,
)
from core.languages.python.testgen.models import TestingMode

# Creating these at module level to avoid memory leak from repeated Config() creation
# Previously creating isort.Config() on every call caused 500k+ object accumulation (1-2GB per request) as per traces in production
_ISORT_CONFIG = isort.Config(float_to_top=True)


def format_and_float_to_top(code: str) -> str:
    """Sort imports and float them to the top.

    NOTE: Black formatting is intentionally disabled because:
    1. This code is machine-generated and immediately executed (never read by humans)
    2. Black creates 100k+ AST objects for 200-500 lines, consuming 5-20MB per call
    3. With multiple calls per request, this causes workers to grow from 200MB to 7GB
    4. Import sorting is sufficient to ensure valid Python execution
    """
    try:
        formatted_code = safe_isort(code, config=_ISORT_CONFIG)
    except Exception:
        formatted_code = code
    return formatted_code.strip()


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="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=lineno + 16,
                ),
                ast.With(
                    items=[
                        ast.withitem(
                            context_expr=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Call(
                                        func=ast.Name(id="Path", ctx=ast.Load()),
                                        args=[
                                            ast.Constant(value="{codeflash_run_tmp_dir_client_side}"),
                                            ast.JoinedStr(
                                                values=[
                                                    ast.Constant(value="test_return_values_"),
                                                    ast.FormattedValue(
                                                        value=ast.Name(id="iteration", ctx=ast.Load()), conversion=-1
                                                    ),
                                                    ast.Constant(value=".bin"),
                                                ]
                                            ),
                                        ],
                                        keywords=[],
                                    ),
                                    attr="open",
                                    ctx=ast.Load(),
                                ),
                                args=[ast.Constant(value="ab")],
                                keywords=[],
                            ),
                            optional_vars=ast.Name(id="f", ctx=ast.Store()),
                        )
                    ],
                    body=[
                        ast.Assign(
                            targets=[ast.Name(id="pickled_values", 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.Tuple(
                                            elts=[
                                                ast.Name(id="args", ctx=ast.Load()),
                                                ast.Name(id="kwargs", ctx=ast.Load()),
                                                ast.Name(id="exception", ctx=ast.Load()),
                                            ],
                                            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.Tuple(
                                            elts=[
                                                ast.Name(id="args", ctx=ast.Load()),
                                                ast.Name(id="kwargs", ctx=ast.Load()),
                                                ast.Name(id="return_value", ctx=ast.Load()),
                                            ],
                                            ctx=ast.Load(),
                                        )
                                    ],
                                    keywords=[],
                                ),
                            ),
                            lineno=lineno + 17,
                        ),
                        ast.Assign(
                            targets=[ast.Name(id="_test_name", ctx=ast.Store())],
                            value=ast.Call(
                                func=ast.Attribute(
                                    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_line_id", ctx=ast.Load()), conversion=-1
                                            ),
                                        ]
                                    ),
                                    attr="encode",
                                    ctx=ast.Load(),
                                ),
                                args=[ast.Constant(value="ascii")],
                                keywords=[],
                            ),
                            lineno=lineno + 18,
                        ),
                        ast.Expr(
                            value=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Name(id="f", ctx=ast.Load()), attr="write", ctx=ast.Load()
                                ),
                                args=[
                                    ast.Call(
                                        func=ast.Attribute(
                                            value=ast.Call(
                                                func=ast.Name(id="len", ctx=ast.Load()),
                                                args=[ast.Name(id="_test_name", ctx=ast.Load())],
                                                keywords=[],
                                            ),
                                            attr="to_bytes",
                                            ctx=ast.Load(),
                                        ),
                                        args=[ast.Constant(value=4)],
                                        keywords=[ast.keyword(arg="byteorder", value=ast.Constant(value="big"))],
                                    )
                                ],
                                keywords=[],
                            )
                        ),
                        ast.Expr(
                            value=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Name(id="f", ctx=ast.Load()), attr="write", ctx=ast.Load()
                                ),
                                args=[ast.Name(id="_test_name", ctx=ast.Load())],
                                keywords=[],
                            )
                        ),
                        ast.Expr(
                            value=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Name(id="f", ctx=ast.Load()), attr="write", ctx=ast.Load()
                                ),
                                args=[
                                    ast.Call(
                                        func=ast.Attribute(
                                            value=ast.Name(id="codeflash_duration", ctx=ast.Load()),
                                            attr="to_bytes",
                                            ctx=ast.Load(),
                                        ),
                                        args=[ast.Constant(value=8)],
                                        keywords=[ast.keyword(arg="byteorder", value=ast.Constant(value="big"))],
                                    )
                                ],
                                keywords=[],
                            )
                        ),
                        ast.Expr(
                            value=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Name(id="f", ctx=ast.Load()), attr="write", ctx=ast.Load()
                                ),
                                args=[
                                    ast.Call(
                                        func=ast.Attribute(
                                            value=ast.Call(
                                                func=ast.Name(id="len", ctx=ast.Load()),
                                                args=[ast.Name(id="pickled_values", ctx=ast.Load())],
                                                keywords=[],
                                            ),
                                            attr="to_bytes",
                                            ctx=ast.Load(),
                                        ),
                                        args=[ast.Constant(value=4)],
                                        keywords=[ast.keyword(arg="byteorder", value=ast.Constant(value="big"))],
                                    )
                                ],
                                keywords=[],
                            )
                        ),
                        ast.Expr(
                            value=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Name(id="f", ctx=ast.Load()), attr="write", ctx=ast.Load()
                                ),
                                args=[ast.Name(id="pickled_values", ctx=ast.Load())],
                                keywords=[],
                            )
                        ),
                        ast.Expr(
                            value=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Name(id="f", ctx=ast.Load()), attr="write", ctx=ast.Load()
                                ),
                                args=[
                                    ast.Call(
                                        func=ast.Attribute(
                                            value=ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
                                            attr="to_bytes",
                                            ctx=ast.Load(),
                                        ),
                                        args=[ast.Constant(value=8)],
                                        keywords=[ast.keyword(arg="byteorder", value=ast.Constant(value="big"))],
                                    )
                                ],
                                keywords=[],
                            )
                        ),
                        ast.Expr(
                            value=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Name(id="f", ctx=ast.Load()), attr="write", ctx=ast.Load()
                                ),
                                args=[
                                    ast.Call(
                                        func=ast.Attribute(
                                            value=ast.Call(
                                                func=ast.Name(id="len", ctx=ast.Load()),
                                                args=[ast.Name(id="invocation_id", ctx=ast.Load())],
                                                keywords=[],
                                            ),
                                            attr="to_bytes",
                                            ctx=ast.Load(),
                                        ),
                                        args=[ast.Constant(value=4)],
                                        keywords=[ast.keyword(arg="byteorder", value=ast.Constant(value="big"))],
                                    )
                                ],
                                keywords=[],
                            )
                        ),
                        ast.Expr(
                            value=ast.Call(
                                func=ast.Attribute(
                                    value=ast.Name(id="f", ctx=ast.Load()), attr="write", ctx=ast.Load()
                                ),
                                args=[
                                    ast.Call(
                                        func=ast.Attribute(
                                            value=ast.Name(id="invocation_id", ctx=ast.Load()),
                                            attr="encode",
                                            ctx=ast.Load(),
                                        ),
                                        args=[ast.Constant(value="ascii")],
                                        keywords=[],
                                    )
                                ],
                                keywords=[],
                            )
                        ),
                    ],
                    lineno=lineno + 17,
                ),
            ]
            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),
            ],
            vararg=ast.arg(arg="args"),
            kwarg=ast.arg(arg="kwargs"),
            posonlyargs=[],
            kwonlyargs=[],
            kw_defaults=[],
            defaults=[],
        ),
        body=wrapper_body,
        lineno=lineno,
        decorator_list=[],
        returns=None,
    )


behavior_logging_and_wrapper_code = """
from __future__ import annotations

import gc
import inspect
import os
import time
import dill as pickle

from pathlib import Path
from typing import Any, Callable, Optional


def codeflash_wrap(
    codeflash_wrapped: Callable[..., Any],
    codeflash_test_module_name: str,
    codeflash_test_class_name: str | None,
    codeflash_test_name: str,
    codeflash_function_name: str,
    codeflash_line_id: str,
    codeflash_loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> Any:
    test_id = f"{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}"
    if not hasattr(codeflash_wrap, "index"):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f"{codeflash_line_id}_{codeflash_test_index}"
    test_stdout_tag = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}"
    print(
        f"!$######{test_stdout_tag}######$!"
    )
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    print(f"!######{test_stdout_tag}######!")
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with Path(
        "{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"
    ).open("ab") as f:
        pickled_values = (
            pickle.dumps((args, kwargs, exception))
            if exception
            else pickle.dumps((args, kwargs, return_value))
        )
        _test_name = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_line_id}".encode(
            "ascii"
        )
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(codeflash_duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(codeflash_loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))
    if exception:
        raise exception
    return return_value
"""


perf_logging_and_wrapper_code = """
from __future__ import annotations

import gc
import os
import time
from typing import Any, Callable, Optional


def codeflash_wrap(
    codeflash_wrapped: Callable[..., Any],
    codeflash_test_module_name: str,
    codeflash_test_class_name: Optional[str],
    codeflash_test_name: str,
    codeflash_function_name: str,
    codeflash_line_id: str,
    codeflash_loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> str:
    test_id = f"{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}"
    if not hasattr(codeflash_wrap, "index"):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f"{codeflash_line_id}_{codeflash_test_index}"
    exception = None
    test_stdout_tag = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}"
    print(f"!$######{test_stdout_tag}######$!")
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    print(f"!######{test_stdout_tag}:{codeflash_duration}######!")
    if exception:
        raise exception
    return return_value
"""


def inject_logging_code(test_code: str, mode: TestingMode, used_frameworks: dict[str, str] | None = None) -> str:
    """Inject logging code using AST-based wrapper with GPU sync support."""
    wrapper_func = create_wrapper_function(mode, used_frameworks)

    imports: list[ast.stmt] = [
        ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0),
        ast.Import(names=[ast.alias(name="gc")]),
        *([ast.Import(names=[ast.alias(name="inspect")])] if mode == TestingMode.BEHAVIOR else []),
        ast.Import(names=[ast.alias(name="os")]),
        ast.Import(names=[ast.alias(name="time")]),
        *(
            [
                ast.ImportFrom(module="pathlib", names=[ast.alias(name="Path")], level=0),
                ast.Import(names=[ast.alias(name="dill", asname="pickle")]),
            ]
            if mode == TestingMode.BEHAVIOR
            else []
        ),
    ]
    # Add framework imports for GPU sync code (needed when framework is only imported via submodule)
    if used_frameworks:
        for framework_name, framework_alias in used_frameworks.items():
            if framework_alias == framework_name:
                imports.append(ast.Import(names=[ast.alias(name=framework_name)]))
            else:
                imports.append(ast.Import(names=[ast.alias(name=framework_name, asname=framework_alias)]))

    wrapper_module = ast.Module(body=[*imports, wrapper_func], type_ignores=[])
    wrapper_code = ast.unparse(wrapper_module)

    return format_and_float_to_top(wrapper_code + "\n" + test_code)
