"""Test generation via the Codeflash AI service.

Provides helpers for generating regression tests, merging AI-generated
tests with existing unit tests, and orchestrating the full test-generation
pipeline.
"""

from __future__ import annotations

import ast
import logging
import os
import re
from pathlib import Path
from typing import TYPE_CHECKING, Any

import attrs
import libcst as cst
from libcst import MetadataWrapper
from libcst.metadata import PositionProvider

from codeflash_core import (
    AIServiceConnectionError,
    AIServiceError,
    performance_gain,
)

from .._constants import LANGUAGE_FIELDS, LANGUAGE_VERSION

if TYPE_CHECKING:
    from codeflash_core import AIClient

    from .._model import FunctionToOptimize
    from .models import InvocationId

logger = logging.getLogger(__name__)


@attrs.frozen
class GeneratedTests:
    """A set of generated test sources for a single function."""

    generated_original_test_source: str
    instrumented_behavior_test_source: str
    instrumented_perf_test_source: str
    behavior_file_path: Path
    perf_file_path: Path
    raw_generated_test_source: str | None = None


@attrs.frozen
class GeneratedTestsList:
    """A collection of generated test sets."""

    generated_tests: tuple[GeneratedTests, ...] = ()


@attrs.frozen
class TestgenPayload:
    """Typed payload for the ``/ai/testgen`` endpoint."""

    __test__ = False

    source_code_being_tested: str
    function_to_optimize: dict[str, object]
    helper_function_names: list[str]
    module_path: str
    test_module_path: str
    test_framework: str
    test_timeout: int
    trace_id: str
    test_index: int
    language_version: str = LANGUAGE_VERSION
    is_numerical_code: bool | None = None
    is_async: bool = False
    class_name: str | None = None
    qualified_name: str = ""
    codeflash_version: str | None = None
    call_sequence: str | None = None

    def to_dict(self) -> dict[str, Any]:
        """
        Serialize to the dict expected by the AI service.
        """
        d: dict[str, Any] = {
            **LANGUAGE_FIELDS,
            "source_code_being_tested": self.source_code_being_tested,
            "function_to_optimize": self.function_to_optimize,
            "helper_function_names": self.helper_function_names,
            "module_path": self.module_path,
            "test_module_path": self.test_module_path,
            "test_framework": self.test_framework,
            "test_timeout": self.test_timeout,
            "trace_id": self.trace_id,
            "test_index": self.test_index,
            "language_version": self.language_version,
            "is_numerical_code": self.is_numerical_code,
            "is_async": self.is_async,
            "class_name": self.class_name,
            "qualified_name": self.qualified_name,
        }
        if self.codeflash_version is not None:
            d["codeflash_version"] = self.codeflash_version
        if self.call_sequence is not None:
            d["call_sequence"] = self.call_sequence
        return d


def generate_regression_tests(
    client: AIClient,
    payload: TestgenPayload,
) -> tuple[str, str | None] | None:
    """
    Call the AI service ``/ai/testgen`` endpoint to generate regression tests.

    Returns *(generated_tests, raw_generated)* or ``None`` on failure.
    """
    data = client.post("/testgen", payload.to_dict())

    generated = data.get("generated_tests", "")
    raw = data.get("raw_generated_tests")

    if not generated:
        return None

    return (generated, raw)


def generate_tests(  # noqa: PLR0913
    client: AIClient,
    source_code_being_tested: str,
    function_to_optimize: FunctionToOptimize,
    helper_function_names: list[str],
    module_path: str,
    test_framework: str,
    test_timeout: int,
    trace_id: str,
    test_index: int,
    test_path: Path,
    test_perf_path: Path,
    test_module_path: str,
    language_version: str,
    is_numerical_code: bool | None = None,  # noqa: FBT001
) -> tuple[str, str | None, Path, Path] | None:
    """
    Generate regression tests for a function via the AI service.

    Returns *(generated_source, raw_source, test_path, test_perf_path)*
    or ``None`` on failure.
    """
    payload = TestgenPayload(
        source_code_being_tested=source_code_being_tested,
        function_to_optimize=function_to_optimize.to_dict(),
        helper_function_names=helper_function_names,
        module_path=module_path,
        test_module_path=test_module_path,
        test_framework=test_framework,
        test_timeout=test_timeout,
        trace_id=trace_id,
        test_index=test_index,
        language_version=language_version,
        is_numerical_code=is_numerical_code,
        is_async=function_to_optimize.is_async,
        class_name=function_to_optimize.class_name,
        qualified_name=function_to_optimize.qualified_name,
    )

    response = generate_regression_tests(
        client=client,
        payload=payload,
    )

    if response is None:
        return None

    generated, raw = response
    return (generated, raw, test_path, test_perf_path)


