from __future__ import annotations

import textwrap
from pathlib import Path

import libcst as cst

from codeflash_python._model import (
    FunctionParent,
    FunctionSource,
    FunctionToOptimize,
)
from codeflash_python.benchmarking._line_profiling import (
    LineProfilerDecoratorAdder,
    LineProfilerImportAdder,
    ProfileEnableTransformer,
    add_decorator_imports,
    add_decorator_to_qualified_function,
    add_profile_enable,
)


def make_function(
    name: str = "target",
    file_path: str = "/dev/null",
    parents: tuple[FunctionParent, ...] = (),
) -> FunctionToOptimize:
    """Create a FunctionToOptimize for testing."""
    return FunctionToOptimize(
        function_name=name,
        file_path=Path(file_path),
        parents=parents,
    )


class TestLineProfilerDecoratorAdder:
    """LineProfilerDecoratorAdder CST transformer."""

    def test_adds_decorator_to_top_level_function(self) -> None:
        """Adds decorator to a matching top-level function."""
        code = textwrap.dedent("""\
            def target():
                pass
        """)
        module = cst.parse_module(code)
        transformer = LineProfilerDecoratorAdder("target", "profile")

        result = module.visit(transformer)

        assert "@profile" in result.code
        assert "def target():" in result.code

    def test_adds_decorator_to_class_method(self) -> None:
        """Adds decorator to a method via qualified name."""
        code = textwrap.dedent("""\
            class MyClass:
                def method(self):
                    pass
        """)
        module = cst.parse_module(code)
        transformer = LineProfilerDecoratorAdder(
            "MyClass.method", "codeflash_line_profile"
        )

        result = module.visit(transformer)

        assert "@codeflash_line_profile" in result.code
        assert "def method(self):" in result.code

    def test_does_not_add_duplicate_decorator(self) -> None:
        """Does not add if the decorator is already present."""
        code = textwrap.dedent("""\
            @profile
            def target():
                pass
        """)
        module = cst.parse_module(code)
        transformer = LineProfilerDecoratorAdder("target", "profile")

        result = module.visit(transformer)

        assert result.code.count("@profile") == 1

    def test_does_not_add_to_wrong_function(self) -> None:
        """Does not add to a function with a different name."""
        code = textwrap.dedent("""\
            def other():
                pass

            def target():
                pass
        """)
        module = cst.parse_module(code)
        transformer = LineProfilerDecoratorAdder("target", "profile")

        result = module.visit(transformer)

        lines = result.code.splitlines()
        for i, line in enumerate(lines):
            if "def other():" in line:
                assert i == 0 or "@profile" not in lines[i - 1]

    def test_handles_nested_class_method(self) -> None:
        """Adds decorator to a method in a nested class."""
        code = textwrap.dedent("""\
            class Outer:
                class Inner:
                    def method(self):
                        pass
        """)
        module = cst.parse_module(code)
        transformer = LineProfilerDecoratorAdder(
            "Outer.Inner.method", "profile"
        )

        result = module.visit(transformer)

        assert "@profile" in result.code
        assert "def method(self):" in result.code

    def test_decorator_is_prepended(self) -> None:
        """New decorator is added before existing decorators."""
        code = textwrap.dedent("""\
            class MyClass:
                @staticmethod
                def method():
                    pass
        """)
        module = cst.parse_module(code)
        transformer = LineProfilerDecoratorAdder("MyClass.method", "profile")

        result = module.visit(transformer)

        profile_pos = result.code.index("@profile")
        static_pos = result.code.index("@staticmethod")
        assert profile_pos < static_pos

    def test_does_not_add_duplicate_call_decorator(self) -> None:
        """Does not add when existing decorator is a Call node."""
        code = textwrap.dedent("""\
            @profile()
            def target():
                pass
        """)
        module = cst.parse_module(code)
        transformer = LineProfilerDecoratorAdder("target", "profile")

        result = module.visit(transformer)

        assert result.code.count("profile") == 1


