"""Tests for high-level pipeline orchestrator (stage 23c)."""

from __future__ import annotations

import textwrap
from typing import TYPE_CHECKING
from unittest.mock import MagicMock

if TYPE_CHECKING:
    from pathlib import Path

from unittest.mock import patch

from codeflash_python._model import FunctionToOptimize
from codeflash_python.pipeline._orchestrator import (
    cleanup_paths,
    find_leftover_instrumented_test_files,
    prepare_module_for_optimization,
    rank_by_dependency_count,
    rank_functions_globally,
    run_benchmarks,
)


def _make_func(
    file_path: Path,
    name: str,
    line: int = 1,
) -> FunctionToOptimize:
    """Create a minimal FunctionToOptimize for testing."""
    return FunctionToOptimize(
        function_name=name,
        file_path=file_path,
        parents=[],
        starting_line=line,
        ending_line=line + 5,
    )


class TestCleanupPaths:
    """Tests for cleanup_paths."""

    def test_removes_file(self, tmp_path: Path) -> None:
        """Existing files are deleted."""
        f = tmp_path / "temp.txt"
        f.write_text("data")
        cleanup_paths([f])
        assert not f.exists()

    def test_removes_directory(self, tmp_path: Path) -> None:
        """Existing directories are removed recursively."""
        d = tmp_path / "subdir"
        d.mkdir()
        (d / "child.txt").write_text("x")
        cleanup_paths([d])
        assert not d.exists()

    def test_ignores_none(self) -> None:
        """None entries are silently skipped."""
        cleanup_paths([None])

    def test_ignores_missing(self, tmp_path: Path) -> None:
        """Non-existent paths are silently skipped."""
        cleanup_paths([tmp_path / "nonexistent"])


class TestFindLeftoverInstrumentedTestFiles:
    """Tests for find_leftover_instrumented_test_files."""

    def test_finds_perf_test(self, tmp_path: Path) -> None:
        """Matches test_*__perf_test_0.py files."""
        f = tmp_path / "test_foo__perf_test_0.py"
        f.write_text("")
        result = find_leftover_instrumented_test_files(tmp_path)
        assert f in result

    def test_finds_perfinstrumented(self, tmp_path: Path) -> None:
        """Matches test_*__perfinstrumented.py files."""
        f = tmp_path / "test_bar__perfinstrumented.py"
        f.write_text("")
        result = find_leftover_instrumented_test_files(tmp_path)
        assert f in result

    def test_finds_perfonlyinstrumented(self, tmp_path: Path) -> None:
        """Matches test_*__perfonlyinstrumented.py files."""
        f = tmp_path / "test_baz__perfonlyinstrumented.py"
        f.write_text("")
        result = find_leftover_instrumented_test_files(tmp_path)
        assert f in result

    def test_finds_unit_test(self, tmp_path: Path) -> None:
        """Matches test_*__unit_test_0.py files."""
        f = tmp_path / "test_qux__unit_test_0.py"
        f.write_text("")
        result = find_leftover_instrumented_test_files(tmp_path)
        assert f in result

    def test_ignores_normal_test(self, tmp_path: Path) -> None:
        """Normal test files are not matched."""
        f = tmp_path / "test_foo.py"
        f.write_text("")
        result = find_leftover_instrumented_test_files(tmp_path)
        assert f not in result

    def test_empty_directory(self, tmp_path: Path) -> None:
        """Empty directory returns no results."""
        result = find_leftover_instrumented_test_files(tmp_path)
        assert [] == result


class TestPrepareModuleForOptimization:
    """Tests for prepare_module_for_optimization."""

    def test_valid_module(self, tmp_path: Path) -> None:
        """A valid module returns validated code and AST."""
        src = tmp_path / "sample.py"
        src.write_text(
            textwrap.dedent("""\
                def foo():
                    return 42
            """)
        )
        result = prepare_module_for_optimization(src, tmp_path)
        assert result is not None
        validated, module_ast = result
        assert src in validated

    def test_syntax_error_returns_none(self, tmp_path: Path) -> None:
        """A module with a syntax error returns None."""
        src = tmp_path / "bad.py"
        src.write_text("def foo(:\n")
        result = prepare_module_for_optimization(src, tmp_path)
        assert result is None


class TestRankFunctionsGlobally:
    """Tests for rank_functions_globally."""

    def test_no_trace_no_graph(self, tmp_path: Path) -> None:
        """Without trace or graph, returns original order."""
        f = tmp_path / "mod.py"
        func_a = _make_func(f, "a", line=1)
        func_b = _make_func(f, "b", line=10)
        file_to_funcs = {f: [func_a, func_b]}

        result = rank_functions_globally(file_to_funcs)
        assert [(f, func_a), (f, func_b)] == result

    def test_no_trace_with_graph(self, tmp_path: Path) -> None:
        """Without trace but with graph, uses dependency ranking."""
        f = tmp_path / "mod.py"
        func_a = _make_func(f, "a", line=1)
        func_b = _make_func(f, "b", line=10)

        mock_graph = MagicMock()
        mock_graph.count_callees_per_function.return_value = {
            (f, "a"): 1,
            (f, "b"): 5,
        }

        result = rank_functions_globally(
            {f: [func_a, func_b]},
            call_graph=mock_graph,
        )
        # b has more callees, so it should be first
        assert result[0][1].qualified_name == "b"
        assert result[1][1].qualified_name == "a"

    def test_nonexistent_trace_file(self, tmp_path: Path) -> None:
        """A non-existent trace file falls back to original order."""
        f = tmp_path / "mod.py"
        func_a = _make_func(f, "a")
        result = rank_functions_globally(
            {f: [func_a]},
            trace_file_path=tmp_path / "missing.trace",
        )
        assert [(f, func_a)] == result