def review_generated_tests(
    client: AIClient,
    payload: dict[str, Any],
) -> list[dict[str, Any]]:
    """
    Review generated tests via the AI service.

    Returns a list of review dicts, or ``[]`` on failure.
    """
    try:
        data = client.post("/testgen_review", payload)
    except (AIServiceError, AIServiceConnectionError):
        return []
    reviews: list[dict[str, Any]] = data.get("reviews", [])
    return reviews


def repair_generated_tests(
    client: AIClient,
    payload: dict[str, Any],
) -> str | None:
    """
    Repair generated tests via the AI service.

    Returns the repaired test source or ``None`` on failure.
    """
    try:
        data = client.post("/testgen_repair", payload)
    except (AIServiceError, AIServiceConnectionError):
        return None
    generated: str = data.get("generated_tests", "")
    if not generated:
        return None
    return generated


class ModifyInspiredTests(ast.NodeTransformer):
    """Extract imports from inspired tests."""

    def __init__(
        self,
        import_list: list[ast.stmt],
    ) -> None:
        """Initialize with an import list to populate."""
        self.import_list = import_list

    def visit_Import(self, node: ast.Import) -> None:
        """Extract top-level import statements."""
        self.import_list.append(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        """Extract top-level from-import statements."""
        self.import_list.append(node)


def delete_multiple_if_name_main(test_ast: ast.Module) -> ast.Module:
    """
    Remove all but the last ``if __name__ == "__main__"`` block.
    """
    if_indexes: list[int] = []
    for index, node in enumerate(test_ast.body):
        if (
            isinstance(node, ast.If)
            and isinstance(node.test, ast.Compare)
            and isinstance(node.test.left, ast.Name)
            and node.test.left.id == "__name__"
            and len(node.test.ops) == 1
            and isinstance(node.test.ops[0], ast.Eq)
            and len(node.test.comparators) == 1
            and isinstance(node.test.comparators[0], ast.Constant)
            and node.test.comparators[0].value == "__main__"
        ):
            if_indexes.append(index)
    for index in list(reversed(if_indexes))[1:]:
        del test_ast.body[index]
    return test_ast


def merge_unit_tests(
    unit_test_source: str,
    inspired_unit_tests: str,
) -> str:
    """Merge existing unit tests with AI-generated (inspired) tests.

    Appends generated tests with ``__inspired`` suffix.
    """
    try:
        inspired_unit_tests_ast = ast.parse(inspired_unit_tests)
        unit_test_source_ast = ast.parse(unit_test_source)
    except SyntaxError:
        return unit_test_source

    import_list: list[ast.stmt] = []
    modified_ast = ModifyInspiredTests(import_list).visit(
        inspired_unit_tests_ast
    )

    for node in ast.iter_child_nodes(modified_ast):
        if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"):
            node.name = node.name + "__inspired"

    unit_test_source_ast.body.extend(modified_ast.body)
    unit_test_source_ast.body = import_list + unit_test_source_ast.body

    return ast.unparse(unit_test_source_ast)


_NS = 1_000
_US = 1_000_000
_MS = 1_000_000_000
_THRESH_LOW = 10
_THRESH_HIGH = 100


def format_time(nanoseconds: int) -> str:
    """Format nanoseconds into a human-readable string."""
    if not isinstance(nanoseconds, int):
        msg = "Input must be an integer."
        raise TypeError(msg)
    if nanoseconds < 0:
        msg = "Input must be a positive integer."
        raise ValueError(msg)

    if nanoseconds < _NS:
        return f"{nanoseconds}ns"
    if nanoseconds < _US:
        value = nanoseconds / _NS
        return _format_value(value, "μs")
    if nanoseconds < _MS:
        value = nanoseconds / _US
        return _format_value(value, "ms")
    value = nanoseconds / _MS
    return _format_value(value, "s")


def _format_value(value: float, unit: str) -> str:
    """Format a numeric value with its unit at appropriate precision."""
    if value < _THRESH_LOW:
        return f"{value:.2f}{unit}"
    if value < _THRESH_HIGH:
        return f"{value:.1f}{unit}"
    return f"{int(value)}{unit}"


def format_perf(percentage: float) -> str:
    """Format percentage with appropriate precision."""
    abs_perc = abs(percentage)
    if abs_perc >= _THRESH_HIGH:
        return f"{percentage:.0f}"
    if abs_perc >= _THRESH_LOW:
        return f"{percentage:.1f}"
    if abs_perc >= 1:
        return f"{percentage:.2f}"
    return f"{percentage:.3f}"


class CommentMapper(ast.NodeVisitor):
    """Map line numbers to runtime comparison comments for generated tests."""

    def __init__(
        self,
        test: GeneratedTests,
        original_runtimes: dict[str, int],
        optimized_runtimes: dict[str, int],
    ) -> None:
        """Initialize with test data and runtime dictionaries."""
        self.results: dict[int, str] = {}
        self.test: GeneratedTests = test
        self.original_runtimes = original_runtimes
        self.optimized_runtimes = optimized_runtimes
        self.abs_path = test.behavior_file_path.with_suffix("")
        self.context_stack: list[str] = []

    def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
        """Visit test methods inside a class and map their line numbers."""
        self.context_stack.append(node.name)
        for inner_node in node.body:
            if isinstance(inner_node, ast.FunctionDef):
                self.visit_FunctionDef(inner_node)
            elif isinstance(inner_node, ast.AsyncFunctionDef):
                self.visit_AsyncFunctionDef(inner_node)
        self.context_stack.pop()
        return node

    def get_comment(self, match_key: str) -> str:
        """Build a runtime comparison comment string for the given key."""
        original_time = self.original_runtimes[match_key]
        optimized_time = self.optimized_runtimes[match_key]
        perf_gain = format_perf(
            abs(
                performance_gain(
                    original_runtime_ns=original_time,
                    optimized_runtime_ns=optimized_time,
                )
                * 100
            )
        )
        status = "slower" if optimized_time > original_time else "faster"
        orig = format_time(original_time)
        opt = format_time(optimized_time)
        return f"# {orig} -> {opt} ({perf_gain}% {status})"

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
        """Map runtime comments for a synchronous test function."""
        self._process_function_def_common(node)
        return node

    def visit_AsyncFunctionDef(
        self, node: ast.AsyncFunctionDef
    ) -> ast.AsyncFunctionDef:
        """Map runtime comments for an asynchronous test function."""
        self._process_function_def_common(node)
        return node

    def _process_function_def_common(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef
    ) -> None:
        """Walk the function body and record runtime comments per line."""
        self.context_stack.append(node.name)
        i = len(node.body) - 1
        test_qualified_name = ".".join(self.context_stack)
        key = test_qualified_name + "#" + str(self.abs_path)
        while i >= 0:
            line_node = node.body[i]
            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]
                    nodes_to_check = [compound_line_node]
                    nodes_to_check.extend(
                        getattr(compound_line_node, "body", [])
                    )
                    for internal_node in nodes_to_check:
                        if isinstance(internal_node, (ast.stmt, ast.Assign)):
                            inv_id = str(i) + "_" + str(j)
                            match_key = key + "#" + inv_id
                            if (
                                match_key in self.original_runtimes
                                and match_key in self.optimized_runtimes
                            ):
                                self.results[internal_node.lineno] = (
                                    self.get_comment(match_key)
                                )
                    j -= 1
            else:
                inv_id = str(i)
                match_key = key + "#" + inv_id
                if (
                    match_key in self.original_runtimes
                    and match_key in self.optimized_runtimes
                ):
                    self.results[line_node.lineno] = self.get_comment(
                        match_key
                    )
            i -= 1
        self.context_stack.pop()


