"""Pytest fixture and marker CST transformations."""

from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING

import libcst as cst

if TYPE_CHECKING:
    from pathlib import Path


def has_autouse_fixture(node: cst.FunctionDef) -> bool:
    """Check if *node* has ``autouse=True`` pytest fixture."""
    for decorator in node.decorators:
        dec = decorator.decorator
        if not isinstance(dec, cst.Call):
            continue
        is_fixture = (
            isinstance(dec.func, cst.Attribute)
            and isinstance(dec.func.value, cst.Name)
            and dec.func.attr.value == "fixture"
            and dec.func.value.value == "pytest"
        ) or (isinstance(dec.func, cst.Name) and dec.func.value == "fixture")
        if is_fixture:
            for arg in dec.args:
                if (
                    arg.keyword
                    and arg.keyword.value == "autouse"
                    and isinstance(arg.value, cst.Name)
                    and arg.value.value == "True"
                ):
                    return True
    return False


class AddRequestArgument(cst.CSTTransformer):
    """Add a ``request`` parameter to autouse fixtures."""

    def leave_FunctionDef(  # noqa: N802
        self,
        original_node: cst.FunctionDef,
        updated_node: cst.FunctionDef,
    ) -> cst.FunctionDef:
        """Insert *request* param if autouse and missing."""
        if not has_autouse_fixture(original_node):
            return updated_node

        args = updated_node.params.params
        arg_names = {arg.name.value for arg in args}

        if "request" in arg_names:
            return updated_node

        request_param = cst.Param(
            name=cst.Name("request"),
        )

        if args:
            first_arg = args[0].name.value
            if first_arg in {"self", "cls"}:
                new_params = [
                    args[0],
                    request_param,
                    *list(args[1:]),
                ]
            else:
                new_params = [
                    request_param,
                    *list(args),
                ]
        else:
            new_params = [request_param]

        new_param_list = updated_node.params.with_changes(
            params=new_params,
        )
        return updated_node.with_changes(
            params=new_param_list,
        )


class PytestMarkAdder(cst.CSTTransformer):
    """Add a custom pytest mark to all test functions."""

    def __init__(self, mark_name: str) -> None:
        """Initialize with the pytest mark name to add."""
        super().__init__()
        self.mark_name = mark_name
        self.has_pytest_import = False

    def visit_Module(  # noqa: N802
        self, node: cst.Module
    ) -> None:
        """Check if pytest is already imported."""
        for statement in node.body:
            if isinstance(statement, cst.SimpleStatementLine):
                for stmt in statement.body:
                    if isinstance(stmt, cst.Import) and isinstance(
                        stmt.names,
                        Sequence,
                    ):
                        for import_alias in stmt.names:
                            if (
                                isinstance(
                                    import_alias,
                                    cst.ImportAlias,
                                )
                                and isinstance(
                                    import_alias.name,
                                    cst.Name,
                                )
                                and import_alias.name.value == "pytest"
                            ):
                                self.has_pytest_import = True

    def leave_Module(  # noqa: N802
        self,
        original_node: cst.Module,
        updated_node: cst.Module,
    ) -> cst.Module:
        """Add ``import pytest`` if not present."""
        if not self.has_pytest_import:
            import_stmt = cst.SimpleStatementLine(
                body=[
                    cst.Import(
                        names=[
                            cst.ImportAlias(
                                name=cst.Name(
                                    "pytest",
                                ),
                            ),
                        ],
                    ),
                ],
            )
            updated_node = updated_node.with_changes(
                body=[
                    import_stmt,
                    *updated_node.body,
                ],
            )
        return updated_node

    def leave_FunctionDef(  # noqa: N802
        self,
        original_node: cst.FunctionDef,
        updated_node: cst.FunctionDef,
    ) -> cst.FunctionDef:
        """Add ``@pytest.mark.<mark_name>`` to tests."""
        for decorator in updated_node.decorators:
            if self._is_pytest_mark(
                decorator.decorator,
                self.mark_name,
            ):
                return updated_node

        mark_decorator = self._create_pytest_mark()
        new_decorators = [
            *list(updated_node.decorators),
            mark_decorator,
        ]
        return updated_node.with_changes(
            decorators=new_decorators,
        )

    def _is_pytest_mark(
        self,
        decorator: cst.BaseExpression,
        mark_name: str,
    ) -> bool:
        """Return True if ``@pytest.mark.<mark_name>``."""
        if isinstance(decorator, cst.Attribute):
            if (
                isinstance(decorator.value, cst.Attribute)
                and isinstance(
                    decorator.value.value,
                    cst.Name,
                )
                and decorator.value.value.value == "pytest"
                and decorator.value.attr.value == "mark"
                and decorator.attr.value == mark_name
            ):
                return True
        elif isinstance(
            decorator,
            cst.Call,
        ) and isinstance(decorator.func, cst.Attribute):
            return self._is_pytest_mark(
                decorator.func,
                mark_name,
            )
        return False

    def _create_pytest_mark(self) -> cst.Decorator:
        """Build ``@pytest.mark.<mark_name>`` decorator."""
        mark_attr = cst.Attribute(
            value=cst.Attribute(
                value=cst.Name("pytest"),
                attr=cst.Name("mark"),
            ),
            attr=cst.Name(self.mark_name),
        )
        return cst.Decorator(decorator=mark_attr)


class AutouseFixtureModifier(cst.CSTTransformer):
    """Wrap autouse fixture bodies to skip on marker."""

    def leave_FunctionDef(  # noqa: N802
        self,
        original_node: cst.FunctionDef,
        updated_node: cst.FunctionDef,
    ) -> cst.FunctionDef:
        """Wrap body in marker check."""
        if not has_autouse_fixture(original_node):
            return updated_node

        else_block = cst.Else(body=updated_node.body)
        if_test = cst.parse_expression(
            'request.node.get_closest_marker("codeflash_no_autouse")',
        )
        yield_statement = cst.parse_statement("yield")
        if_body = cst.IndentedBlock(
            body=[yield_statement],
        )
        new_if = cst.If(
            test=if_test,
            body=if_body,
            orelse=else_block,
        )
        return updated_node.with_changes(
            body=cst.IndentedBlock(body=[new_if]),
        )


def disable_autouse(test_path: Path) -> str:
    """Modify *test_path* to disable autouse fixtures.

    Returns the original file content so it can be restored.
    """
    file_content = test_path.read_text(encoding="utf-8")
    module = cst.parse_module(file_content)
    modified_module = module.visit(
        AddRequestArgument(),
    )
    modified_module = modified_module.visit(
        AutouseFixtureModifier(),
    )
    test_path.write_text(modified_module.code, encoding="utf-8")
    return file_content


def modify_autouse_fixture(
    conftest_files: list[Path],
) -> dict[Path, str]:
    """Disable autouse fixtures in *conftest_files*.

    Returns a mapping from file path to original content.
    """
    file_content_map: dict[Path, str] = {}
    for cf_file in conftest_files:
        original_content = disable_autouse(cf_file)
        file_content_map[cf_file] = original_content
    return file_content_map


def add_custom_marker_to_all_tests(
    test_paths: list[Path],
) -> None:
    """Add ``@pytest.mark.codeflash_no_autouse`` to tests."""
    for test_path in test_paths:
        file_content = test_path.read_text(
            encoding="utf-8",
        )
        module = cst.parse_module(file_content)
        pytest_mark_adder = PytestMarkAdder(
            "codeflash_no_autouse",
        )
        modified_module = module.visit(
            pytest_mark_adder,
        )
        test_path.write_text(
            modified_module.code,
            encoding="utf-8",
        )