class TestRankByDependencyCount:
    """Tests for rank_by_dependency_count."""

    def test_sorts_by_callee_count(self, tmp_path: Path) -> None:
        """Functions with more callees rank higher."""
        f = tmp_path / "mod.py"
        func_a = _make_func(f, "a", line=1)
        func_b = _make_func(f, "b", line=10)
        func_c = _make_func(f, "c", line=20)

        mock_graph = MagicMock()
        mock_graph.count_callees_per_function.return_value = {
            (f, "a"): 2,
            (f, "b"): 10,
            (f, "c"): 5,
        }

        result = rank_by_dependency_count(
            [(f, func_a), (f, func_b), (f, func_c)],
            mock_graph,
        )
        names = [func.qualified_name for _, func in result]
        assert ["b", "c", "a"] == names

    def test_preserves_order_on_tie(self, tmp_path: Path) -> None:
        """Equal callee counts preserve original order."""
        f = tmp_path / "mod.py"
        func_a = _make_func(f, "a", line=1)
        func_b = _make_func(f, "b", line=10)

        mock_graph = MagicMock()
        mock_graph.count_callees_per_function.return_value = {
            (f, "a"): 3,
            (f, "b"): 3,
        }

        result = rank_by_dependency_count(
            [(f, func_a), (f, func_b)],
            mock_graph,
        )
        names = [func.qualified_name for _, func in result]
        assert ["a", "b"] == names

    def test_empty_input(self) -> None:
        """Empty input returns empty result."""
        mock_graph = MagicMock()
        mock_graph.count_callees_per_function.return_value = {}
        result = rank_by_dependency_count([], mock_graph)
        assert [] == result


_BENCH_MOD = "codeflash_python.benchmarking._benchmarking"
_BENCH_PLUGIN = (
    "codeflash_python.benchmarking._benchmark_plugin.CodeFlashBenchmarkPlugin"
)
_TRACE_RUNNER = (
    "codeflash_python.testing._subprocess_runners.trace_benchmarks_pytest"
)


class TestRunBenchmarks:
    """Tests for run_benchmarks."""

    def test_no_benchmarks_returns_empty(
        self,
        tmp_path: Path,
    ) -> None:
        """When replay_count is 0, returns empty dicts."""
        src = tmp_path / "mod.py"
        src.write_text("def foo(): pass\n")
        func = _make_func(src, "foo")
        benchmarks_root = tmp_path / "benchmarks"
        benchmarks_root.mkdir()

        with (
            patch(
                f"{_BENCH_MOD}.instrument_codeflash_trace_decorator",
            ),
            patch(_TRACE_RUNNER),
            patch(
                f"{_BENCH_MOD}.generate_replay_test",
                return_value=0,
            ),
        ):
            fn_timings, total_timings, replay_dir = run_benchmarks(
                {src: [func]},
                benchmarks_root,
                tmp_path / "tests",
                tmp_path,
            )

        assert {} == fn_timings
        assert {} == total_timings
        assert replay_dir is not None

    def test_restores_source_on_error(
        self,
        tmp_path: Path,
    ) -> None:
        """Original source is restored even when tracing fails."""
        src = tmp_path / "mod.py"
        original = "def foo(): return 42\n"
        src.write_text(original)
        func = _make_func(src, "foo")
        benchmarks_root = tmp_path / "benchmarks"
        benchmarks_root.mkdir()

        with (
            patch(
                f"{_BENCH_MOD}.instrument_codeflash_trace_decorator",
                side_effect=lambda _: src.write_text("MODIFIED"),
            ),
            patch(
                _TRACE_RUNNER,
                side_effect=RuntimeError("boom"),
            ),
        ):
            run_benchmarks(
                {src: [func]},
                benchmarks_root,
                tmp_path / "tests",
                tmp_path,
            )

        assert original == src.read_text()

    def test_calls_plugin_on_success(
        self,
        tmp_path: Path,
    ) -> None:
        """Timing extraction is called when replays are generated."""
        src = tmp_path / "mod.py"
        src.write_text("def foo(): pass\n")
        func = _make_func(src, "foo")
        benchmarks_root = tmp_path / "benchmarks"
        benchmarks_root.mkdir()

        mock_plugin = MagicMock()
        mock_plugin.get_function_benchmark_timings.return_value = {
            "mod.foo": {"bench_key": 100},
        }
        mock_plugin.get_benchmark_timings.return_value = {
            "bench_key": 200,
        }

        with (
            patch(
                f"{_BENCH_MOD}.instrument_codeflash_trace_decorator",
            ),
            patch(_TRACE_RUNNER),
            patch(
                f"{_BENCH_MOD}.generate_replay_test",
                return_value=3,
            ),
            patch(_BENCH_PLUGIN, mock_plugin),
        ):
            fn_timings, total_timings, replay_dir = run_benchmarks(
                {src: [func]},
                benchmarks_root,
                tmp_path / "tests",
                tmp_path,
            )

        assert {"mod.foo": {"bench_key": 100}} == fn_timings
        assert {"bench_key": 200} == total_timings
        assert replay_dir is not None
        mock_plugin.get_function_benchmark_timings.assert_called_once()
        mock_plugin.get_benchmark_timings.assert_called_once()
