"""Tests for discover_helpers (Jedi-based helper discovery)."""

from __future__ import annotations

import textwrap
from typing import TYPE_CHECKING

from codeflash_python._model import FunctionParent, FunctionToOptimize
from codeflash_python.context.helpers import discover_helpers

if TYPE_CHECKING:
    from pathlib import Path


def helper_names(
    result: dict[Path, set[object]],
) -> set[str]:
    """Extract qualified names from a discover_helpers result."""
    return {
        fs.qualified_name  # type: ignore[union-attr]
        for sources in result.values()
        for fs in sources
    }


class TestDiscoverHelpers:
    """Tests for discover_helpers."""

    def test_direct_helper_in_same_file(
        self,
        tmp_path: Path,
    ) -> None:
        """A function calling a local helper discovers that helper."""
        src = tmp_path / "sample.py"
        src.write_text(
            textwrap.dedent("""\
                def helper():
                    return 42

                def target():
                    return helper()
            """),
        )

        fn = FunctionToOptimize(
            function_name="target",
            file_path=src,
            starting_line=4,
            ending_line=5,
        )

        result = discover_helpers(fn, tmp_path)
        assert "helper" in helper_names(result)

    def test_stdlib_not_discovered(
        self,
        tmp_path: Path,
    ) -> None:
        """Standard library calls are filtered out."""
        src = tmp_path / "sample.py"
        src.write_text(
            textwrap.dedent("""\
                import os

                def target():
                    return os.getcwd()
            """),
        )

        fn = FunctionToOptimize(
            function_name="target",
            file_path=src,
            starting_line=3,
            ending_line=4,
        )

        result = discover_helpers(fn, tmp_path)
        assert helper_names(result) == set()

    def test_no_helpers(self, tmp_path: Path) -> None:
        """A function with no calls returns no helpers."""
        src = tmp_path / "sample.py"
        src.write_text(
            textwrap.dedent("""\
                def target():
                    return 42
            """),
        )

        fn = FunctionToOptimize(
            function_name="target",
            file_path=src,
            starting_line=1,
            ending_line=2,
        )

        result = discover_helpers(fn, tmp_path)
        assert result == {}

    def test_transitive_helpers(
        self,
        tmp_path: Path,
    ) -> None:
        """Helpers-of-helpers are discovered at the second level."""
        src = tmp_path / "sample.py"
        src.write_text(
            textwrap.dedent("""\
                def deep_helper():
                    return 42

                def helper():
                    return deep_helper()

                def target():
                    return helper()
            """),
        )

        fn = FunctionToOptimize(
            function_name="target",
            file_path=src,
            starting_line=7,
            ending_line=8,
        )

        result = discover_helpers(fn, tmp_path)
        names = helper_names(result)
        assert "helper" in names
        assert "deep_helper" in names

    def test_method_discovers_helper(
        self,
        tmp_path: Path,
    ) -> None:
        """A class method calling a helper discovers that helper."""
        src = tmp_path / "sample.py"
        src.write_text(
            textwrap.dedent("""\
                def compute(x):
                    return x * 2

                class Processor:
                    def run(self):
                        return compute(21)
            """),
        )

        fn = FunctionToOptimize(
            function_name="run",
            file_path=src,
            parents=(
                FunctionParent(
                    name="Processor",
                    type="ClassDef",
                ),
            ),
            starting_line=5,
            ending_line=6,
            is_method=True,
        )

        result = discover_helpers(fn, tmp_path)
        assert "compute" in helper_names(result)

    def test_cross_file_helper(
        self,
        tmp_path: Path,
    ) -> None:
        """Helpers imported from another project file are discovered."""
        utils = tmp_path / "utils.py"
        utils.write_text(
            textwrap.dedent("""\
                def format_name(name):
                    return name.strip().title()
            """),
        )

        main = tmp_path / "main.py"
        main.write_text(
            textwrap.dedent("""\
                from utils import format_name

                def greet(name):
                    return f"Hello, {format_name(name)}!"
            """),
        )

        fn = FunctionToOptimize(
            function_name="greet",
            file_path=main,
            starting_line=3,
            ending_line=4,
        )

        result = discover_helpers(fn, tmp_path)
        assert "format_name" in helper_names(result)

    def test_recursive_call_excluded(
        self,
        tmp_path: Path,
    ) -> None:
        """Recursive calls to self are not reported as helpers."""
        src = tmp_path / "sample.py"
        src.write_text(
            textwrap.dedent("""\
                def factorial(n):
                    if n <= 1:
                        return 1
                    return n * factorial(n - 1)
            """),
        )

        fn = FunctionToOptimize(
            function_name="factorial",
            file_path=src,
            starting_line=1,
            ending_line=4,
        )

        result = discover_helpers(fn, tmp_path)
        assert "factorial" not in helper_names(result)