class TestProfileEnableTransformer:
    """ProfileEnableTransformer CST transformer."""

    def test_inserts_enable_after_aliased_import(self) -> None:
        """Inserts enable() after the aliased import."""
        code = textwrap.dedent("""\
            from line_profiler import profile as codeflash_line_profile
            def target():
                pass
        """)
        module = cst.parse_module(code)
        transformer = ProfileEnableTransformer("output.lprof")

        result = module.visit(transformer)

        assert "codeflash_line_profile.enable(" in result.code
        assert "output_prefix='output.lprof'" in result.code

    def test_does_not_modify_when_import_absent(self) -> None:
        """Returns module unchanged when no matching import."""
        code = textwrap.dedent("""\
            import os
            def target():
                pass
        """)
        module = cst.parse_module(code)
        transformer = ProfileEnableTransformer("output.lprof")

        result = module.visit(transformer)

        assert "codeflash_line_profile" not in result.code
        assert code == result.code

    def test_works_with_bare_profile_import(self) -> None:
        """Inserts enable() when import has no alias."""
        code = textwrap.dedent("""\
            from line_profiler import profile
            def target():
                pass
        """)
        module = cst.parse_module(code)
        transformer = ProfileEnableTransformer("results.lprof")

        result = module.visit(transformer)

        assert "codeflash_line_profile.enable(" in result.code
        assert "output_prefix='results.lprof'" in result.code

    def test_enable_contains_filename(self) -> None:
        """The enable() call includes the provided filename."""
        code = textwrap.dedent("""\
            from line_profiler import profile as codeflash_line_profile
            x = 1
        """)
        module = cst.parse_module(code)
        transformer = ProfileEnableTransformer("my_output")

        result = module.visit(transformer)

        assert "output_prefix='my_output'" in result.code

    def test_enable_inserted_right_after_import(self) -> None:
        """enable() appears immediately after the import."""
        code = textwrap.dedent("""\
            import os
            from line_profiler import profile as codeflash_line_profile
            x = 1
        """)
        module = cst.parse_module(code)
        transformer = ProfileEnableTransformer("out.lprof")

        result = module.visit(transformer)

        lines = result.code.splitlines()
        import_idx = next(
            i for i, ln in enumerate(lines) if "from line_profiler" in ln
        )
        enable_line = lines[import_idx + 1]
        assert "codeflash_line_profile.enable(" in enable_line


class TestLineProfilerImportAdder:
    """LineProfilerImportAdder CST transformer."""

    def test_adds_import_to_module(self) -> None:
        """Adds the import to a module that lacks it."""
        code = textwrap.dedent("""\
            def target():
                pass
        """)
        module = cst.parse_module(code)
        stmt = "from line_profiler import profile as codeflash_line_profile"
        transformer = LineProfilerImportAdder(stmt)

        result = module.visit(transformer)

        assert "from line_profiler import profile" in result.code

    def test_does_not_add_if_already_present(self) -> None:
        """Does not duplicate import if already present."""
        code = textwrap.dedent("""\
            from line_profiler import profile
            def target():
                pass
        """)
        module = cst.parse_module(code)
        stmt = "from line_profiler import profile as codeflash_line_profile"
        transformer = LineProfilerImportAdder(stmt)

        result = module.visit(transformer)

        assert result.code.count("from line_profiler") == 1

    def test_import_appears_at_beginning(self) -> None:
        """Added import is placed at the start of the module."""
        code = textwrap.dedent("""\
            import os
            def target():
                pass
        """)
        module = cst.parse_module(code)
        stmt = "from line_profiler import profile as codeflash_line_profile"
        transformer = LineProfilerImportAdder(stmt)

        result = module.visit(transformer)

        lines = result.code.strip().splitlines()
        assert "from line_profiler import profile" in lines[0]


class TestAddDecoratorToQualifiedFunction:
    """add_decorator_to_qualified_function convenience wrapper."""

    def test_adds_decorator_to_simple_function(self) -> None:
        """Adds a decorator to a top-level function."""
        code = textwrap.dedent("""\
            def target():
                pass
        """)
        module = cst.parse_module(code)

        result = add_decorator_to_qualified_function(
            module, "target", "profile"
        )

        assert "@profile" in result.code

    def test_adds_decorator_to_class_method(self) -> None:
        """Adds a decorator to a class method."""
        code = textwrap.dedent("""\
            class Calculator:
                def add(self, a, b):
                    return a + b
        """)
        module = cst.parse_module(code)

        result = add_decorator_to_qualified_function(
            module, "Calculator.add", "codeflash_line_profile"
        )

        assert "@codeflash_line_profile" in result.code
        assert "def add(self, a, b):" in result.code

    def test_returns_cst_module(self) -> None:
        """Returns a cst.Module instance."""
        code = "def f(): pass\n"
        module = cst.parse_module(code)

        result = add_decorator_to_qualified_function(module, "f", "profile")

        assert isinstance(result, cst.Module)


class TestAddProfileEnable:
    """add_profile_enable convenience wrapper."""

    def test_returns_modified_code_with_enable(self) -> None:
        """Returns code with enable() inserted."""
        code = textwrap.dedent("""\
            from line_profiler import profile as codeflash_line_profile
            def target():
                pass
        """)

        result = add_profile_enable(code, "output.lprof")

        assert "codeflash_line_profile.enable(" in result
        assert "output_prefix='output.lprof'" in result

    def test_returns_unmodified_when_no_import(self) -> None:
        """Returns original code when no matching import."""
        code = textwrap.dedent("""\
            import os
            def target():
                pass
        """)

        result = add_profile_enable(code, "output.lprof")

        assert code == result

    def test_returns_string(self) -> None:
        """The return value is a string, not a CST node."""
        code = textwrap.dedent("""\
            from line_profiler import profile as codeflash_line_profile
            x = 1
        """)

        result = add_profile_enable(code, "out.lprof")

        assert isinstance(result, str)


