"""Tests for token-limit fallback logic."""

from __future__ import annotations

import textwrap
from typing import TYPE_CHECKING

import libcst as cst
import pytest

from codeflash_python._model import FunctionSource
from codeflash_python.context.dependencies import (
    remove_unused_definitions_by_function_names,
)
from codeflash_python.context.fallback import (
    apply_token_limits,
    encoded_tokens_len,
    re_extract_from_cache,
)
from codeflash_python.context.imports import gather_source_imports
from codeflash_python.context.models import (
    AllContextResults,
    CodeContextType,
    CodeString,
    CodeStringsMarkdown,
    FileContextCache,
)
from codeflash_python.context.orchestration import (
    extract_all_contexts,
    extract_contexts_for_file,
)

if TYPE_CHECKING:
    from pathlib import Path


def build_cache(  # noqa: PLR0913
    code: str,
    file_path: Path,
    project_root: Path,
    *,
    fto_names: set[str] | None = None,
    hoh_names: set[str] | None = None,
    helper_fqns: set[str] | None = None,
) -> FileContextCache:
    """Build a FileContextCache from source code."""
    if fto_names is None:
        fto_names = set()
    if hoh_names is None:
        hoh_names = set()
    if helper_fqns is None:
        helper_fqns = set()

    original = cst.parse_module(code)
    all_names = fto_names | hoh_names
    cleaned = remove_unused_definitions_by_function_names(
        original,
        all_names,
    )
    gathered = gather_source_imports(
        original,
        file_path,
        project_root,
    )
    return FileContextCache(
        original_module=original,
        cleaned_module=cleaned,
        fto_names=fto_names,
        hoh_names=hoh_names,
        helper_fqns=helper_fqns,
        file_path=file_path,
        relative_path=file_path.relative_to(project_root),
        gathered_imports=gathered,
    )


class TestEncodedTokensLen:
    """Tests for encoded_tokens_len."""

    def test_empty_string(self):
        """Empty string has zero tokens."""
        assert 0 == encoded_tokens_len("")

    def test_short_string(self):
        """Four-character string yields one token."""
        assert 1 == encoded_tokens_len("test")

    def test_longer_string(self):
        """Twelve-character string yields three tokens."""
        assert 3 == encoded_tokens_len("hello world!")


class TestReExtractFromCache:
    """Tests for re_extract_from_cache."""

    def test_strips_docstrings(self, tmp_path):
        """
        Docstrings are stripped when remove_docstrings is True.
        """
        code = textwrap.dedent("""\
            def helper():
                \"\"\"A helper docstring.\"\"\"
                return 42
        """)
        mod_file = tmp_path / "mod.py"
        mod_file.write_text(code)

        cache = build_cache(
            code,
            mod_file,
            tmp_path,
            hoh_names={"helper"},
            helper_fqns={"mod.helper"},
        )
        result = re_extract_from_cache(
            [cache],
            CodeContextType.READ_ONLY,
            tmp_path,
            remove_docstrings=True,
        )
        assert isinstance(result, CodeStringsMarkdown)
        assert len(result.code_strings) > 0
        combined = "\n".join(cs.code for cs in result.code_strings)
        assert "def helper" in combined
        assert "helper docstring" not in combined

    def test_preserves_docstrings_when_disabled(
        self,
        tmp_path,
    ):
        """
        Docstrings are kept when remove_docstrings is False.
        """
        code = textwrap.dedent("""\
            def helper():
                \"\"\"A helper docstring.\"\"\"
                return 42
        """)
        mod_file = tmp_path / "mod.py"
        mod_file.write_text(code)

        cache = build_cache(
            code,
            mod_file,
            tmp_path,
            hoh_names={"helper"},
            helper_fqns={"mod.helper"},
        )
        result = re_extract_from_cache(
            [cache],
            CodeContextType.READ_ONLY,
            tmp_path,
            remove_docstrings=False,
        )
        combined = "\n".join(cs.code for cs in result.code_strings)
        assert "def helper" in combined
        assert "helper docstring" in combined

    def test_empty_caches_returns_empty(
        self,
        tmp_path,
    ):
        """
        Empty cache list returns an empty CodeStringsMarkdown.
        """
        result = re_extract_from_cache(
            [],
            CodeContextType.READ_ONLY,
            tmp_path,
        )
        assert isinstance(result, CodeStringsMarkdown)
        assert [] == result.code_strings

    def test_hashing_normalizes_via_ast(self, tmp_path):
        """
        HASHING context type normalizes output via ast.unparse.
        """
        code = textwrap.dedent("""\
            def  target( ):
                x  =  1
                return   x
        """)
        mod_file = tmp_path / "mod.py"
        mod_file.write_text(code)

        cache = build_cache(
            code,
            mod_file,
            tmp_path,
            fto_names={"target"},
            helper_fqns={"mod.target"},
        )
        result = re_extract_from_cache(
            [cache],
            CodeContextType.HASHING,
            tmp_path,
        )
        combined = "\n".join(cs.code for cs in result.code_strings)
        assert "def target" in combined
        assert "  =  " not in combined

    def test_result_has_file_paths(self, tmp_path):
        """
        Each CodeString in the result has a file_path set.
        """
        code = textwrap.dedent("""\
            def helper():
                return 42
        """)
        mod_file = tmp_path / "mod.py"
        mod_file.write_text(code)

        cache = build_cache(
            code,
            mod_file,
            tmp_path,
            hoh_names={"helper"},
            helper_fqns={"mod.helper"},
        )
        result = re_extract_from_cache(
            [cache],
            CodeContextType.READ_ONLY,
            tmp_path,
        )
        assert len(result.code_strings) > 0
        for cs in result.code_strings:
            assert cs.file_path is not None


