from __future__ import annotations

import ast
import platform
from dataclasses import dataclass
from typing import Optional

from aiservice.models.functions_to_optimize import FunctionParent, FunctionToOptimize


@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": "seed",
    "torch": "manual_seed",
    "tensorflow": "set_seed",
    # Add more modules as needed
}


class FunctionImportedAsVisitor(ast.NodeVisitor):
    """This checks if a function has been imported as an alias. We only care about the alias then.
    from numpy import array as np_array
    np_array is what we want
    """

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

    # TODO: Validate if the function imported is actually from the right module
    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        for alias in node.names:
            if alias.name == self.to_match and hasattr(alias, "asname") and alias.asname is not None:
                if self.function.parents:
                    self.imported_as = FunctionToOptimize(
                        function_name=self.function.function_name,
                        parents=[FunctionParent(alias.asname, "ClassDef")],
                        file_path=self.function.file_path,
                        starting_line=self.function.starting_line,
                        ending_line=self.function.ending_line,
                    )
                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,
    ) -> 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()  # Tracks original module names or aliases from 'import' statements
        self.present_aliases = set()  # Tracks aliases from 'from ... import ...' statements
        self.additional_imports = []  # Stores additional import nodes to be added

        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

    def visit_ImportFrom(self, node: ast.ImportFrom):
        if any([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 = alias.asname or imported_name
            # Track the alias from 'from ... import ...' statements
            self.present_aliases.add(module_alias)
            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
            self.present_modules.add(module_alias)
            if module_name in RNG_MODULES_SEEDS:
                self.random_test = True
                self.random_modules[module_alias] = module_name
        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) -> Optional[ast.ClassDef]:
        # 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) -> bool:
        return isinstance(node, ast.Call) and (
            (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)
        )

    def update_line_node(
        self,
        test_node: ast.Call,
        test_node_arguments: FunctionCallNodeArguments,
        node_name: str,
        index: str,
        class_name: Optional[str] = None,
    ) -> list[ast.Assign | ast.Expr]:
        if isinstance(test_node.func, ast.Name):
            function_id = test_node.func.id
        else:
            function_id = test_node.func.attr

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

        return [
            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=[test_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,
            ),
            ast.Assign(
                targets=[
                    ast.Tuple(
                        elts=[
                            ast.Name(id="codeflash_return_value", ctx=ast.Store()),
                            ast.Name(id="codeflash_duration", ctx=ast.Store()),
                            ast.Name(id="codeflash_invocation_id", ctx=ast.Store()),
                        ],
                        ctx=ast.Store(),
                    )
                ],
                value=ast.Call(
                    func=ast.Name(id="codeflash_wrap", ctx=ast.Load()),
                    args=[
                        ast.Name(id=function_id, ctx=ast.Load()),
                        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()),
                    ],
                    keywords=[
                        ast.keyword(
                            value=ast.Attribute(
                                value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                attr="arguments",
                                ctx=ast.Load(),
                            )
                        )
                    ],
                ),
                lineno=test_node.lineno + 2,
                col_offset=test_node.col_offset,
            ),
            ast.Expr(
                ast.Call(
                    func=ast.Name(id="_log__test__values", ctx=ast.Load()),
                    args=[
                        ast.Tuple(
                            elts=[
                                ast.Attribute(
                                    value=ast.Name(id="_call__bound__arguments", ctx=ast.Load()),
                                    attr="arguments",
                                    ctx=ast.Load(),
                                ),
                                ast.Name(id="codeflash_return_value", ctx=ast.Load()),
                            ],
                            ctx=ast.Load(),
                        ),
                        ast.Name(id="codeflash_duration", ctx=ast.Load()),
                        ast.Constant(
                            value=(
                                f"{self.test_module_path}:{class_name or ''}{'.' if class_name is not None else ''}"
                                f"{node_name}:{self.only_function_name}:{index}"
                            )
                        ),
                        ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
                        ast.Name(id="codeflash_invocation_id", ctx=ast.Load()),
                    ],
                    keywords=[],
                ),
                lineno=test_node.lineno + 3,
                col_offset=test_node.col_offset,
            ),
        ]

    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:
            if alias in self.present_modules:
                module_ref = alias
            else:
                module_ref = full_module_name

            # Build the seed call using the module name
            module_parts = module_ref.split(".")
            module = 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 = 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
            seed_call = ast.Expr(value=ast.Call(func=seed_func, args=[ast.Constant(value=42)], keywords=[]))
            return seed_call
        return None

    def find_target_function_call(self, node: ast.AST) -> ast.Call | None:
        for child in ast.walk(node):
            if isinstance(child, ast.Call) and self.is_target_function_node(child):
                return child
        return None

    def process_statements(
        self,
        statements: list[ast.stmt],
        node_name: str,
        class_name: str | None,
        offset: int = 0,
        parent_index: Optional[str] = None,
    ) -> None:
        i = 0
        while i < len(statements):
            stmt = statements[i]
            if parent_index is None:
                current_index = str(i + offset)
            else:
                current_index = 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)
                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)
                if isinstance(stmt, ast.Try) and stmt.handlers:
                    for idx, handler in enumerate(stmt.handlers):
                        self.process_statements(
                            handler.body, node_name, class_name, offset=0, parent_index=current_index
                        )
            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
                    updated_nodes = self.update_line_node(
                        test_node, get_call_arguments(test_node), node_name, current_index, class_name
                    )
                    statements[i : i + 1] = updated_nodes
                    i += len(updated_nodes) - 1  # Adjust index for inserted nodes

            i += 1

    def visit_FunctionDef(self, node: ast.FunctionDef, class_name: Optional[str] = None) -> Optional[ast.FunctionDef]:
        random_index = 0
        random_test_function = False
        random_importfrom = False
        if node.name == self.only_function_name or node.name in self.helper_function_names:
            return None  # Remove the re-definition of the function and its dependencies from the test generation code
        if node.name.startswith("test_"):
            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)

            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)
            offset = len(new_body)
            # Add the original function body
            # new_body.extend(node.body)

            # Process each statement in the function body
            for stmt in 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
                    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

                # For other statements, process them recursively
                self.generic_visit(stmt)
                stmt_list = [stmt]
                self.process_statements(stmt_list, node.name, class_name, offset=offset)
                new_body.extend(stmt_list)
                offset += len(stmt_list)  # Update offset for each new statement added
            # Update the function body
            node.body = new_body

            if self.test_framework == "unittest":
                # TODO: Make sure that if the test times out,
                #  the test's time is excluded from the total time calculation and comparison
                node.decorator_list.append(
                    ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id="timeout_decorator", ctx=ast.Load()), attr="timeout", ctx=ast.Load()
                        ),
                        args=[ast.Constant(value=self.individual_test_timeout)],
                        keywords=[ast.keyword(arg="use_signals", value=ast.Constant(value=True))],
                    )
                )
                # TODO: The value of use_signals should be False,
                #  otherwise this will fail in multiprocessing environments
        return node