class TestAddDecoratorImports:
    """add_decorator_imports orchestration function."""

    def test_adds_decorator_to_function_in_file(self, tmp_path: Path) -> None:
        """Adds @codeflash_line_profile to the target function."""
        src = tmp_path / "module.py"
        src.write_text(
            textwrap.dedent("""\
                def target():
                    pass
            """),
            encoding="utf-8",
        )
        fto = make_function(name="target", file_path=str(src))

        add_decorator_imports(fto, [])

        content = src.read_text("utf-8")
        assert "@codeflash_line_profile" in content

    def test_adds_import_to_file(self, tmp_path: Path) -> None:
        """Adds the line_profiler import to the file."""
        src = tmp_path / "module.py"
        src.write_text(
            textwrap.dedent("""\
                def target():
                    pass
            """),
            encoding="utf-8",
        )
        fto = make_function(name="target", file_path=str(src))

        add_decorator_imports(fto, [])

        content = src.read_text("utf-8")
        assert "from line_profiler import profile" in content

    def test_adds_enable_to_main_file(self, tmp_path: Path) -> None:
        """Adds enable() call to the main file."""
        src = tmp_path / "module.py"
        src.write_text(
            textwrap.dedent("""\
                def target():
                    pass
            """),
            encoding="utf-8",
        )
        fto = make_function(name="target", file_path=str(src))

        add_decorator_imports(fto, [])

        content = src.read_text("utf-8")
        assert "codeflash_line_profile.enable(" in content

    def test_returns_path(self, tmp_path: Path) -> None:
        """Returns a Path for the line profiler output."""
        src = tmp_path / "module.py"
        src.write_text(
            textwrap.dedent("""\
                def target():
                    pass
            """),
            encoding="utf-8",
        )
        fto = make_function(name="target", file_path=str(src))

        result = add_decorator_imports(fto, [])

        assert isinstance(result, Path)

    def test_handles_helper_in_separate_file(self, tmp_path: Path) -> None:
        """Adds decorator and import to helper files."""
        main_src = tmp_path / "main.py"
        main_src.write_text(
            textwrap.dedent("""\
                def target():
                    pass
            """),
            encoding="utf-8",
        )
        helper_src = tmp_path / "helper.py"
        helper_src.write_text(
            textwrap.dedent("""\
                def helper_func():
                    pass
            """),
            encoding="utf-8",
        )
        fto = make_function(name="target", file_path=str(main_src))
        helper = FunctionSource(
            file_path=helper_src,
            qualified_name="helper_func",
            fully_qualified_name="helper.helper_func",
            source_code="def helper_func():\n    pass\n",
        )

        add_decorator_imports(fto, [helper])

        helper_content = helper_src.read_text("utf-8")
        assert "@codeflash_line_profile" in helper_content
        assert "from line_profiler import profile" in helper_content

    def test_enable_only_in_main_file(self, tmp_path: Path) -> None:
        """enable() is only in the main file, not helpers."""
        main_src = tmp_path / "main.py"
        main_src.write_text(
            textwrap.dedent("""\
                def target():
                    pass
            """),
            encoding="utf-8",
        )
        helper_src = tmp_path / "helper.py"
        helper_src.write_text(
            textwrap.dedent("""\
                def helper_func():
                    pass
            """),
            encoding="utf-8",
        )
        fto = make_function(name="target", file_path=str(main_src))
        helper = FunctionSource(
            file_path=helper_src,
            qualified_name="helper_func",
            fully_qualified_name="helper.helper_func",
            source_code="def helper_func():\n    pass\n",
        )

        add_decorator_imports(fto, [helper])

        main_content = main_src.read_text("utf-8")
        helper_content = helper_src.read_text("utf-8")
        assert "codeflash_line_profile.enable(" in main_content
        assert "codeflash_line_profile.enable(" not in helper_content

    def test_adds_decorator_to_class_method_in_file(
        self, tmp_path: Path
    ) -> None:
        """Adds decorator to a class method via parents."""
        src = tmp_path / "module.py"
        src.write_text(
            textwrap.dedent("""\
                class MyClass:
                    def method(self):
                        pass
            """),
            encoding="utf-8",
        )
        fto = FunctionToOptimize(
            function_name="method",
            file_path=src,
            parents=(FunctionParent(name="MyClass", type="ClassDef"),),
        )

        add_decorator_imports(fto, [])

        content = src.read_text("utf-8")
        assert "@codeflash_line_profile" in content
