"""Tests for module preparation (stage 23a)."""

from __future__ import annotations

import ast
import textwrap
from typing import TYPE_CHECKING

from codeflash_python._model import FunctionParent
from codeflash_python.pipeline._module_prep import (
    ValidCode,
    prepare_python_module,
    resolve_python_function_ast,
)

if TYPE_CHECKING:
    from pathlib import Path


class TestValidCode:
    """Tests for the ValidCode frozen data class."""

    def test_fields(self) -> None:
        """ValidCode stores source and normalized code."""
        vc = ValidCode(source_code="x = 1", normalized_code="x=1")
        assert "x = 1" == vc.source_code
        assert "x=1" == vc.normalized_code

    def test_frozen(self) -> None:
        """ValidCode instances are immutable."""
        import attrs

        vc = ValidCode(source_code="a", normalized_code="b")
        assert attrs.has(type(vc))


class TestPreparePythonModule:
    """Tests for prepare_python_module."""

    def test_basic_module(self, tmp_path: Path) -> None:
        """A valid module produces ValidCode with normalized form."""
        src = tmp_path / "sample.py"
        code = textwrap.dedent("""\
            def foo():
                return 42
        """)
        src.write_text(code)

        result = prepare_python_module(code, src, tmp_path)
        assert result is not None
        validated, module_ast = result
        assert src in validated
        assert validated[src].source_code == code
        assert isinstance(module_ast, ast.Module)

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

        result = prepare_python_module(code, src, tmp_path)
        assert result is None

    def test_callee_included(self, tmp_path: Path) -> None:
        """Imported internal modules are included in the result."""
        helper = tmp_path / "helper.py"
        helper.write_text(
            textwrap.dedent("""\
                def add(a, b):
                    return a + b
            """),
        )

        main = tmp_path / "main.py"
        main_code = textwrap.dedent("""\
            from helper import add

            def compute():
                return add(1, 2)
        """)
        main.write_text(main_code)

        result = prepare_python_module(main_code, main, tmp_path)
        assert result is not None
        validated, _ = result
        assert main in validated
        assert helper.resolve() in validated

    def test_callee_syntax_error_returns_none(self, tmp_path: Path) -> None:
        """If a callee module has a syntax error, return None."""
        helper = tmp_path / "broken.py"
        helper.write_text("def foo(:\n")

        main = tmp_path / "main.py"
        main_code = textwrap.dedent("""\
            from broken import foo

            def run():
                return foo()
        """)
        main.write_text(main_code)

        result = prepare_python_module(main_code, main, tmp_path)
        assert result is None


class TestResolvePythonFunctionAst:
    """Tests for resolve_python_function_ast."""

    def test_top_level_function(self) -> None:
        """A top-level function is resolved by name."""
        code = textwrap.dedent("""\
            def foo():
                return 42

            def bar():
                return 0
        """)
        module_ast = ast.parse(code)

        result = resolve_python_function_ast("foo", [], module_ast)
        assert result is not None
        assert result.name == "foo"

    def test_method_in_class(self) -> None:
        """A method inside a class is resolved via parents."""
        code = textwrap.dedent("""\
            class MyClass:
                def my_method(self):
                    return 1
        """)
        module_ast = ast.parse(code)
        parents = [FunctionParent(name="MyClass", type="ClassDef")]

        result = resolve_python_function_ast("my_method", parents, module_ast)
        assert result is not None
        assert result.name == "my_method"

    def test_missing_function_returns_none(self) -> None:
        """A function not in the module returns None."""
        code = "x = 1\n"
        module_ast = ast.parse(code)

        result = resolve_python_function_ast("nonexistent", [], module_ast)
        assert result is None

    def test_async_function(self) -> None:
        """An async function is resolved correctly."""
        code = textwrap.dedent("""\
            async def fetch():
                return await something()
        """)
        module_ast = ast.parse(code)

        result = resolve_python_function_ast("fetch", [], module_ast)
        assert result is not None
        assert isinstance(result, ast.AsyncFunctionDef)
