from __future__ import annotations

from pathlib import Path

import pytest

from codeflash_python._model import FunctionToOptimize
from codeflash_python.analysis._discovery import discover_functions
from codeflash_python.analysis._extraction import extract_function_source

SOURCE = """\
def greet(name):
    return f"hello {name}"

class Formatter:
    def bold(self, text):
        return f"**{text}**"

    def italic(self, text):
        return f"*{text}*"

def no_return():
    print("side effect")
"""


@pytest.fixture(name="source_file")
def _source_file(tmp_path: Path) -> Path:
    p = tmp_path / "sample.py"
    p.write_text(SOURCE)
    return p


class TestExtractFunctionSource:
    """Tests for extract_function_source."""

    def test_top_level_function(self, source_file: Path) -> None:
        """A top-level function is extracted exactly."""
        funcs = discover_functions(SOURCE, source_file)
        greet = next(f for f in funcs if f.function_name == "greet")

        result = extract_function_source(greet)

        assert 'def greet(name):\n    return f"hello {name}"\n' == result

    def test_method(self, source_file: Path) -> None:
        """A class method is extracted with its indentation."""
        funcs = discover_functions(SOURCE, source_file)
        bold = next(f for f in funcs if f.function_name == "bold")

        result = extract_function_source(bold)

        assert "    def bold(self, text):\n" in result
        assert '        return f"**{text}**"\n' in result

    def test_all_discovered_functions_extractable(
        self,
        source_file: Path,
    ) -> None:
        """Every discovered function can be extracted without error."""
        funcs = discover_functions(SOURCE, source_file)

        for fn in funcs:
            result = extract_function_source(fn)
            assert fn.function_name in result

    def test_missing_line_numbers_raises(self) -> None:
        """Functions without line numbers cannot be extracted."""
        fn = FunctionToOptimize(
            function_name="orphan",
            file_path=Path("/dev/null"),
        )

        with pytest.raises(ValueError, match="missing line numbers"):
            extract_function_source(fn)

    def test_round_trip_with_real_file(self, tmp_path: Path) -> None:
        """Extracted source compiles to valid Python."""
        src = "def add(a, b):\n    return a + b\n"
        p = tmp_path / "add.py"
        p.write_text(src)
        funcs = discover_functions(src, p)

        result = extract_function_source(funcs[0])

        compile(result, "<test>", "exec")
        assert "def add(a, b):" in result
