"""AST transformers for test instrumentation.

Provides the ``InjectPerfOnly`` transformer that rewrites existing test
functions to wrap target-function calls with timing and capture logic,
and supporting transformers for async functions.

This module re-exports the full public API from its sub-modules so that
existing callers can continue to import from ``_instrumentation`` without
changes.
"""

from __future__ import annotations

import ast
import logging
from pathlib import Path
from typing import TYPE_CHECKING

from .._model import (
    TestingMode,
)
from ..analysis._formatter import (
    sort_imports as sort_imports,  # noqa: PLC0414
)
from ..runtime._codeflash_wrap_decorator import (
    get_run_tmp_file as get_run_tmp_file,  # noqa: PLC0414
)
from ._instrument_async import (
    ASYNC_HELPER_FILENAME as ASYNC_HELPER_FILENAME,  # noqa: PLC0414
)
from ._instrument_async import (
    ASYNC_HELPER_INLINE_CODE as ASYNC_HELPER_INLINE_CODE,  # noqa: PLC0414
)
from ._instrument_async import (
    AsyncCallInstrumenter as AsyncCallInstrumenter,  # noqa: PLC0414
)
from ._instrument_async import (
    AsyncDecoratorAdder as AsyncDecoratorAdder,  # noqa: PLC0414
)
from ._instrument_async import (
    add_async_decorator_to_function as add_async_decorator_to_function,  # noqa: PLC0414
)
from ._instrument_async import (
    get_decorator_name_for_mode as get_decorator_name_for_mode,  # noqa: PLC0414
)
from ._instrument_async import (
    inject_async_profiling_into_existing_test as inject_async_profiling_into_existing_test,  # noqa: PLC0414
)
from ._instrument_async import (
    write_async_helper_file as write_async_helper_file,  # noqa: PLC0414
)
from ._instrument_capture import (
    InitDecorator as InitDecorator,  # noqa: PLC0414
)
from ._instrument_capture import (
    add_codeflash_capture_to_init as add_codeflash_capture_to_init,  # noqa: PLC0414
)
from ._instrument_capture import (
    create_instrumented_source_module_path as create_instrumented_source_module_path,  # noqa: PLC0414
)
from ._instrument_capture import (
    instrument_codeflash_capture as instrument_codeflash_capture,  # noqa: PLC0414
)
from ._instrument_capture import (
    revert_instrumented_files as revert_instrumented_files,  # noqa: PLC0414
)
from ._instrument_core import (
    FunctionCallNodeArguments as FunctionCallNodeArguments,  # noqa: PLC0414
)
from ._instrument_core import (
    FunctionImportedAsVisitor as FunctionImportedAsVisitor,  # noqa: PLC0414
)
from ._instrument_core import (
    create_device_sync_precompute_statements as create_device_sync_precompute_statements,  # noqa: PLC0414
)
from ._instrument_core import (
    create_device_sync_statements as create_device_sync_statements,  # noqa: PLC0414
)
from ._instrument_core import (
    create_wrapper_function as create_wrapper_function,  # noqa: PLC0414
)
from ._instrument_core import (
    detect_frameworks_from_code as detect_frameworks_from_code,  # noqa: PLC0414
)
from ._instrument_core import (
    get_call_arguments as get_call_arguments,  # noqa: PLC0414
)
from ._instrument_core import (
    is_argument_name as is_argument_name,  # noqa: PLC0414
)
from ._instrument_core import (
    node_in_call_position as node_in_call_position,  # noqa: PLC0414
)

if TYPE_CHECKING:
    from collections.abc import Iterable

    from .._model import FunctionToOptimize
    from ..test_discovery.models import CodePosition

log = logging.getLogger(__name__)


class InjectPerfOnly(ast.NodeTransformer):
    """Inject performance profiling into existing test functions."""

    def __init__(
        self,
        function: FunctionToOptimize,
        module_path: str,
        call_positions: list[CodePosition],
        mode: TestingMode = TestingMode.BEHAVIOR,
    ) -> None:
        """Initialize with the target function, module path, and testing mode."""
        self.mode: TestingMode = mode
        self.function_object = function
        self.class_name: str | None = 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.parents[0].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:
        """Find and rewrite target function calls within a test statement."""
        # ast.walk is expensive for big trees and only
        # checks for ast.Call, so visit nodes manually.
        # Only descend into expressions/statements.

        # Helper for manual walk
        def iter_ast_calls(node: ast.AST) -> Iterable[ast.Call]:
            """Yield all ast.Call nodes reachable from the given node."""
            # Yield each ast.Call in test_node
            stack = [node]
            while stack:
                n = stack.pop()
                if isinstance(n, ast.Call):
                    yield n
                # Specialized BFS instead of ast.walk
                # for less overhead
                for _field, value in ast.iter_fields(n):
                    if isinstance(value, list):
                        stack.extend(
                            item
                            for item in reversed(value)
                            if isinstance(item, ast.AST)
                        )
                    elif isinstance(value, ast.AST):
                        stack.append(value)

        # Single stack instead of O(N) stack-frames
        # per child-node, less Python call overhead.
        return_statement = [test_node]
        call_node = None

        # Convert mode, function_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,
        # look up AST objects 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 (perf)
                # 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 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:
        """Visit test methods inside a class definition."""
        # TODO: Ensure this class inherits from
        # unittest.TestCase.
        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:
        """Instrument a test function by wrapping target function calls."""
        if node.name.startswith("test_"):
            did_update = False
            i = len(node.body) - 1
            while i >= 0:
                line_node = node.body[i]
                # TODO: Validate that the call
                # did not raise 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


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]:
    """Inject instrumentation into an existing test file.

    For sync functions, applies the ``InjectPerfOnly`` transformer.
    For async functions, delegates to async-specific instrumentation.
    Returns *(did_instrument, modified_source)*.
    """
    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:
        log.exception("Syntax error in code in file - %s", test_path)
        return False, None

    from ..test_discovery.linking import (  # noqa: PLC0415
        module_name_from_file_path,
    )

    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: list[ast.stmt] = [
        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")]),
            ]
        )
    for framework_name, framework_alias in used_frameworks.items():
        if framework_alias == framework_name:
            new_imports.append(
                ast.Import(names=[ast.alias(name=framework_name)])
            )
        else:
            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)
