"""Async-specific instrumentation: AST transformers, decorators, and helpers.

Provides ``AsyncCallInstrumenter`` for injecting ``CODEFLASH_CURRENT_LINE_ID``
assignments before ``await`` calls, ``AsyncDecoratorAdder`` for adding
async performance/behavior decorators via libcst, the inline async helper
code, and high-level functions for instrumenting async test and source files.
"""

from __future__ import annotations

import ast
import logging
from typing import TYPE_CHECKING

import libcst as cst

from .._model import (
    FunctionToOptimize,
    TestingMode,
)
from ..analysis._formatter import sort_imports
from ..test_discovery.linking import module_name_from_file_path
from ._instrument_core import (
    FunctionImportedAsVisitor,
    node_in_call_position,
)

if TYPE_CHECKING:
    from pathlib import Path

    from ..test_discovery.models import CodePosition

log = logging.getLogger(__name__)


class AsyncCallInstrumenter(ast.NodeTransformer):
    """AST transformer for async function instrumentation."""

    def __init__(
        self,
        function: FunctionToOptimize,
        module_path: str,
        call_positions: list[CodePosition],
        mode: TestingMode = TestingMode.BEHAVIOR,
    ) -> None:
        """Initialize with the target async function and testing mode."""
        self.mode = 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
        self.did_instrument = False
        self.async_call_counter: dict[str, int] = {}
        if (
            len(function.parents) == 1
            and function.parents[0].type == "ClassDef"
        ):
            self.class_name = function.parents[0].name

    def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
        """Recurse into class bodies to find test methods."""
        return self.generic_visit(node)  # type: ignore[return-value]

    def visit_AsyncFunctionDef(
        self, node: ast.AsyncFunctionDef
    ) -> ast.AsyncFunctionDef:
        """Instrument async test functions that call the target function."""
        if not node.name.startswith("test_"):
            return node

        return self._process_test_function(node)  # type: ignore[return-value]

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
        """Instrument sync test functions that call the target async function."""
        # Only process test functions
        if not node.name.startswith("test_"):
            return node

        return self._process_test_function(node)  # type: ignore[return-value]

    def _process_test_function(
        self,
        node: ast.AsyncFunctionDef | ast.FunctionDef,
    ) -> ast.AsyncFunctionDef | ast.FunctionDef:
        """Add CODEFLASH_CURRENT_LINE_ID assignments before target await calls."""
        # Initialize counter for this test function
        if node.name not in self.async_call_counter:
            self.async_call_counter[node.name] = 0

        new_body: list[ast.stmt] = []

        # Scan only relevant nodes instead of
        # full ast.walk in _instrument_statement
        for _i, stmt in enumerate(node.body):
            transformed_stmt, added_env_assignment = (
                self._optimized_instrument_statement(stmt)
            )

            if added_env_assignment:
                current_call_index = 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=f"{current_call_index}"),
                    lineno=stmt.lineno if hasattr(stmt, "lineno") else 1,
                )
                new_body.append(env_assignment)
                self.did_instrument = True

            new_body.append(transformed_stmt)

        node.body = new_body
        return node

    def _instrument_statement(
        self, stmt: ast.stmt, _node_name: str
    ) -> tuple[ast.stmt, bool]:
        """Check whether a statement contains an awaited target call."""
        for node in ast.walk(stmt):
            if (
                isinstance(node, ast.Await)
                and isinstance(node.value, ast.Call)
                and self._is_target_call(node.value)
                and self._call_in_positions(node.value)
            ):
                # Check if this call is in one of our target positions
                return (
                    stmt,
                    True,
                )  # Return original statement but signal we added env var

        return stmt, False

    def _is_target_call(self, call_node: ast.Call) -> bool:
        """Check if this call node is calling our target async function."""
        if isinstance(call_node.func, ast.Name):
            return call_node.func.id == self.function_object.function_name
        if isinstance(call_node.func, ast.Attribute):
            return call_node.func.attr == self.function_object.function_name
        return False

    def _call_in_positions(self, call_node: ast.Call) -> bool:
        """Return True if the call node is at one of the tracked positions."""
        if not hasattr(call_node, "lineno") or not hasattr(
            call_node, "col_offset"
        ):
            return False

        return node_in_call_position(call_node, self.call_positions)

    # Optimized version: only walk child nodes for Await
    def _optimized_instrument_statement(
        self, stmt: ast.stmt
    ) -> tuple[ast.stmt, bool]:
        """Stack-based search for awaited target calls in a statement."""
        # Stack-based DFS, manual for relevant Await nodes
        stack: list[ast.AST] = [stmt]
        while stack:
            node = stack.pop()
            # Favor direct ast.Await detection
            if isinstance(node, ast.Await):
                val = node.value
                if (
                    isinstance(val, ast.Call)
                    and self._is_target_call(val)
                    and self._call_in_positions(val)
                ):
                    return stmt, True
            # Use _fields instead of ast.walk for less allocations
            for fname in getattr(node, "_fields", ()):
                child = getattr(node, fname, None)
                if isinstance(child, list):
                    stack.extend(child)
                elif isinstance(child, ast.AST):
                    stack.append(child)
        return stmt, False


