"""Safety tests for get_optimized_code_for_module() fallback chain.

These tests verify the matching logic that maps AI-generated code blocks
to the correct source file, including all fallback strategies.
"""

from __future__ import annotations

from pathlib import Path

from codeflash_python.codegen._replacement import (
    get_optimized_code_for_module,
)
from codeflash_python.context.models import CodeString, CodeStringsMarkdown


def make_optimized_code(file_to_code: dict[str, str]) -> CodeStringsMarkdown:
    """Create a CodeStringsMarkdown with a given file_to_code mapping."""
    return CodeStringsMarkdown(
        code_strings=[
            CodeString(
                code=code,
                file_path=Path(path) if path != "None" else None,
            )
            for path, code in file_to_code.items()
        ],
    )


class TestGetOptimizedCodeForModule:
    """Test the fallback chain in get_optimized_code_for_module."""

    def test_exact_path_match(self) -> None:
        """When the relative path matches exactly, return that code."""
        code = make_optimized_code({"src/pkg/foo.py": "class Foo: pass"})
        result = get_optimized_code_for_module(Path("src/pkg/foo.py"), code)
        assert "class Foo: pass" == result

    def test_none_key_fallback(self) -> None:
        """When there's a single code block with 'None' key, use it."""
        code = make_optimized_code({"None": "class Foo: optimized"})
        result = get_optimized_code_for_module(Path("src/pkg/foo.py"), code)
        assert "class Foo: optimized" == result

    def test_basename_match(self) -> None:
        """When the AI returns just 'algorithms.py', match by basename."""
        code = make_optimized_code({"algorithms.py": "def fast(): pass"})
        result = get_optimized_code_for_module(
            Path("src/pkg/algorithms.py"), code
        )
        assert "def fast(): pass" == result

    def test_basename_match_with_different_prefix(self) -> None:
        """Basename match should work even with a different directory prefix."""
        code = make_optimized_code({"other/foo.py": "class Foo: v2"})
        result = get_optimized_code_for_module(Path("src/pkg/foo.py"), code)
        assert "class Foo: v2" == result

    def test_single_block_wrong_path_does_not_match(self) -> None:
        """A single code block with wrong path should NOT match for Python."""
        code = make_optimized_code({"wrong/path/bar.py": "def bar(): pass"})
        result = get_optimized_code_for_module(Path("src/foo.py"), code)
        assert "" == result

    def test_no_match_returns_empty(self) -> None:
        """When multiple blocks exist and none match, return empty string."""
        code = make_optimized_code(
            {
                "other/file1.py": "class File1: pass",
                "other/file2.py": "class File2: pass",
            }
        )
        result = get_optimized_code_for_module(Path("src/pkg/foo.py"), code)
        assert "" == result

    def test_none_key_with_multiple_blocks_no_match(self) -> None:
        """When there are multiple blocks including 'None', don't use None fallback."""
        code = make_optimized_code(
            {
                "None": "class Default: pass",
                "other/file.py": "class File: pass",
            }
        )
        result = get_optimized_code_for_module(Path("src/pkg/foo.py"), code)
        # With multiple blocks, the None-key fallback should NOT trigger
        assert "" == result
