"""Tests for the top-level context extraction pipeline."""

from __future__ import annotations

import textwrap

import pytest

from codeflash_python._model import FunctionToOptimize
from codeflash_python.context.enrichment import build_testgen_context
from codeflash_python.context.models import (
    CodeString,
    CodeStringsMarkdown,
)
from codeflash_python.context.pipeline import (
    get_code_optimization_context,
)


class TestBuildTestgenContext:
    """Tests for build_testgen_context."""

    def test_passes_through_without_enrichment(self, tmp_path):
        """
        With enrichment disabled, base context is returned unchanged.
        """
        base = CodeStringsMarkdown(
            code_strings=[CodeString(code="def helper(): return 1")],
        )
        result = build_testgen_context(
            base,
            tmp_path,
            include_enrichment=False,
        )
        assert 1 == len(result.code_strings)
        assert "def helper" in result.code_strings[0].code

    def test_enrichment_enabled_returns_at_least_base(
        self,
        tmp_path,
    ):
        """
        With enrichment enabled but no resolvable imports,
        the base context is preserved.
        """
        base = CodeStringsMarkdown(
            code_strings=[CodeString(code="x = 1")],
        )
        result = build_testgen_context(
            base,
            tmp_path,
        )
        assert len(result.code_strings) >= 1
        assert "x = 1" in result.code_strings[0].code


class TestGetCodeOptimizationContext:
    """Tests for get_code_optimization_context."""

    def test_simple_function(self, tmp_path):
        """
        Simple function produces a valid CodeOptimizationContext
        with populated read_writable and hashing fields.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
                def target():
                    return 42
            """)
        )

        fto = FunctionToOptimize(
            function_name="target",
            file_path=mod,
        )
        result = get_code_optimization_context(fto, tmp_path)

        assert "def target" in result.read_writable
        assert "def target" in result.hashing
        assert 64 == len(result.hashing_hash)

    def test_function_with_helper(self, tmp_path):
        """
        Function calling a helper in the same file includes the
        helper in read_writable context.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
                def target():
                    return helper()

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

        fto = FunctionToOptimize(
            function_name="target",
            file_path=mod,
        )
        result = get_code_optimization_context(fto, tmp_path)

        assert "def target" in result.read_writable
        assert "def helper" in result.read_writable

    def test_target_file_first_in_rw(self, tmp_path):
        """
        The target file's code block appears first in read_writable.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
                def target():
                    return 1
            """)
        )

        fto = FunctionToOptimize(
            function_name="target",
            file_path=mod,
        )
        result = get_code_optimization_context(fto, tmp_path)

        assert len(result.read_writable_code.code_strings) > 0
        first_block = result.read_writable_code.code_strings[0]
        assert "def target" in first_block.code

    def test_hashing_hash_is_sha256(self, tmp_path):
        """
        hashing_hash is a SHA256 hex digest of the hashing markdown.
        """
        import hashlib

        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
                def target():
                    return 1
            """)
        )

        fto = FunctionToOptimize(
            function_name="target",
            file_path=mod,
        )
        result = get_code_optimization_context(fto, tmp_path)

        expected = hashlib.sha256(
            result.hashing.encode("utf-8"),
        ).hexdigest()
        assert expected == result.hashing_hash

    def test_helper_fqns_populated(self, tmp_path):
        """
        testgen_helper_fqns contains fully qualified names
        of discovered helpers.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
                def target():
                    return 1
            """)
        )

        fto = FunctionToOptimize(
            function_name="target",
            file_path=mod,
        )
        result = get_code_optimization_context(fto, tmp_path)

        assert isinstance(result.testgen_helper_fqns, list)
        assert len(result.testgen_helper_fqns) > 0

    def test_preexisting_objects_populated(self, tmp_path):
        """
        preexisting_objects contains function/class names from
        the read_writable context.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
                def target():
                    return 1
            """)
        )

        fto = FunctionToOptimize(
            function_name="target",
            file_path=mod,
        )
        result = get_code_optimization_context(fto, tmp_path)

        names = {name for name, _ in result.preexisting_objects}
        assert "target" in names

    def test_rw_exceeds_limit_raises(self, tmp_path):
        """
        ValueError is raised when read_writable alone exceeds
        the token limit.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
                def target():
                    return 1
            """)
        )

        fto = FunctionToOptimize(
            function_name="target",
            file_path=mod,
        )
        with pytest.raises(ValueError, match=r"(?i)read.writable"):
            get_code_optimization_context(
                fto,
                tmp_path,
                optim_token_limit=1,
            )