def get_fn_call_linenos(
    test: GeneratedTests,
    original_runtimes: dict[str, int],
    optimized_runtimes: dict[str, int],
) -> dict[int, str]:
    """Return a mapping of line numbers to runtime comments for a test."""
    line_comment_ast_mapper = CommentMapper(
        test, original_runtimes, optimized_runtimes
    )
    source_code = test.generated_original_test_source
    tree = ast.parse(source_code)
    line_comment_ast_mapper.visit(tree)
    return line_comment_ast_mapper.results


class CommentAdder(cst.CSTTransformer):
    """Transformer that adds comments to specified lines."""

    METADATA_DEPENDENCIES = (PositionProvider,)

    def __init__(self, line_to_comments: dict[int, str]) -> None:
        """Initialize with a mapping of line numbers to comment strings."""
        self.line_to_comments = line_to_comments
        super().__init__()

    def leave_SimpleStatementLine(  # noqa: N802
        self,
        original_node: cst.SimpleStatementLine,
        updated_node: cst.SimpleStatementLine,
    ) -> cst.SimpleStatementLine:
        """Append a trailing comment if this line has a mapped runtime comment."""
        pos = self.get_metadata(PositionProvider, original_node)
        if pos and pos.start.line in self.line_to_comments:
            comment = cst.TrailingWhitespace(
                whitespace=cst.SimpleWhitespace(" "),
                comment=cst.Comment(self.line_to_comments[pos.start.line]),
            )
            return updated_node.with_changes(trailing_whitespace=comment)
        return updated_node

    def leave_SimpleStatementSuite(  # noqa: N802
        self,
        original_node: cst.SimpleStatementSuite,
        updated_node: cst.SimpleStatementSuite,
    ) -> cst.SimpleStatementSuite:
        """Append a trailing comment if this suite line has a mapped runtime comment."""
        pos = self.get_metadata(PositionProvider, original_node)
        if pos and pos.start.line in self.line_to_comments:
            comment = cst.TrailingWhitespace(
                whitespace=cst.SimpleWhitespace("  "),
                comment=cst.Comment(self.line_to_comments[pos.start.line]),
            )
            return updated_node.with_changes(trailing_whitespace=comment)
        return updated_node