class TestApplyTokenLimits:
    """Tests for apply_token_limits."""

    def test_under_limit_returns_unchanged(self, tmp_path):
        """
        Context under the token limit is returned unchanged.
        """
        rw = CodeStringsMarkdown(
            code_strings=[
                CodeString(code="def target(): return 1"),
            ],
        )
        ro = CodeStringsMarkdown(
            code_strings=[
                CodeString(code="def helper(): return 2"),
            ],
        )
        all_results = AllContextResults(
            read_writable=rw,
            read_only=ro,
            hashing=CodeStringsMarkdown(),
            testgen=CodeStringsMarkdown(),
            file_caches=[],
        )
        result = apply_token_limits(
            all_results,
            tmp_path,
            optim_token_limit=100000,
        )
        assert "def helper(): return 2" in result.read_only

    def test_rw_exceeds_limit_raises(self, tmp_path):
        """
        ValueError is raised when read_writable alone exceeds
        the token limit.
        """
        large_rw = CodeStringsMarkdown(
            code_strings=[CodeString(code="x" * 1000)],
        )
        all_results = AllContextResults(
            read_writable=large_rw,
            read_only=CodeStringsMarkdown(
                code_strings=[CodeString(code="small")],
            ),
            hashing=CodeStringsMarkdown(),
            testgen=CodeStringsMarkdown(),
            file_caches=[],
        )
        with pytest.raises(ValueError, match=r"(?i)read.writable"):
            apply_token_limits(
                all_results,
                tmp_path,
                optim_token_limit=10,
            )

    def test_strips_docstrings_when_over_limit(
        self,
        tmp_path,
    ):
        """
        Docstrings are stripped from read_only when the
        combined context exceeds the limit.
        """
        long_doc = (
            "This is a very long docstring that takes up "
            "a lot of tokens and should be stripped."
        )
        code = textwrap.dedent(f"""\
            def helper():
                \"\"\"{long_doc}\"\"\"
                return 42
        """)
        mod_file = tmp_path / "mod.py"
        mod_file.write_text(code)

        cache = build_cache(
            code,
            mod_file,
            tmp_path,
            hoh_names={"helper"},
            helper_fqns={"mod.helper"},
        )

        rw = CodeStringsMarkdown(
            code_strings=[
                CodeString(code="def target(): return 1"),
            ],
        )
        ro = CodeStringsMarkdown(
            code_strings=[
                CodeString(
                    code=code.strip(),
                    file_path=mod_file.relative_to(tmp_path),
                ),
            ],
        )
        all_results = AllContextResults(
            read_writable=rw,
            read_only=ro,
            hashing=CodeStringsMarkdown(),
            testgen=CodeStringsMarkdown(),
            file_caches=[cache],
        )

        rw_tokens = encoded_tokens_len(rw.markdown)
        ro_tokens = encoded_tokens_len(ro.markdown)
        limit = rw_tokens + ro_tokens - 5

        result = apply_token_limits(
            all_results,
            tmp_path,
            optim_token_limit=limit,
        )
        assert "def helper" in result.read_only
        assert "very long docstring" not in result.read_only

    def test_removes_read_only_when_still_over(
        self,
        tmp_path,
    ):
        """
        read_only is set to empty string when stripping
        docstrings is not enough to fit within limits.
        """
        code = textwrap.dedent("""\
            def helper():
                return 42
        """)
        mod_file = tmp_path / "mod.py"
        mod_file.write_text(code)

        cache = build_cache(
            code,
            mod_file,
            tmp_path,
            hoh_names={"helper"},
            helper_fqns={"mod.helper"},
        )

        rw = CodeStringsMarkdown(
            code_strings=[CodeString(code="x" * 100)],
        )
        ro = CodeStringsMarkdown(
            code_strings=[CodeString(code="y" * 100)],
        )
        all_results = AllContextResults(
            read_writable=rw,
            read_only=ro,
            hashing=CodeStringsMarkdown(),
            testgen=CodeStringsMarkdown(),
            file_caches=[cache],
        )
        rw_tokens = encoded_tokens_len(rw.markdown)
        result = apply_token_limits(
            all_results,
            tmp_path,
            optim_token_limit=rw_tokens + 1,
        )
        assert "" == result.read_only