class AsyncDecoratorAdder(cst.CSTTransformer):
    """Transformer that adds async decorator to async function definitions."""

    def __init__(
        self,
        function: FunctionToOptimize,
        mode: TestingMode = TestingMode.BEHAVIOR,
    ) -> None:
        """Initialize the transformer.

        Args:
        ----
            function: Target async function.
            mode: Testing mode for decorator.

        """
        super().__init__()
        self.function = function
        self.mode = mode
        self.qualified_name_parts = function.qualified_name.split(".")
        self.context_stack: list[str] = []
        self.added_decorator = False

        # Choose decorator based on mode
        if mode == TestingMode.BEHAVIOR:
            self.decorator_name = "codeflash_behavior_async"
        elif mode == TestingMode.CONCURRENCY:
            self.decorator_name = "codeflash_concurrency_async"
        else:
            self.decorator_name = "codeflash_performance_async"

    def visit_ClassDef(  # noqa: N802
        self, node: cst.ClassDef
    ) -> None:
        """Push class name onto the context stack."""
        # Track when we enter a class
        self.context_stack.append(node.name.value)

    def leave_ClassDef(  # noqa: N802
        self,
        original_node: cst.ClassDef,
        updated_node: cst.ClassDef,
    ) -> cst.ClassDef:
        """Pop class name from the context stack."""
        # Pop the context when we leave a class
        self.context_stack.pop()
        return updated_node

    def visit_FunctionDef(  # noqa: N802
        self, node: cst.FunctionDef
    ) -> None:
        """Push function name onto the context stack."""
        # Track when we enter a function
        self.context_stack.append(node.name.value)

    def leave_FunctionDef(  # noqa: N802
        self,
        original_node: cst.FunctionDef,
        updated_node: cst.FunctionDef,
    ) -> cst.FunctionDef:
        """Add the async decorator if the function matches the target."""
        # Check if this is an async function and matches our target
        if (
            original_node.asynchronous is not None
            and self.context_stack == self.qualified_name_parts
        ):
            # Check if the decorator is already present
            has_decorator = any(
                self._is_target_decorator(decorator.decorator)
                for decorator in original_node.decorators
            )

            # Only add the decorator if it's not already there
            if not has_decorator:
                new_decorator = cst.Decorator(
                    decorator=cst.Name(value=self.decorator_name)
                )

                # Add our new decorator to the existing decorators
                updated_decorators = [
                    new_decorator,
                    *list(updated_node.decorators),
                ]
                updated_node = updated_node.with_changes(
                    decorators=tuple(updated_decorators)
                )
                self.added_decorator = True

        # Pop the context when we leave a function
        self.context_stack.pop()
        return updated_node

    def _is_target_decorator(self, decorator_node: cst.BaseExpression) -> bool:
        """Check if a decorator matches our target decorator name."""
        if isinstance(decorator_node, cst.Name):
            return decorator_node.value in {
                "codeflash_trace_async",
                "codeflash_behavior_async",
                "codeflash_performance_async",
                "codeflash_concurrency_async",
            }
        if isinstance(decorator_node, cst.Call) and isinstance(
            decorator_node.func, cst.Name
        ):
            return decorator_node.func.value in {
                "codeflash_trace_async",
                "codeflash_behavior_async",
                "codeflash_performance_async",
                "codeflash_concurrency_async",
            }
        return False


