from __future__ import annotations

import ast
import platform
from collections import deque
from dataclasses import dataclass

import sentry_sdk

from aiservice.models.functions_to_optimize import FunctionParent, FunctionToOptimize
from core.languages.python.testgen.ast_utils.test_detection import is_test_function_name
from core.languages.python.testgen.instrumentation.device_sync import detect_frameworks_from_code
from core.languages.python.testgen.instrumentation.wrapper import format_and_float_to_top, inject_logging_code
from core.languages.python.testgen.models import TestingMode

plat_str = platform.python_version_tuple()
platform_tuple: tuple[int, int, int] = int(plat_str[0]), int(plat_str[1]), int(plat_str[2])


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


# Mapping of RNG modules to their seed-setting functions
RNG_MODULES_SEEDS = {
    "random": "seed",
    "numpy.random": "seed",
    "numpy": "random.seed",
    "torch": "manual_seed",
    "tensorflow": "random.set_seed",
    "tensorflow.random": "set_seed",
    # Add more modules as needed
}


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

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

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

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


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


class InjectPerfAndLogging(ast.NodeTransformer):
    def __init__(
        self,
        function: FunctionToOptimize,
        helper_function_names: list[str],
        test_module_path: str,
        test_framework: str = "pytest",
        test_timeout: int = 15,
        mode: TestingMode = TestingMode.BEHAVIOR,
    ) -> None:
        self.function_object = function
        self.class_name = None
        self.only_function_name = function.function_name
        self.test_framework = test_framework
        self.individual_test_timeout = test_timeout
        self.test_module_path = test_module_path
        self.random_test = False
        self.random_modules: dict[str, str] = {}  # alias -> module_name
        self.present_modules: set[str] = set()  # Tracks original module names or aliases from 'import' statements
        self.present_aliases: set[str] = set()  # Tracks aliases from 'from ... import ...' statements
        self.additional_imports: list[ast.Import] = []  # Stores additional import nodes to be added
        self.mode = mode
        self.async_call_counter: dict[str, int] = {}

        if len(function.parents) == 1 and function.parents[0].type == "ClassDef":
            self.class_name = function.top_level_parent_name
        self.helper_function_names = helper_function_names  # Other functional dependencies that were injected
        self.removed_helper_functions: set[str] = (
            set()
        )  # Track which helper functions were removed during instrumentation

        # Cache for parent mappings to avoid repeated ast.walk() calls
        self._parent_map: dict[ast.AST, ast.AST] | None = None

    def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.ImportFrom | None:
        if all(name.name in [self.only_function_name, self.class_name] for name in node.names):
            return None  # Remove the import of the function the test generation code
        module_name = node.module  # e.g., 'numpy.random'
        if module_name is None:
            # Handle relative imports or cases where module_name is None
            return node

        for alias in node.names:  # e.g., 'rand'
            imported_name = alias.name  # e.g., 'random', 'randint'
            module_alias = module_name
            self.present_aliases.add(alias.asname or imported_name)
            if module_name in RNG_MODULES_SEEDS:
                self.random_test = True
                self.random_modules[module_alias] = module_name

                if module_name not in self.present_modules:
                    # Add an import statement for the module if it's not already imported
                    self.present_modules.add(module_name)
                    import_node = ast.Import(names=[ast.alias(name=module_name)])
                    self.additional_imports.append(import_node)
        return node

    def visit_Import(self, node: ast.Import) -> ast.Import:
        for alias in node.names:
            module_name = alias.name
            module_alias = alias.asname or module_name
            full_module_path = module_name
            if full_module_path in RNG_MODULES_SEEDS:
                self.random_test = True
                self.random_modules[module_alias] = full_module_path
        return node

    def visit_Module(self, node: ast.Module) -> ast.Module:
        # Reset additional_imports at the start of each module visit
        self.additional_imports = []
        self.generic_visit(node)
        # Insert additional imports at the beginning of the module to avoid conflicts
        node.body = self.additional_imports + node.body
        return node

    def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef | None:
        # If the original class exists during testing, then remove it all together
        if node.name == self.class_name:
            return None  # Remove the re-definition of the class and its dependencies from the test generation code
        for inner_node in ast.walk(node):
            if isinstance(inner_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                self.visit_FunctionDef(inner_node, node.name)
        return node

    def is_target_function_node(self, node: ast.AST) -> bool:
        # Check for regular call nodes
        if isinstance(node, ast.Call):
            return (isinstance(node.func, ast.Name) and node.func.id == self.only_function_name) or (
                isinstance(node.func, ast.Attribute) and node.func.attr == self.only_function_name
            )
        # Check for await expressions (for async functions)
        if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
            return (isinstance(node.value.func, ast.Name) and node.value.func.id == self.only_function_name) or (
                isinstance(node.value.func, ast.Attribute) and node.value.func.attr == self.only_function_name
            )
        return False

    def update_line_node(
        self,
        test_node: ast.Call | ast.Await,
        test_node_arguments: FunctionCallNodeArguments,
        node_name: str,
        index: str,
        class_name: str | None = None,
        parent_statement: ast.stmt | None = None,
    ) -> list[ast.stmt]:
        if self.function_object.is_async:
            if node_name not in self.async_call_counter:
                self.async_call_counter[node_name] = 0
            current_call_index = str(self.async_call_counter[node_name])
            self.async_call_counter[node_name] += 1

            env_assignment = ast.Assign(
                targets=[
                    ast.Subscript(
                        value=ast.Attribute(value=ast.Name(id="os", ctx=ast.Load()), attr="environ", ctx=ast.Load()),
                        slice=ast.Constant(value="CODEFLASH_CURRENT_LINE_ID"),
                        ctx=ast.Store(),
                    )
                ],
                value=ast.Constant(value=current_call_index),
                lineno=test_node.lineno if hasattr(test_node, "lineno") else 1,
                col_offset=test_node.col_offset if hasattr(test_node, "col_offset") else 0,
            )
            return (
                [env_assignment, parent_statement] if parent_statement else [env_assignment, ast.Expr(value=test_node)]
            )

        call_node = test_node.value if isinstance(test_node, ast.Await) else test_node
        if not isinstance(call_node, ast.Call):
            return []
        if isinstance(call_node.func, ast.Name):
            function_id = call_node.func.id
        elif isinstance(call_node.func, ast.Attribute):
            function_id = call_node.func.attr

        if isinstance(call_node.func, ast.Attribute) and isinstance(call_node.func.value, ast.Name):
            function_id = call_node.func.value.id + "." + function_id

        new_nodes = [
            *(
                [
                    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=ast.Name(id="inspect", ctx=ast.Load()), attr="signature", ctx=ast.Load()
                                    ),
                                    args=[call_node.func],
                                    keywords=[],
                                ),
                                attr="bind",
                                ctx=ast.Load(),
                            ),
                            args=test_node_arguments.args,
                            keywords=test_node_arguments.keywords,
                        ),
                        lineno=test_node.lineno,
                        col_offset=test_node.col_offset,
                    ),
                    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,
                    ),
                ]
                if self.mode == TestingMode.BEHAVIOR
                else []
            ),
            ast.Assign(
                targets=[ast.Name(id="codeflash_return_value", ctx=ast.Store())],
                value=ast.Call(
                    func=ast.Name(id="codeflash_wrap", ctx=ast.Load()),
                    args=[
                        call_node.func,
                        ast.Constant(value=self.test_module_path),
                        ast.Constant(value=class_name or None),
                        ast.Constant(value=node_name),
                        ast.Constant(value=self.only_function_name),
                        ast.Constant(value=index),
                        ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
                        *(
                            call_node.args
                            if self.mode == TestingMode.PERFORMANCE
                            else [
                                ast.Starred(
                                    value=ast.Attribute(
                                        value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                        attr="args",
                                        ctx=ast.Load(),
                                    ),
                                    ctx=ast.Load(),
                                )
                            ]
                        ),
                    ],
                    keywords=[
                        ast.keyword(
                            value=ast.Attribute(
                                value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                attr="kwargs",
                                ctx=ast.Load(),
                            )
                        )
                    ]
                    if self.mode == TestingMode.BEHAVIOR
                    else call_node.keywords,
                ),
                lineno=test_node.lineno + 2,
                col_offset=test_node.col_offset,
            ),
        ]
        if parent_statement and isinstance(parent_statement, ast.Assign) and parent_statement.targets:
            if len(parent_statement.targets) == 1:
                new_nodes.append(
                    ast.Assign(
                        targets=parent_statement.targets,
                        value=ast.Name(id="codeflash_return_value", ctx=ast.Load()),
                        lineno=test_node.lineno + 3,
                        col_offset=test_node.col_offset,
                    )
                )
            else:
                new_nodes.append(
                    ast.Assign(
                        targets=parent_statement.targets,
                        value=ast.Tuple(
                            elts=[
                                ast.Subscript(
                                    value=ast.Name(id="codeflash_return_value", ctx=ast.Load()),
                                    slice=ast.Constant(value=i),
                                    ctx=ast.Load(),
                                )
                                for i in range(len(parent_statement.targets))
                            ]
                        ),
                        lineno=test_node.lineno + 3,
                        col_offset=test_node.col_offset,
                    )
                )

        return new_nodes

    def create_seed_statement(self, alias: str, full_module_name: str) -> ast.stmt | None:
        seed_method = RNG_MODULES_SEEDS.get(full_module_name)
        if seed_method:
            module_alias = alias
            module_parts = module_alias.split(".")
            module: ast.Name | ast.Attribute = ast.Name(id=module_parts[0], ctx=ast.Load())
            for part in module_parts[1:]:
                module = ast.Attribute(value=module, attr=part, ctx=ast.Load())
            # Split the seed method if it's nested (e.g., 'random.seed')
            seed_parts = seed_method.split(".")
            # Build the attribute chain
            seed_func: ast.Name | ast.Attribute = module
            for part in seed_parts:
                seed_func = ast.Attribute(value=seed_func, attr=part, ctx=ast.Load())
            # Create the call to set the seed
            return ast.Expr(value=ast.Call(func=seed_func, args=[ast.Constant(value=42)], keywords=[]))
        return None

    def is_inside_comprehension(self, node: ast.AST, target_call: ast.AST) -> bool:
        """Check if target_call is inside a comprehension within node."""
        # Build parent map once per root node
        if self._parent_map is None:
            self._parent_map = self._build_parent_map(node)

        # Traverse upward from target_call using parent map
        current = target_call
        while current in self._parent_map:
            current = self._parent_map[current]
            if isinstance(current, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)):
                return True
            # Stop if we've reached the original node boundary
            if current is node:
                break

        return False

    def find_target_function_call(self, node: ast.AST) -> ast.Call | ast.Await | None:
        # Build a parent map with a single BFS traversal to avoid quadratic ast.walk calls
        queue: deque[ast.AST] = deque([node])
        parent: dict[ast.AST, ast.AST | None] = {node: None}

        while queue:
            current = queue.popleft()

            if self.is_target_function_node(current) and isinstance(current, (ast.Call, ast.Await)):
                # Walk up ancestors to check if inside a comprehension
                ancestor = parent.get(current)
                while ancestor is not None:
                    if isinstance(ancestor, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)):
                        break
                    ancestor = parent.get(ancestor)
                else:
                    # Not inside comprehension
                    # For chained calls like action_wrap(X)(Y), walk up to find the outermost call
                    # where the target function call is used as the func of another call
                    result = current
                    while True:
                        p = parent.get(result)
                        if isinstance(p, ast.Call) and p.func is result:
                            result = p
                        else:
                            break
                    return result

            # Enqueue children and record their parent
            for child in ast.iter_child_nodes(current):
                if child not in parent:
                    parent[child] = current
                    queue.append(child)

        return None

    def process_statements(
        self,
        statements: list[ast.stmt],
        node_name: str,
        class_name: str | None,
        offset: int = 0,
        parent_index: str | None = None,
    ) -> None:
        i = 0
        while i < len(statements):
            stmt = statements[i]
            current_index = str(i + offset) if parent_index is None else parent_index + "_" + str(i)
            if isinstance(stmt, (ast.With, ast.For, ast.While, ast.If, ast.Try, ast.AsyncWith, ast.AsyncFor)):
                # Recursively process the body of the control flow statement
                self.process_statements(stmt.body, node_name, class_name, offset=0, parent_index=current_index)
                # Handle else, finally blocks, and exception handlers similarly
                if hasattr(stmt, "orelse"):  # Some statements (like if, for, while) can have an else block
                    self.process_statements(stmt.orelse, node_name, class_name, offset=0, parent_index=current_index)  # type: ignore[attr-defined]
                if hasattr(stmt, "finalbody"):  # try statements can have a finally block
                    self.process_statements(stmt.finalbody, node_name, class_name, offset=0, parent_index=current_index)  # type: ignore[attr-defined]
                if isinstance(stmt, ast.Try) and stmt.handlers:
                    for handler in stmt.handlers:
                        self.process_statements(
                            handler.body, node_name, class_name, offset=0, parent_index=current_index
                        )
            elif isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                # Skip nested function and class definitions - don't instrument calls inside them
                # The nested definition will still execute at runtime, preserving its definition
                pass
            else:
                # Check if the statement contains a call to the target function
                test_node = self.find_target_function_call(stmt)
                if test_node is not None:
                    # Replace the statement with the updated nodes
                    call_node = test_node.value if isinstance(test_node, ast.Await) else test_node
                    if not isinstance(call_node, ast.Call):
                        continue
                    updated_nodes = self.update_line_node(
                        test_node,
                        get_call_arguments(call_node),
                        node_name,
                        current_index,
                        class_name,
                        parent_statement=stmt,
                    )
                    statements[i : i + 1] = updated_nodes
                    i += len(updated_nodes) - 1  # Adjust index for inserted nodes

                elif isinstance(stmt, ast.Expr) and isinstance(
                    stmt.value, ast.Call
                ):  # handles https://linear.app/codeflash-ai/issue/CF-500/numpy-compatibility#comment-6d956d3e
                    numpy_test_funcs = (
                        "assert_array_almost_equal",
                        "assert_array_equal",
                        "assert_allclose",
                        "assert_almost_equal",
                        "assert_equal",
                        "assert_array_almost_equal_nulp",
                        "assert_array_max_ulp",
                        "assert_array_less",
                        "assert_string_equal",
                        "assert_approx_equal",
                    )
                    if (
                        isinstance(stmt.value.func, ast.Attribute)
                        and stmt.value.func.attr in numpy_test_funcs
                        and stmt.value.args
                    ):
                        stmt.value.args[0] = ast.Name(id="codeflash_return_value", ctx=ast.Load())
            i += 1

    def visit_FunctionDef(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef, class_name: str | None = None
    ) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
        if node.name == self.only_function_name:
            return None  # Remove the re-definition of the function being tested
        if node.name in self.helper_function_names:
            self.removed_helper_functions.add(node.name)  # Track removed helpers so we can add imports for them
            return None  # Remove the re-definition of helper dependencies from the test generation code
        if is_test_function_name(node.name):
            self.present_modules = set()  # Reset tracked modules
            self.present_aliases = set()  # Reset tracked aliases
            new_body = []
            # Insert seed-setting statements
            for alias, full_module_name in self.random_modules.items():
                seed_statement = self.create_seed_statement(alias, full_module_name)
                if seed_statement:
                    new_body.append(seed_statement)

            # Only add codeflash_loop_index instrumentation for sync functions being tested
            # For async functions, instrumentation happens at the function definition site
            if not self.function_object.is_async:
                loop_index_assign = 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 + 1,
                    col_offset=node.col_offset,
                )
                new_body.append(loop_index_assign)

            # Modules that are added at module level during instrumentation.
            # We must remove these from inside function bodies to avoid UnboundLocalError.
            # Python treats `import X` inside a function as making X a local variable,
            # so any reference to X before that import will fail.
            modules_added_at_module_level = {"inspect"} if self.mode == TestingMode.BEHAVIOR else set()
            if self.function_object.is_async:
                modules_added_at_module_level.add("os")

            # Process each statement in the function body
            for node_body_index, stmt in enumerate(node.body):
                # Check if the statement is an import of a random module
                if isinstance(stmt, (ast.Import, ast.ImportFrom)):
                    self.visit(stmt)  # To populate self.random_modules
                    # Skip imports that are added at module level to avoid UnboundLocalError
                    if isinstance(stmt, ast.Import) and self._is_import_of_modules(stmt, modules_added_at_module_level):
                        continue
                    new_body.append(stmt)
                    # Insert seed-setting statement after import
                    alias_list = list(self.random_modules.keys())
                    if alias_list:
                        alias = alias_list[-1]  # Get the last alias added
                        full_module_name = self.random_modules[alias]
                        seed_statement = self.create_seed_statement(alias, full_module_name)
                        if seed_statement:
                            new_body.append(seed_statement)
                        # Remove the alias from random_modules to prevent duplicate seeds
                        del self.random_modules[alias]
                    continue

                # Skip generic_visit for nested class/function definitions to preserve their bodies
                # (generic_visit would remove methods matching self.only_function_name)
                if not isinstance(stmt, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
                    self.generic_visit(stmt)
                stmt_list = [stmt]
                # The offset is the index of the statement in the original test function body
                self.process_statements(stmt_list, node.name, class_name, offset=node_body_index)
                new_body.extend(stmt_list)
            # Update the function body
            node.body = new_body
        return node

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef | None:
        result = self.visit_FunctionDef(node)
        if result is None or isinstance(result, ast.AsyncFunctionDef):
            return result
        return None

    def _is_import_of_modules(self, stmt: ast.Import, module_names: set[str]) -> bool:
        """Check if an import statement imports any of the given module names.

        Returns True if all aliases in the import are for modules in module_names.
        For example, `import inspect` returns True if "inspect" is in module_names.
        `import inspect, sys` returns False if only "inspect" is in module_names.
        """
        if not stmt.names:
            return False
        return all(alias.name in module_names for alias in stmt.names)

    def _build_parent_map(self, node: ast.AST) -> dict[ast.AST, ast.AST]:
        """Build a mapping from child nodes to their parents."""
        parent_map = {}
        for parent in ast.walk(node):
            for child in ast.iter_child_nodes(parent):
                parent_map[child] = parent
        return parent_map


def instrument_test_source(
    test_source: str,
    function_to_optimize: FunctionToOptimize,
    helper_function_names: list[str],
    module_path: str,
    test_module_path: str,
    test_framework: str,
    test_timeout: int,
    python_version: tuple[int, int, int] = platform_tuple,
    mode: TestingMode = TestingMode.BEHAVIOR,
) -> str | None:
    # Detect GPU/device frameworks from the test source code
    used_frameworks = detect_frameworks_from_code(test_source)

    try:
        module_node = ast.parse(test_source, feature_version=python_version[:2])
    except Exception as e:
        e.add_note("Failed to parse test source code for instrumentation:\n" + test_source)
        sentry_sdk.capture_exception(e)
        return None

    import_visitor = FunctionImportedAsVisitor(function_to_optimize)
    import_visitor.visit(module_node)
    func = import_visitor.imported_as
    transformer = InjectPerfAndLogging(
        function=func,
        helper_function_names=helper_function_names,
        test_module_path=test_module_path,
        test_framework=test_framework,
        test_timeout=test_timeout,
        mode=mode,
    )
    new_module_node = transformer.visit(module_node)

    # Build imports: include helper functions that were removed so they can be imported from the source module
    import_names = [ast.alias(name=function_to_optimize.top_level_parent_name)]
    import_names.extend(ast.alias(name=helper_name) for helper_name in transformer.removed_helper_functions)

    new_imports = [
        *([ast.Import(names=[ast.alias(name="inspect")])] if mode == TestingMode.BEHAVIOR else []),
        *([ast.Import(names=[ast.alias(name="os")])] if function_to_optimize.is_async else []),
        # This adds the function and any removed helper functions to the namespace
        ast.ImportFrom(module=module_path, names=import_names, col_offset=0, lineno=1, level=0),
    ]
    new_module_node.body = [*new_imports, *new_module_node.body]
    new_tests = ast.unparse(new_module_node)
    if function_to_optimize.is_async:
        return format_and_float_to_top(new_tests)
    if mode not in (TestingMode.BEHAVIOR, TestingMode.PERFORMANCE):
        sentry_sdk.capture_message(f"Unknown testing mode: {mode}")
        return None
    return inject_logging_code(new_tests, mode, used_frameworks)
