from __future__ import annotations

from pathlib import Path
from unittest.mock import patch

import pytest

from codeflash_python._model import FunctionToOptimize
from codeflash_python.pipeline._orchestrator import (
    rank_by_dependency_count,
    rank_functions_globally,
)
from codeflash_python.test_discovery.discovery import existing_unit_test_count
from codeflash_python.test_discovery.models import (
    CodePosition,
    FunctionCalledInTest,
    TestsInFile,
    TestType,
)


def make_func(name: str, project_root: Path) -> FunctionToOptimize:
    """Create a minimal FunctionToOptimize for testing."""
    return FunctionToOptimize(
        function_name=name, file_path=project_root / "mod.py"
    )


def make_test(
    test_type: TestType, test_name: str = "test_something"
) -> FunctionCalledInTest:
    """Create a minimal FunctionCalledInTest for testing."""
    return FunctionCalledInTest(
        tests_in_file=TestsInFile(
            test_file=Path("/tests/test_mod.py"),
            test_class=None,
            test_function=test_name,
            test_type=test_type,
        ),
        position=CodePosition(line_no=1, col_no=0),
    )


def build_test_count_cache(
    funcs: list[FunctionToOptimize],
    project_root: Path,
    function_to_tests: dict[str, set[FunctionCalledInTest]],
) -> dict[tuple[Path, str], int]:
    """Build a test-count cache from a list of functions and their tests."""
    return {
        (func.file_path, func.qualified_name): existing_unit_test_count(
            func, project_root, function_to_tests
        )
        for func in funcs
    }


@pytest.fixture
def project_root(tmp_path: Path) -> Path:
    """Create a temporary project root with a dummy module."""
    root = tmp_path / "project"
    root.mkdir()
    (root / "mod.py").write_text(
        "def foo(): pass\ndef bar(): pass\ndef baz(): pass\n"
    )
    return root


def test_no_tests(project_root: Path) -> None:
    """existing_unit_test_count returns 0 when there are no tests."""
    func = make_func("foo", project_root)
    assert 0 == existing_unit_test_count(func, project_root, {})


def test_no_matching_key(project_root: Path) -> None:
    """existing_unit_test_count returns 0 when the key doesn't match."""
    func = make_func("foo", project_root)
    tests = {"other_module.bar": {make_test(TestType.EXISTING_UNIT_TEST)}}
    assert 0 == existing_unit_test_count(func, project_root, tests)


def test_only_replay_tests(project_root: Path) -> None:
    """existing_unit_test_count returns 0 for replay-only tests."""
    func = make_func("foo", project_root)
    key = func.qualified_name_with_modules_from_root(project_root)
    tests = {key: {make_test(TestType.REPLAY_TEST)}}
    assert 0 == existing_unit_test_count(func, project_root, tests)


def test_single_existing_test(project_root: Path) -> None:
    """existing_unit_test_count returns 1 for a single existing test."""
    func = make_func("foo", project_root)
    key = func.qualified_name_with_modules_from_root(project_root)
    tests = {key: {make_test(TestType.EXISTING_UNIT_TEST)}}
    assert 1 == existing_unit_test_count(func, project_root, tests)


def test_multiple_existing_tests(project_root: Path) -> None:
    """existing_unit_test_count counts all existing unit tests."""
    func = make_func("foo", project_root)
    key = func.qualified_name_with_modules_from_root(project_root)
    tests = {
        key: {
            make_test(TestType.EXISTING_UNIT_TEST, "test_one"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_two"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_three"),
        }
    }
    assert 3 == existing_unit_test_count(func, project_root, tests)