ASYNC_HELPER_INLINE_CODE = """import asyncio
import gc
import os
import sqlite3
import time
from functools import wraps
from pathlib import Path
from tempfile import TemporaryDirectory

import dill as pickle


def get_run_tmp_file(file_path):
    if not hasattr(get_run_tmp_file, "tmpdir"):
        get_run_tmp_file.tmpdir = TemporaryDirectory(prefix="codeflash_")
    return Path(get_run_tmp_file.tmpdir.name) / file_path


def extract_test_context_from_env():
    test_module = os.environ["CODEFLASH_TEST_MODULE"]
    test_class = os.environ.get("CODEFLASH_TEST_CLASS", None)
    test_function = os.environ["CODEFLASH_TEST_FUNCTION"]
    if test_module and test_function:
        return (test_module, test_class if test_class else None, test_function)
    raise RuntimeError(
        "Test context environment variables not set"
        " - ensure tests are run through"
        " codeflash test runner"
    )


def codeflash_behavior_async(func):
    @wraps(func)
    async def async_wrapper(*args, **kwargs):
        loop = asyncio.get_running_loop()
        function_name = func.__name__
        line_id = os.environ["CODEFLASH_CURRENT_LINE_ID"]
        loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
        (test_module_name, test_class_name,
         test_name) = extract_test_context_from_env()
        test_id = (
            f"{test_module_name}:{test_class_name}"
            f":{test_name}:{line_id}:{loop_index}"
        )
        if not hasattr(async_wrapper, "index"):
            async_wrapper.index = {}
        if test_id in async_wrapper.index:
            async_wrapper.index[test_id] += 1
        else:
            async_wrapper.index[test_id] = 0
        codeflash_test_index = async_wrapper.index[test_id]
        invocation_id = f"{line_id}_{codeflash_test_index}"
        class_prefix = (
            (test_class_name + ".") if test_class_name else ""
        )
        test_stdout_tag = (
            f"{test_module_name}:{class_prefix}"
            f"{test_name}:{function_name}"
            f":{loop_index}:{invocation_id}"
        )
        print(f"!$######{test_stdout_tag}######$!")
        iteration = os.environ.get(
            "CODEFLASH_TEST_ITERATION", "0"
        )
        db_path = get_run_tmp_file(
            Path(f"test_return_values_{iteration}.sqlite")
        )
        codeflash_con = sqlite3.connect(db_path)
        codeflash_cur = codeflash_con.cursor()
        codeflash_cur.execute(
            "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)"
        )
        exception = None
        counter = loop.time()
        gc.disable()
        try:
            ret = func(*args, **kwargs)
            counter = loop.time()
            return_value = await ret
            codeflash_duration = int(
                (loop.time() - counter) * 1_000_000_000
            )
        except Exception as e:
            codeflash_duration = int(
                (loop.time() - counter) * 1_000_000_000
            )
            exception = e
        finally:
            gc.enable()
        print(f"!######{test_stdout_tag}######!")
        pickled_return_value = (
            pickle.dumps(exception) if exception
            else pickle.dumps(
                (args, kwargs, return_value)
            )
        )
        codeflash_cur.execute(
            "INSERT INTO test_results VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                test_module_name,
                test_class_name,
                test_name,
                function_name,
                loop_index,
                invocation_id,
                codeflash_duration,
                pickled_return_value,
                "function_call",
            ),
        )
        codeflash_con.commit()
        codeflash_con.close()
        if exception:
            raise exception
        return return_value
    return async_wrapper


def codeflash_performance_async(func):
    @wraps(func)
    async def async_wrapper(*args, **kwargs):
        loop = asyncio.get_running_loop()
        function_name = func.__name__
        line_id = os.environ["CODEFLASH_CURRENT_LINE_ID"]
        loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
        (test_module_name, test_class_name,
         test_name) = extract_test_context_from_env()
        test_id = (
            f"{test_module_name}:{test_class_name}"
            f":{test_name}:{line_id}:{loop_index}"
        )
        if not hasattr(async_wrapper, "index"):
            async_wrapper.index = {}
        if test_id in async_wrapper.index:
            async_wrapper.index[test_id] += 1
        else:
            async_wrapper.index[test_id] = 0
        codeflash_test_index = async_wrapper.index[test_id]
        invocation_id = f"{line_id}_{codeflash_test_index}"
        class_prefix = (
            (test_class_name + ".") if test_class_name else ""
        )
        test_stdout_tag = (
            f"{test_module_name}:{class_prefix}"
            f"{test_name}:{function_name}"
            f":{loop_index}:{invocation_id}"
        )
        print(f"!$######{test_stdout_tag}######$!")
        exception = None
        counter = loop.time()
        gc.disable()
        try:
            ret = func(*args, **kwargs)
            counter = loop.time()
            return_value = await ret
            codeflash_duration = int((loop.time() - counter) * 1_000_000_000)
        except Exception as e:
            codeflash_duration = int((loop.time() - counter) * 1_000_000_000)
            exception = e
        finally:
            gc.enable()
        print(f"!######{test_stdout_tag}:{codeflash_duration}######!")
        if exception:
            raise exception
        return return_value
    return async_wrapper


def codeflash_concurrency_async(func):
    @wraps(func)
    async def async_wrapper(*args, **kwargs):
        function_name = func.__name__
        concurrency_factor = int(os.environ.get(
            "CODEFLASH_CONCURRENCY_FACTOR", "10"
        ))
        test_module_name = os.environ.get("CODEFLASH_TEST_MODULE", "")
        test_class_name = os.environ.get("CODEFLASH_TEST_CLASS", "")
        test_function = os.environ.get("CODEFLASH_TEST_FUNCTION", "")
        loop_index = os.environ.get("CODEFLASH_LOOP_INDEX", "0")
        gc.disable()
        try:
            seq_start = time.perf_counter_ns()
            for _ in range(concurrency_factor):
                result = await func(*args, **kwargs)
            sequential_time = time.perf_counter_ns() - seq_start
        finally:
            gc.enable()
        gc.disable()
        try:
            conc_start = time.perf_counter_ns()
            tasks = [func(*args, **kwargs) for _ in range(concurrency_factor)]
            await asyncio.gather(*tasks)
            concurrent_time = time.perf_counter_ns() - conc_start
        finally:
            gc.enable()
        tag = (
            f"{test_module_name}:{test_class_name}"
            f":{test_function}:{function_name}"
            f":{loop_index}"
        )
        print(
            f"!@######CONC:{tag}"
            f":{sequential_time}:{concurrent_time}"
            f":{concurrency_factor}######@!"
        )
        return result
    return async_wrapper
"""