class TestExtractContextsForFileCache:
    """Tests for extract_contexts_for_file returning cache."""

    def test_returns_cache_on_success(self, tmp_path):
        """
        Fifth element is a FileContextCache on success.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
                def target():
                    return 1
            """)
        )

        result = 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 5 == len(result)
        cache = result[4]
        assert isinstance(cache, FileContextCache)
        assert mod == cache.file_path
        assert {"target"} == cache.fto_names

    def test_returns_code_strings_with_paths(self, tmp_path):
        """
        Non-None results are CodeString with file_path set.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
                def target():
                    return 1
            """)
        )

        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 isinstance(rw, CodeString)
        assert rw.file_path is not None
        assert "def target" in rw.code

    def test_returns_none_cache_on_parse_error(
        self,
        tmp_path,
    ):
        """
        Fifth element is None when the file cannot be parsed.
        """
        bad = tmp_path / "bad.py"
        bad.write_text("def (broken syntax\n")

        result = extract_contexts_for_file(
            file_path=bad,
            fto_names={"target"},
            hoh_names=set(),
            rw_helper_fqns=set(),
            all_helper_fqns=set(),
            project_root=tmp_path,
        )
        assert 5 == len(result)
        assert result[4] is None


class TestExtractAllContextsReturnsAllContextResults:
    """Tests for extract_all_contexts returning AllContextResults."""

    def test_returns_all_context_results(self, tmp_path):
        """
        extract_all_contexts returns an AllContextResults with
        CodeStringsMarkdown fields and non-empty file_caches.
        """
        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={},
            project_root=tmp_path,
        )
        assert isinstance(result, AllContextResults)
        assert isinstance(result.read_writable, CodeStringsMarkdown)
        assert isinstance(result.read_only, CodeStringsMarkdown)
        assert isinstance(result.hashing, CodeStringsMarkdown)
        assert isinstance(result.testgen, CodeStringsMarkdown)
        assert len(result.read_writable.code_strings) > 0
        rw_code = "\n".join(
            cs.code for cs in result.read_writable.code_strings
        )
        assert "def target" in rw_code
        assert isinstance(result.file_caches, list)
        assert len(result.file_caches) > 0

    def test_code_strings_have_file_paths(self, tmp_path):
        """
        CodeString entries have relative file_path set.
        """
        mod = tmp_path / "mod.py"
        mod.write_text(
            textwrap.dedent("""\
                def target():
                    return 1
            """)
        )

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

        result = extract_all_contexts(
            helpers_of_fto={mod: {target_src}},
            helpers_of_helpers={},
            project_root=tmp_path,
        )
        for cs in result.read_writable.code_strings:
            assert cs.file_path is not None


class TestCodeStringsMarkdown:
    """Tests for CodeStringsMarkdown.markdown property."""

    def test_empty_produces_empty_string(self):
        """
        Empty code_strings produces an empty markdown string.
        """
        csm = CodeStringsMarkdown()
        assert "" == csm.markdown

    def test_single_block_with_path(self, tmp_path):
        """
        Single code string with path produces a code block
        with file path suffix.
        """
        from pathlib import PurePosixPath

        csm = CodeStringsMarkdown(
            code_strings=[
                CodeString(
                    code="def foo(): pass",
                    file_path=PurePosixPath("src/mod.py"),
                ),
            ],
        )
        md = csm.markdown
        assert "```python:src/mod.py" in md
        assert "def foo(): pass" in md
        assert md.endswith("```")

    def test_single_block_without_path(self):
        """
        Code string without path produces a bare code block.
        """
        csm = CodeStringsMarkdown(
            code_strings=[CodeString(code="x = 1")],
        )
        md = csm.markdown
        assert md.startswith("```python\n")
        assert "x = 1" in md

    def test_multiple_blocks(self, tmp_path):
        """
        Multiple code strings produce multiple code blocks.
        """
        from pathlib import PurePosixPath

        csm = CodeStringsMarkdown(
            code_strings=[
                CodeString(
                    code="def a(): pass",
                    file_path=PurePosixPath("a.py"),
                ),
                CodeString(
                    code="def b(): pass",
                    file_path=PurePosixPath("b.py"),
                ),
            ],
        )
        md = csm.markdown
        assert "```python:a.py" in md
        assert "```python:b.py" in md
        assert "def a(): pass" in md
        assert "def b(): pass" in md