def _is_python_file(file_path: Path) -> bool:
    """Check if a file is a Python file."""
    return file_path.suffix == ".py"


def unique_inv_id(
    inv_id_runtimes: dict[InvocationId, list[int]], tests_project_rootdir: Path
) -> dict[str, int]:
    """Collapse invocation runtimes into unique string-keyed minimum runtimes."""
    unique_inv_ids: dict[str, int] = {}
    for inv_id, runtimes in inv_id_runtimes.items():
        test_qualified_name = (
            inv_id.test_class_name + "." + inv_id.test_function_name  # type: ignore[operator]
            if inv_id.test_class_name
            else inv_id.test_function_name
        )

        test_module_path = inv_id.test_module_path
        if "/" in test_module_path or "\\" in test_module_path:
            abs_path = tests_project_rootdir / Path(test_module_path)
        else:
            abs_path = tests_project_rootdir / Path(
                test_module_path.replace(".", os.sep)
            ).with_suffix(".py")

        abs_path_str = str(abs_path.resolve().with_suffix(""))
        if (
            "__unit_test_" not in abs_path_str
            and "__perf_test_" not in abs_path_str
        ) or not test_qualified_name:
            continue
        key = test_qualified_name + "#" + abs_path_str
        id_parts = inv_id.iteration_id.split("_")  # type: ignore[union-attr]
        cur_invid = (
            id_parts[0]
            if len(id_parts) < 3  # noqa: PLR2004
            else "_".join(id_parts[:-1])
        )
        match_key = key + "#" + cur_invid
        if match_key not in unique_inv_ids:
            unique_inv_ids[match_key] = 0
        unique_inv_ids[match_key] += min(runtimes)
    return unique_inv_ids


def add_runtime_comments_to_generated_tests(
    generated_tests: GeneratedTestsList,
    original_runtimes: dict[InvocationId, list[int]],
    optimized_runtimes: dict[InvocationId, list[int]],
    tests_project_rootdir: Path | None = None,
) -> GeneratedTestsList:
    """Add runtime comments to generated tests."""
    original_runtimes_dict = unique_inv_id(
        original_runtimes, tests_project_rootdir or Path()
    )
    optimized_runtimes_dict = unique_inv_id(
        optimized_runtimes, tests_project_rootdir or Path()
    )
    modified_tests = []
    for test in generated_tests.generated_tests:
        is_python = _is_python_file(test.behavior_file_path)

        if is_python:
            try:
                tree = cst.parse_module(test.generated_original_test_source)
                wrapper = MetadataWrapper(tree)
                line_to_comments = get_fn_call_linenos(
                    test, original_runtimes_dict, optimized_runtimes_dict
                )
                comment_adder = CommentAdder(line_to_comments)
                modified_tree = wrapper.visit(comment_adder)
                modified_source = modified_tree.code
                modified_test = GeneratedTests(
                    generated_original_test_source=modified_source,
                    instrumented_behavior_test_source=test.instrumented_behavior_test_source,
                    instrumented_perf_test_source=test.instrumented_perf_test_source,
                    behavior_file_path=test.behavior_file_path,
                    perf_file_path=test.perf_file_path,
                )
                modified_tests.append(modified_test)
            except Exception:  # noqa: BLE001
                logger.debug("Failed to add runtime comments to test")
                modified_tests.append(test)
        else:
            modified_tests.append(test)

    return GeneratedTestsList(generated_tests=tuple(modified_tests))