ASYNC_HELPER_FILENAME = "codeflash_async_wrapper.py"


def get_decorator_name_for_mode(
    mode: TestingMode,
) -> str:
    """Return the async decorator function name for the given testing mode."""
    if mode == TestingMode.BEHAVIOR:
        return "codeflash_behavior_async"
    if mode == TestingMode.CONCURRENCY:
        return "codeflash_concurrency_async"
    return "codeflash_performance_async"


def write_async_helper_file(
    target_dir: Path,
) -> Path:
    """Write the async decorator helper file to the target directory."""
    helper_path = target_dir / ASYNC_HELPER_FILENAME
    if not helper_path.exists():
        helper_path.write_text(ASYNC_HELPER_INLINE_CODE, "utf-8")
    return helper_path


def inject_async_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 profiling for async function calls in a test file."""
    with test_path.open(encoding="utf8") as f:
        test_code = f.read()

    try:
        tree = ast.parse(test_code)
    except SyntaxError:
        log.exception(
            "Syntax error in code in file - %s",
            test_path,
        )
        return False, None

    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

    async_instrumenter = AsyncCallInstrumenter(
        func,
        test_module_path,
        call_positions,
        mode=mode,
    )
    tree = async_instrumenter.visit(tree)

    if not async_instrumenter.did_instrument:
        return False, None

    new_imports = [ast.Import(names=[ast.alias(name="os")])]
    tree.body = [*new_imports, *tree.body]
    return True, sort_imports(ast.unparse(tree), float_to_top=True)


def add_async_decorator_to_function(
    source_path: Path,
    function: FunctionToOptimize,
    mode: TestingMode = TestingMode.BEHAVIOR,
    project_root: Path | None = None,
) -> tuple[bool, dict[Path, str]]:
    """Add an async instrumentation decorator to *function*.

    Writes the async helper file and adds the appropriate import
    and decorator.  Returns ``(True, originals)`` if the decorator
    was added, where *originals* maps each modified file to its
    content before modification.  Callers should pass *originals*
    to :func:`revert_instrumented_files` when done.
    """
    if not function.is_async:
        return False, {}

    try:
        with source_path.open(encoding="utf8") as f:
            source_code = f.read()

        module = cst.parse_module(source_code)
        decorator_transformer = AsyncDecoratorAdder(function, mode)
        module = module.visit(decorator_transformer)

        if decorator_transformer.added_decorator:
            helper_dir = (
                project_root
                if project_root is not None
                else source_path.parent
            )
            write_async_helper_file(helper_dir)
            decorator_name = get_decorator_name_for_mode(mode)
            import_node = cst.parse_statement(
                f"from codeflash_async_wrapper import {decorator_name}"
            )
            module = module.with_changes(
                body=[import_node, *list(module.body)]
            )

        modified_code = sort_imports(code=module.code, float_to_top=True)
    except Exception:
        log.exception(
            "Error adding async decorator to function %s",
            function.qualified_name,
        )
        return False, {}
    else:
        if decorator_transformer.added_decorator:
            originals: dict[Path, str] = {source_path: source_code}
            with source_path.open("w", encoding="utf8") as f:
                f.write(modified_code)
            return True, originals
        return False, {}
