"""Tests for context extraction orchestration."""

from __future__ import annotations

import textwrap

from codeflash_python._model import FunctionSource
from codeflash_python.context.orchestration import (
    extract_all_contexts,
    extract_contexts_for_file,
)


class TestExtractContextsForFile:
    """Tests for extract_contexts_for_file."""

    def test_fto_file_produces_rw_and_hash(
        self,
        tmp_path,
    ):
        """
        FTO file produces read-writable and hashing
        contexts with imports.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
            import math

            def target():
                return math.sqrt(2)

            def unrelated():
                return 0
        """)
        )

        rw, _ro, hsh, _tg, _cache = extract_contexts_for_file(
            file_path=mod,
            fto_names={"target"},
            hoh_names=set(),
            rw_helper_fqns={"mod.target"},
            all_helper_fqns={"mod.target"},
            project_root=tmp_path,
        )
        assert rw is not None
        assert "def target" in rw.code
        assert "def unrelated" not in rw.code
        assert "import math" in rw.code
        assert hsh is not None
        assert "def target" in hsh.code

    def test_hoh_only_file_produces_ro_and_testgen(
        self,
        tmp_path,
    ):
        """
        HoH-only file produces RO and TESTGEN but no RW.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
            def helper():
                return 42
        """)
        )

        rw, ro, _hsh, tg, _cache = extract_contexts_for_file(
            file_path=mod,
            fto_names=set(),
            hoh_names={"helper"},
            rw_helper_fqns=set(),
            all_helper_fqns={"mod.helper"},
            project_root=tmp_path,
        )
        assert rw is None
        assert ro is not None
        assert "def helper" in ro.code
        assert tg is not None
        assert "def helper" in tg.code

    def test_missing_file_returns_empty(self, tmp_path):
        """
        Missing file returns four empty strings.
        """
        missing = tmp_path / "missing.py"
        rw, ro, hsh, tg, cache = extract_contexts_for_file(
            file_path=missing,
            fto_names={"target"},
            hoh_names=set(),
            rw_helper_fqns=set(),
            all_helper_fqns=set(),
            project_root=tmp_path,
        )
        assert rw is None
        assert ro is None
        assert hsh is None
        assert tg is None
        assert cache is None

    def test_fto_with_class_method(self, tmp_path):
        """
        Class method FTO produces correct RW context.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
            class MyClass:
                def __init__(self):
                    self.x = 1

                def target(self):
                    return self.x
        """)
        )

        rw, _ro, _hsh, _tg, _cache = extract_contexts_for_file(
            file_path=mod,
            fto_names={"MyClass.target"},
            hoh_names=set(),
            rw_helper_fqns={"mod.MyClass.target"},
            all_helper_fqns={"mod.MyClass.target"},
            project_root=tmp_path,
        )
        assert rw is not None
        assert "class MyClass" in rw.code
        assert "def target" in rw.code
        assert "def __init__" in rw.code


class TestExtractAllContexts:
    """Tests for extract_all_contexts."""

    def test_single_fto_file(self, tmp_path):
        """
        Single FTO file produces correct contexts.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
            import math

            X = 42

            def target():
                \"\"\"Target function.\"\"\"
                return helper() + X

            def helper():
                return math.sqrt(2)

            def unrelated():
                return 0
        """)
        )

        target_src = FunctionSource(
            file_path=mod,
            qualified_name="target",
            fully_qualified_name="mod.target",
            source_code="",
        )
        helper_src = FunctionSource(
            file_path=mod,
            qualified_name="helper",
            fully_qualified_name="mod.helper",
            source_code="",
        )

        result = extract_all_contexts(
            helpers_of_fto={
                mod: {target_src, helper_src},
            },
            helpers_of_helpers={},
            project_root=tmp_path,
        )
        rw_md = result.read_writable.markdown
        assert "def target" in rw_md
        assert "def helper" in rw_md
        assert "X = 42" in rw_md
        assert "def unrelated" not in rw_md
        assert "import math" in rw_md

        # Hashing has targets without docstrings
        hash_md = result.hashing.markdown
        assert "def target" in hash_md
        assert "Target function" not in hash_md

    def test_hoh_deduplication(self, tmp_path):
        """
        HoH entries overlapping with FTO are deduplicated.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
            def target():
                return helper()

            def helper():
                return 42
        """)
        )

        target_src = FunctionSource(
            file_path=mod,
            qualified_name="target",
            fully_qualified_name="mod.target",
            source_code="",
        )
        helper_src = FunctionSource(
            file_path=mod,
            qualified_name="helper",
            fully_qualified_name="mod.helper",
            source_code="",
        )

        result = extract_all_contexts(
            helpers_of_fto={
                mod: {target_src, helper_src},
            },
            helpers_of_helpers={
                mod: {helper_src},
            },
            project_root=tmp_path,
        )
        # helper is in FTO so HoH dedup removes it
        assert "" == result.read_only.markdown

    def test_multi_file_contexts(self, tmp_path):
        """
        Contexts from multiple files are combined.
        """
        target_mod = tmp_path / "target_mod.py"
        target_mod.write_text(
            textwrap.dedent("""\
            def target():
                return 1
        """)
        )

        helper_mod = tmp_path / "helper_mod.py"
        helper_mod.write_text(
            textwrap.dedent("""\
            def helper_func():
                return 42
        """)
        )

        target_src = FunctionSource(
            file_path=target_mod,
            qualified_name="target",
            fully_qualified_name="target_mod.target",
            source_code="",
        )
        helper_src = FunctionSource(
            file_path=helper_mod,
            qualified_name="helper_func",
            fully_qualified_name="helper_mod.helper_func",
            source_code="",
        )

        result = extract_all_contexts(
            helpers_of_fto={
                target_mod: {target_src},
            },
            helpers_of_helpers={
                helper_mod: {helper_src},
            },
            project_root=tmp_path,
        )
        assert "def target" in result.read_writable.markdown
        assert "def helper_func" in result.read_only.markdown
        assert "def helper_func" in result.testgen.markdown

    def test_empty_inputs(self, tmp_path):
        """
        Empty inputs produce empty contexts with a hash.
        """
        result = extract_all_contexts(
            helpers_of_fto={},
            helpers_of_helpers={},
            project_root=tmp_path,
        )
        assert "" == result.read_writable.markdown
        assert "" == result.read_only.markdown
        assert "" == result.hashing.markdown
        assert "" == result.testgen.markdown
