"""Test multi-file code replacement with helper functions in separate files."""

import re
from pathlib import Path

from codeflash_python.analysis._discovery import FunctionToOptimize
from codeflash_python.codegen._replacement import replace_functions_in_file
from codeflash_python.context.models import CodeString, CodeStringsMarkdown
from codeflash_python.context.pipeline import get_code_optimization_context


def test_multi_file_replcement01() -> None:
    """Verify replacement works for a helper function in a separate file.

    The optimized code contains two code blocks: one for the helper
    file (with an optimized ``_estimate_string_tokens``) and one for
    the main file (unchanged).  ``replace_functions_in_file`` should
    update the helper file while leaving the main file unchanged.
    """
    root_dir = Path(__file__).parent.resolve()
    helper_file = (root_dir / "code_to_optimize/temp_helper.py").resolve()

    original_helper = """\
import re
from collections.abc import Sequence

from pydantic_ai_slim.pydantic_ai.messages import BinaryContent, UserContent

_TOKEN_SPLIT_RE = re.compile(r'[\\s",.:]+')

def _estimate_string_tokens(content: str | Sequence[UserContent]) -> int:
    if not content:
        return 0

    if isinstance(content, str):
        return len(_TOKEN_SPLIT_RE.split(content.strip()))

    tokens = 0
    for part in content:
        if isinstance(part, str):
            tokens += len(_TOKEN_SPLIT_RE.split(part.strip()))
        elif isinstance(part, BinaryContent):
            tokens += len(part.data)
        # TODO(Marcelo): We need to study how we can estimate the tokens for AudioUrl or ImageUrl.

    return tokens
"""
    helper_file.write_text(original_helper, encoding="utf-8")

    main_file = (root_dir / "code_to_optimize/temp_main.py").resolve()

    original_main = """\
from temp_helper import _estimate_string_tokens
from pydantic_ai_slim.pydantic_ai.usage import Usage

def _get_string_usage(text: str) -> Usage:
    response_tokens = _estimate_string_tokens(text)
    return Usage(response_tokens=response_tokens, total_tokens=response_tokens)
"""
    main_file.write_text(original_main, encoding="utf-8")

    try:
        # Optimized code for the helper file.
        optimized_helper = """\
import re
from collections.abc import Sequence

from pydantic_ai_slim.pydantic_ai.messages import BinaryContent, UserContent

_TOKEN_SPLIT_RE = re.compile(r'[\\s",.:]+')
_translate_table = {ord(c): ord(' ') for c in ' \\t\\n\\r\\x0b\\x0c",.:'}

def _estimate_string_tokens(content: str | Sequence[UserContent]) -> int:
    if not content:
        return 0

    if isinstance(content, str):
        # Fast path using translate and split instead of regex when separat
        s = content.strip()
        if s:
            s = s.translate(_translate_table)
            # Split on whitespace (default). This handles multiple consecut
            return len(s.split())
        return 0

    tokens = 0
    for part in content:
        if isinstance(part, str):
            s = part.strip()
            if s:
                s = s.translate(_translate_table)
                tokens += len(s.split())
        elif isinstance(part, BinaryContent):
            tokens += len(part.data)

    return tokens
"""

        # Get code context.
        func = FunctionToOptimize(
            function_name="_get_string_usage",
            parents=[],
            file_path=main_file,
        )
        code_context = get_code_optimization_context(
            function_to_optimize=func,
            project_root=root_dir,
        )

        # Verify helper function was discovered.
        helper_fqns = {
            h.fully_qualified_name for h in code_context.helper_functions
        }
        assert any("_estimate_string_tokens" in fqn for fqn in helper_fqns), (
            f"Expected _estimate_string_tokens in helpers, got: {helper_fqns}"
        )

        # Replace the function in the helper file.
        new_helper = replace_functions_in_file(
            source_code=original_helper,
            original_function_names=["_estimate_string_tokens"],
            optimized_code=optimized_helper,
            preexisting_objects=code_context.preexisting_objects,
        )
        helper_file.write_text(new_helper, encoding="utf-8")

        # Main file should remain unchanged (no replacement needed).
        new_main = main_file.read_text(encoding="utf-8")

        assert new_main.rstrip() == original_main.rstrip()
        assert "_translate_table" in new_helper
        assert "s.translate(_translate_table)" in new_helper
    finally:
        helper_file.unlink(missing_ok=True)
        main_file.unlink(missing_ok=True)