def _compile_function_patterns(
    test_functions_to_remove: list[str],
) -> list[re.Pattern[str]]:
    """Compile regex patterns to match test function definitions by name."""
    return [
        re.compile(
            rf"(@pytest\.mark\.parametrize\(.*?\)\s*)?(async\s+)?def\s+{re.escape(func)}\(.*?\):.*?(?=\n(async\s+)?def\s|$)",
            re.DOTALL,
        )
        for func in test_functions_to_remove
    ]


def remove_functions_from_generated_tests(
    generated_tests: GeneratedTestsList, test_functions_to_remove: list[str]
) -> GeneratedTestsList:
    """Remove specified test functions from generated test sources."""
    function_patterns = _compile_function_patterns(test_functions_to_remove)
    new_generated_tests = []

    for gt in generated_tests.generated_tests:
        source = gt.generated_original_test_source

        for pattern in function_patterns:
            for match in pattern.finditer(source):
                if "@pytest.mark.parametrize" in match.group(0):
                    continue
                start, end = match.span()
                source = source[:start] + source[end:]
                break

        new_generated_tests.append(
            GeneratedTests(
                generated_original_test_source=source,
                instrumented_behavior_test_source=gt.instrumented_behavior_test_source,
                instrumented_perf_test_source=gt.instrumented_perf_test_source,
                behavior_file_path=gt.behavior_file_path,
                perf_file_path=gt.perf_file_path,
            )
        )

    return GeneratedTestsList(generated_tests=tuple(new_generated_tests))


def _is_trivial_statement(
    stmt: cst.BaseStatement,
) -> bool:
    """Return True if *stmt* is a ``pass`` or a bare docstring."""
    if not isinstance(stmt, cst.SimpleStatementLine):
        return False
    if len(stmt.body) != 1:
        return False
    item = stmt.body[0]
    if isinstance(item, cst.Pass):
        return True
    return isinstance(item, cst.Expr) and isinstance(
        item.value, (cst.SimpleString, cst.ConcatenatedString)
    )


class _TestFunctionRemover(cst.CSTTransformer):
    """Remove targeted function/method definitions from the CST."""

    def __init__(
        self,
        bare_names: set[str],
        qualified_names: set[str],
    ) -> None:
        self.bare_names = bare_names
        self.qualified_names = qualified_names
        self.class_stack: list[str] = []
        self.emptied_classes: set[str] = set()

    def visit_ClassDef(self, node: cst.ClassDef) -> bool:  # noqa: N802
        """Track class nesting."""
        self.class_stack.append(node.name.value)
        return True

    def leave_ClassDef(  # noqa: N802
        self,
        original_node: cst.ClassDef,
        updated_node: cst.ClassDef,
    ) -> cst.ClassDef | cst.RemovalSentinel:
        """Remove class if all meaningful body was stripped."""
        class_name = self.class_stack.pop()
        if class_name in self.emptied_classes:
            self.emptied_classes.discard(class_name)
            body = updated_node.body
            if isinstance(body, cst.IndentedBlock) and all(
                _is_trivial_statement(s) for s in body.body
            ):
                return cst.RemovalSentinel.REMOVE
        return updated_node

    def leave_FunctionDef(  # noqa: N802
        self,
        original_node: cst.FunctionDef,
        updated_node: cst.FunctionDef,
    ) -> cst.FunctionDef | cst.RemovalSentinel:
        """Remove the function if its name is in the target set."""
        fn_name = original_node.name.value
        if fn_name in self.bare_names and not self.class_stack:
            return cst.RemovalSentinel.REMOVE
        if self.class_stack:
            qualified = f"{self.class_stack[-1]}.{fn_name}"
            if qualified in self.qualified_names:
                self.emptied_classes.add(self.class_stack[-1])
                return cst.RemovalSentinel.REMOVE
        return updated_node


def remove_test_functions(
    test_source: str, functions_to_remove: list[str]
) -> str:
    """Remove specific test functions from Python test source.

    Bare names (e.g. ``"test_foo"``) match only module-level functions.
    Qualified names (e.g. ``"TestSuite.test_bar"``) match methods inside the
    named class.  When every method in a class is removed and only docstrings
    or ``pass`` remain, the entire class is removed as well.

    If *test_source* cannot be parsed, the original string is returned
    unchanged.
    """
    bare_names: set[str] = set()
    qualified_names: set[str] = set()
    for name in functions_to_remove:
        if "." in name:
            qualified_names.add(name)
        else:
            bare_names.add(name)

    try:
        tree = cst.parse_module(test_source)
        modified = tree.visit(
            _TestFunctionRemover(bare_names, qualified_names)
        )
    except Exception:  # noqa: BLE001
        return test_source
    else:
        return modified.code