def test_mixed_test_types(project_root: Path) -> None:
    """existing_unit_test_count only counts EXISTING_UNIT_TEST entries."""
    func = make_func("foo", project_root)
    key = func.qualified_name_with_modules_from_root(project_root)
    tests = {
        key: {
            make_test(TestType.EXISTING_UNIT_TEST, "test_one"),
            make_test(TestType.REPLAY_TEST, "test_replay"),
            make_test(TestType.GENERATED_REGRESSION, "test_gen"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_two"),
        }
    }
    assert 2 == existing_unit_test_count(func, project_root, tests)


def test_truthiness_for_boolean_usage(project_root: Path) -> None:
    """existing_unit_test_count is falsy for 0 and truthy otherwise."""
    func = make_func("foo", project_root)
    key = func.qualified_name_with_modules_from_root(project_root)
    assert not existing_unit_test_count(func, project_root, {})
    assert existing_unit_test_count(
        func, project_root, {key: {make_test(TestType.EXISTING_UNIT_TEST)}}
    )


def test_functions_with_more_tests_rank_higher(project_root: Path) -> None:
    """Functions with more existing unit tests sort first."""
    funcs = [make_func(name, project_root) for name in ("foo", "bar", "baz")]
    function_to_tests: dict[str, set[FunctionCalledInTest]] = {
        funcs[0].qualified_name_with_modules_from_root(project_root): {
            make_test(TestType.EXISTING_UNIT_TEST, "test_one")
        },
        funcs[1].qualified_name_with_modules_from_root(project_root): {
            make_test(TestType.EXISTING_UNIT_TEST, "test_one"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_two"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_three"),
        },
        # baz has no tests
    }

    ranked = sorted(
        funcs,
        key=lambda f: (
            -existing_unit_test_count(f, project_root, function_to_tests)
        ),
    )

    assert "bar" == ranked[0].function_name  # 3 tests
    assert "foo" == ranked[1].function_name  # 1 test
    assert "baz" == ranked[2].function_name  # 0 tests


def test_stable_sort_preserves_order_for_equal_counts(
    project_root: Path,
) -> None:
    """Equal test counts preserve the original insertion order."""
    funcs = [make_func(name, project_root) for name in ("foo", "bar", "baz")]
    function_to_tests: dict[str, set[FunctionCalledInTest]] = {
        f.qualified_name_with_modules_from_root(project_root): {
            make_test(TestType.EXISTING_UNIT_TEST)
        }
        for f in funcs
    }

    ranked = sorted(
        funcs,
        key=lambda f: (
            -existing_unit_test_count(f, project_root, function_to_tests)
        ),
    )

    assert ["foo", "bar", "baz"] == [f.function_name for f in ranked]


def test_parametrized_tests_deduplication(project_root: Path) -> None:
    """Parametrized test variants are deduplicated to a single base name."""
    func = make_func("foo", project_root)
    key = func.qualified_name_with_modules_from_root(project_root)
    tests = {
        key: {
            make_test(TestType.EXISTING_UNIT_TEST, "test_foo[0]"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_foo[1]"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_foo[2]"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_bar"),
        }
    }
    assert 2 == existing_unit_test_count(func, project_root, tests)


def test_trace_ranking_keeps_addressable_time_primary_over_test_count(
    project_root: Path, tmp_path: Path
) -> None:
    """Addressable time is the primary sort key; test count is secondary."""
    funcs = [make_func(name, project_root) for name in ("foo", "bar", "baz")]
    trace_file = tmp_path / "trace.db"
    trace_file.touch()

    ranked_functions = [funcs[0], funcs[1], funcs[2]]
    addressable_times = {"foo": 100.0, "bar": 20.0, "baz": 5.0}
    function_to_tests: dict[str, set[FunctionCalledInTest]] = {
        funcs[1].qualified_name_with_modules_from_root(project_root): {
            make_test(TestType.EXISTING_UNIT_TEST, "test_one"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_two"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_three"),
        }
    }

    class FakeRanker:
        """Stub FunctionRanker that returns pre-set rankings."""

        def __init__(self, _trace_file: Path) -> None:
            pass

        def rank_functions(
            self, _functions: list[FunctionToOptimize]
        ) -> list[FunctionToOptimize]:
            """Return the pre-set ranked functions."""
            return ranked_functions

        def get_function_addressable_time(
            self, function: FunctionToOptimize
        ) -> float:
            """Return the pre-set addressable time."""
            return addressable_times[function.function_name]

    with patch(
        "codeflash_python.analysis._function_ranking.FunctionRanker",
        FakeRanker,
    ):
        ranked = rank_functions_globally(
            {project_root / "mod.py": funcs},
            trace_file,
            test_count_cache=build_test_count_cache(
                funcs, project_root, function_to_tests
            ),
        )

    assert ["foo", "bar", "baz"] == [func.function_name for _, func in ranked]


def test_trace_ranking_uses_test_count_as_tiebreaker(
    project_root: Path, tmp_path: Path
) -> None:
    """When addressable time is equal, test count breaks the tie."""
    funcs = [make_func(name, project_root) for name in ("foo", "bar", "baz")]
    trace_file = tmp_path / "trace.db"
    trace_file.touch()

    ranked_functions = [funcs[0], funcs[1], funcs[2]]
    addressable_times = {"foo": 100.0, "bar": 100.0, "baz": 5.0}
    function_to_tests: dict[str, set[FunctionCalledInTest]] = {
        funcs[0].qualified_name_with_modules_from_root(project_root): {
            make_test(TestType.EXISTING_UNIT_TEST, "test_one")
        },
        funcs[1].qualified_name_with_modules_from_root(project_root): {
            make_test(TestType.EXISTING_UNIT_TEST, "test_one"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_two"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_three"),
        },
    }

    class FakeRanker:
        """Stub FunctionRanker that returns pre-set rankings."""

        def __init__(self, _trace_file: Path) -> None:
            pass

        def rank_functions(
            self, _functions: list[FunctionToOptimize]
        ) -> list[FunctionToOptimize]:
            """Return the pre-set ranked functions."""
            return ranked_functions

        def get_function_addressable_time(
            self, function: FunctionToOptimize
        ) -> float:
            """Return the pre-set addressable time."""
            return addressable_times[function.function_name]

    with patch(
        "codeflash_python.analysis._function_ranking.FunctionRanker",
        FakeRanker,
    ):
        ranked = rank_functions_globally(
            {project_root / "mod.py": funcs},
            trace_file,
            test_count_cache=build_test_count_cache(
                funcs, project_root, function_to_tests
            ),
        )

    assert ["bar", "foo", "baz"] == [func.function_name for _, func in ranked]


def test_dependency_count_ranking_keeps_callee_count_primary(
    project_root: Path,
) -> None:
    """Callee count is the primary sort key for dependency ranking."""
    funcs = [make_func(name, project_root) for name in ("foo", "bar")]
    function_to_tests: dict[str, set[FunctionCalledInTest]] = {
        funcs[1].qualified_name_with_modules_from_root(project_root): {
            make_test(TestType.EXISTING_UNIT_TEST, "test_one"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_two"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_three"),
        }
    }

    class FakeResolver:
        """Stub call graph with pre-set callee counts."""

        def count_callees_per_function(
            self, _mapping: dict[Path, set[str]]
        ) -> dict[tuple[Path, str], int]:
            """Return pre-set callee counts."""
            return {
                (project_root / "mod.py", "foo"): 5,
                (project_root / "mod.py", "bar"): 1,
            }

    ranked = rank_by_dependency_count(
        [
            (project_root / "mod.py", funcs[0]),
            (project_root / "mod.py", funcs[1]),
        ],
        FakeResolver(),
        test_count_cache=build_test_count_cache(
            funcs, project_root, function_to_tests
        ),
    )

    assert ["foo", "bar"] == [func.function_name for _, func in ranked]


def test_dependency_count_ranking_uses_test_count_as_tiebreaker(
    project_root: Path,
) -> None:
    """When callee counts are equal, test count breaks the tie."""
    funcs = [make_func(name, project_root) for name in ("foo", "bar")]
    function_to_tests: dict[str, set[FunctionCalledInTest]] = {
        funcs[0].qualified_name_with_modules_from_root(project_root): {
            make_test(TestType.EXISTING_UNIT_TEST, "test_one")
        },
        funcs[1].qualified_name_with_modules_from_root(project_root): {
            make_test(TestType.EXISTING_UNIT_TEST, "test_one"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_two"),
            make_test(TestType.EXISTING_UNIT_TEST, "test_three"),
        },
    }

    class FakeResolver:
        """Stub call graph with equal callee counts."""

        def count_callees_per_function(
            self, _mapping: dict[Path, set[str]]
        ) -> dict[tuple[Path, str], int]:
            """Return pre-set callee counts."""
            return {
                (project_root / "mod.py", "foo"): 2,
                (project_root / "mod.py", "bar"): 2,
            }

    ranked = rank_by_dependency_count(
        [
            (project_root / "mod.py", funcs[0]),
            (project_root / "mod.py", funcs[1]),
        ],
        FakeResolver(),
        test_count_cache=build_test_count_cache(
            funcs, project_root, function_to_tests
        ),
    )

    assert ["bar", "foo"] == [func.function_name for _, func in ranked]