def inject_logging_code(test_code: str) -> str:
    # TODO : Port this to Sqlite3 - makes it standard and easier to use
    logging_code = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))
"""
    return logging_code + create_wrapper_function() + test_code


def create_wrapper_function() -> str:
    return """
def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
"""


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] = tuple([int(ver) for ver in platform.python_version_tuple()]),
) -> str:
    module_node = ast.parse(test_source, feature_version=python_version[:2])
    import_visitor = FunctionImportedAsVisitor(function_to_optimize)
    import_visitor.visit(module_node)
    func = import_visitor.imported_as
    new_module_node = InjectPerfAndLogging(
        function=func,
        helper_function_names=helper_function_names,
        test_module_path=test_module_path,
        test_framework=test_framework,
        test_timeout=test_timeout,
    ).visit(module_node)
    new_imports = [
        ast.Import(names=[ast.alias(name="inspect")]),
        # This adds the function to the namespace
        ast.ImportFrom(
            module=module_path,
            names=[ast.alias(name=function_to_optimize.top_level_parent_name)],
            col_offset=0,
            lineno=1,
            level=0,
        ),
    ]
    if test_framework == "unittest":
        new_imports += [ast.Import(names=[ast.alias(name="timeout_decorator")])]
    new_module_node.body = [*new_imports, *new_module_node.body]
    new_tests = ast.unparse(new_module_node)
    modified_new_tests = inject_logging_code(new_tests)
    return modified_new_tests
