"""Test replace_function_and_helpers_with_optimized_code with mock candidate from mock_candidate.txt."""

import tempfile
from pathlib import Path

import pytest

from codeflash_python._model import FunctionParent
from codeflash_python.analysis._discovery import FunctionToOptimize
from codeflash_python.context.models import CodeStringsMarkdown
from codeflash_python.context.pipeline import get_code_optimization_context
from codeflash_python.pipeline._function_optimizer import (
    replace_function_and_helpers,
)
from codeflash_python.verification._unused_helpers import (
    detect_unused_helper_functions,
)

ORIGINAL_SOURCE = """\
import contextlib
from typing import BinaryIO, TypeVar, Union

_SymbolT = TypeVar("_SymbolT", PSLiteral, PSKeyword)


PSLiteralTable = PSSymbolTable(PSLiteral)
PSKeywordTable = PSSymbolTable(PSKeyword)
LIT = PSLiteralTable.intern
KWD = PSKeywordTable.intern
KEYWORD_DICT_BEGIN = KWD(b"<<")
KEYWORD_DICT_END = KWD(b">>")


PSBaseParserToken = Union[float, bool, PSLiteral, PSKeyword, bytes]


class PSBaseParser:

    def __init__(self, fp: BinaryIO) -> None:
        self.fp = fp
        self.eof = False
        self.seek(0)

    def _parse_main(self, s: bytes, i: int) -> int:
        m = NONSPC.search(s, i)
        if not m:
            return len(s)
        j = m.start(0)
        c = s[j : j + 1]
        self._curtokenpos = self.bufpos + j
        if c == b"%":
            self._curtoken = b"%"
            self._parse1 = self._parse_comment
            return j + 1
        elif c == b"/":
            self._curtoken = b""
            self._parse1 = self._parse_literal
            return j + 1
        elif c in b"-+" or c.isdigit():
            self._curtoken = c
            self._parse1 = self._parse_number
            return j + 1
        elif c == b".":
            self._curtoken = c
            self._parse1 = self._parse_float
            return j + 1
        elif c.isalpha():
            self._curtoken = c
            self._parse1 = self._parse_keyword
            return j + 1
        elif c == b"(":
            self._curtoken = b""
            self.paren = 1
            self._parse1 = self._parse_string
            return j + 1
        elif c == b"<":
            self._curtoken = b""
            self._parse1 = self._parse_wopen
            return j + 1
        elif c == b">":
            self._curtoken = b""
            self._parse1 = self._parse_wclose
            return j + 1
        elif c == b"\\x00":
            return j + 1
        else:
            self._add_token(KWD(c))
            return j + 1

    def _add_token(self, obj: PSBaseParserToken) -> None:
        self._tokens.append((self._curtokenpos, obj))

    def _parse_comment(self, s: bytes, i: int) -> int:
        m = EOL.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        self._parse1 = self._parse_main
        return j

    def _parse_literal(self, s: bytes, i: int) -> int:
        m = END_LITERAL.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        c = s[j : j + 1]
        if c == b"#":
            self.hex = b""
            self._parse1 = self._parse_literal_hex
            return j + 1
        try:
            name: str | bytes = str(self._curtoken, "utf-8")
        except Exception:
            name = self._curtoken
        self._add_token(LIT(name))
        self._parse1 = self._parse_main
        return j

    def _parse_number(self, s: bytes, i: int) -> int:
        m = END_NUMBER.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        c = s[j : j + 1]
        if c == b".":
            self._curtoken += b"."
            self._parse1 = self._parse_float
            return j + 1
        with contextlib.suppress(ValueError):
            self._add_token(int(self._curtoken))
        self._parse1 = self._parse_main
        return j

    def _parse_float(self, s: bytes, i: int) -> int:
        m = END_NUMBER.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        with contextlib.suppress(ValueError):
            self._add_token(float(self._curtoken))
        self._parse1 = self._parse_main
        return j

    def _parse_keyword(self, s: bytes, i: int) -> int:
        m = END_KEYWORD.search(s, i)
        if m:
            j = m.start(0)
            self._curtoken += s[i:j]
        else:
            self._curtoken += s[i:]
            return len(s)
        if self._curtoken == b"true":
            token: bool | PSKeyword = True
        elif self._curtoken == b"false":
            token = False
        else:
            token = KWD(self._curtoken)
        self._add_token(token)
        self._parse1 = self._parse_main
        return j

    def _parse_string(self, s: bytes, i: int) -> int:
        m = END_STRING.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        c = s[j : j + 1]
        if c == b"\\\\":
            self.oct = b""
            self._parse1 = self._parse_string_1
            return j + 1
        if c == b"(":
            self.paren += 1
            self._curtoken += c
            return j + 1
        if c == b")":
            self.paren -= 1
            if self.paren:
                self._curtoken += c
                return j + 1
        self._add_token(self._curtoken)
        self._parse1 = self._parse_main
        return j + 1

    def _parse_wopen(self, s: bytes, i: int) -> int:
        c = s[i : i + 1]
        if c == b"<":
            self._add_token(KEYWORD_DICT_BEGIN)
            self._parse1 = self._parse_main
            i += 1
        else:
            self._parse1 = self._parse_hexstring
        return i

    def _parse_wclose(self, s: bytes, i: int) -> int:
        c = s[i : i + 1]
        if c == b">":
            self._add_token(KEYWORD_DICT_END)
            i += 1
        self._parse1 = self._parse_main
        return i
"""

MOCK_CANDIDATE_MARKDOWN = """\
```python
#!/usr/bin/env python3


import contextlib
from typing import BinaryIO, TypeVar, Union

_SymbolT = TypeVar("_SymbolT", PSLiteral, PSKeyword)


PSLiteralTable = PSSymbolTable(PSLiteral)
PSKeywordTable = PSSymbolTable(PSKeyword)
LIT = PSLiteralTable.intern
KWD = PSKeywordTable.intern
KEYWORD_DICT_BEGIN = KWD(b"<<")
KEYWORD_DICT_END = KWD(b">>")


PSBaseParserToken = Union[float, bool, PSLiteral, PSKeyword, bytes]


class PSBaseParser:

    def __init__(self, fp: BinaryIO) -> None:
        self.fp = fp
        self.eof = False
        self.seek(0)

    def _parse_main(self, s: bytes, i: int) -> int:
        m = NONSPC.search(s, i)
        if not m:
            return len(s)
        j = m.start(0)
        # Use integer byte access to avoid creating a new one-byte bytes object.
        c_int = s[j]
        c_byte = bytes((c_int,))
        self._curtokenpos = self.bufpos + j
        if c_int == 37:  # b"%"
            self._curtoken = b"%"
            self._parse1 = self._parse_comment
            return j + 1
        elif c_int == 47:  # b"/"
            self._curtoken = b""
            self._parse1 = self._parse_literal
            return j + 1
        # b"-" is 45, b"+" is 43
        elif c_int == 45 or c_int == 43 or (48 <= c_int <= 57):
            self._curtoken = c_byte
            self._parse1 = self._parse_number
            return j + 1
        elif c_int == 46:  # b"."
            self._curtoken = c_byte
            self._parse1 = self._parse_float
            return j + 1
        # ASCII alphabetic check
        elif (65 <= c_int <= 90) or (97 <= c_int <= 122):
            self._curtoken = c_byte
            self._parse1 = self._parse_keyword
            return j + 1
        elif c_int == 40:  # b"("
            self._curtoken = b""
            self.paren = 1
            self._parse1 = self._parse_string
            return j + 1
        elif c_int == 60:  # b"<"
            self._curtoken = b""
            self._parse1 = self._parse_wopen
            return j + 1
        elif c_int == 62:  # b">"
            self._curtoken = b""
            self._parse1 = self._parse_wclose
            return j + 1
        elif c_int == 0:  # b"\\x00"
            return j + 1
        else:
            self._add_token(KWD(c_byte))
            return j + 1

    def _add_token(self, obj: PSBaseParserToken) -> None:
        self._tokens.append((self._curtokenpos, obj))

    def _parse_comment(self, s: bytes, i: int) -> int:
        m = EOL.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        self._parse1 = self._parse_main
        # We ignore comments.
        # self._tokens.append(self._curtoken)
        return j

    def _parse_literal(self, s: bytes, i: int) -> int:
        m = END_LITERAL.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        c_int = s[j]
        if c_int == 35:  # b"#"
            self.hex = b""
            self._parse1 = self._parse_literal_hex
            return j + 1
        try:
            name: str | bytes = str(self._curtoken, "utf-8")
        except Exception:
            name = self._curtoken
        self._add_token(LIT(name))
        self._parse1 = self._parse_main
        return j

    def _parse_number(self, s: bytes, i: int) -> int:
        m = END_NUMBER.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        c_int = s[j]
        if c_int == 46:  # b"."
            self._curtoken += b"."
            self._parse1 = self._parse_float
            return j + 1
        with contextlib.suppress(ValueError):
            self._add_token(int(self._curtoken))
        self._parse1 = self._parse_main
        return j

    def _parse_float(self, s: bytes, i: int) -> int:
        m = END_NUMBER.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        with contextlib.suppress(ValueError):
            self._add_token(float(self._curtoken))
        self._parse1 = self._parse_main
        return j

    def _parse_keyword(self, s: bytes, i: int) -> int:
        m = END_KEYWORD.search(s, i)
        if m:
            j = m.start(0)
            self._curtoken += s[i:j]
        else:
            self._curtoken += s[i:]
            return len(s)
        if self._curtoken == b"true":
            token: bool | PSKeyword = True
        elif self._curtoken == b"false":
            token = False
        else:
            token = KWD(self._curtoken)
        self._add_token(token)
        self._parse1 = self._parse_main
        return j

    def _parse_string(self, s: bytes, i: int) -> int:
        m = END_STRING.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        c_int = s[j]
        if c_int == 92:  # b"\\\\"
            self.oct = b""
            self._parse1 = self._parse_string_1
            return j + 1
        if c_int == 40:  # b"("
            self.paren += 1
            # append the literal "(" byte
            self._curtoken += b"("
            return j + 1
        if c_int == 41:  # b")"
            self.paren -= 1
            if self.paren:
                # WTF, they said balanced parens need no special treatment.
                self._curtoken += b")"
                return j + 1
        self._add_token(self._curtoken)
        self._parse1 = self._parse_main
        return j + 1

    def _parse_wopen(self, s: bytes, i: int) -> int:
        c_int = s[i]
        if c_int == 60:  # b"<"
            self._add_token(KEYWORD_DICT_BEGIN)
            self._parse1 = self._parse_main
            i += 1
        else:
            self._parse1 = self._parse_hexstring
        return i

    def _parse_wclose(self, s: bytes, i: int) -> int:
        c_int = s[i]
        if c_int == 62:  # b">"
            self._add_token(KEYWORD_DICT_END)
            i += 1
        self._parse1 = self._parse_main
        return i
```
"""

EXPECTED_OUTPUT = """\
import contextlib
from typing import BinaryIO, TypeVar, Union

_SymbolT = TypeVar("_SymbolT", PSLiteral, PSKeyword)


PSLiteralTable = PSSymbolTable(PSLiteral)
PSKeywordTable = PSSymbolTable(PSKeyword)
LIT = PSLiteralTable.intern
KWD = PSKeywordTable.intern
KEYWORD_DICT_BEGIN = KWD(b"<<")
KEYWORD_DICT_END = KWD(b">>")


PSBaseParserToken = Union[float, bool, PSLiteral, PSKeyword, bytes]


class PSBaseParser:

    def __init__(self, fp: BinaryIO) -> None:
        self.fp = fp
        self.eof = False
        self.seek(0)

    def _parse_main(self, s: bytes, i: int) -> int:
        m = NONSPC.search(s, i)
        if not m:
            return len(s)
        j = m.start(0)
        # Use integer byte access to avoid creating a new one-byte bytes object.
        c_int = s[j]
        c_byte = bytes((c_int,))
        self._curtokenpos = self.bufpos + j
        if c_int == 37:  # b"%"
            self._curtoken = b"%"
            self._parse1 = self._parse_comment
            return j + 1
        elif c_int == 47:  # b"/"
            self._curtoken = b""
            self._parse1 = self._parse_literal
            return j + 1
        # b"-" is 45, b"+" is 43
        elif c_int == 45 or c_int == 43 or (48 <= c_int <= 57):
            self._curtoken = c_byte
            self._parse1 = self._parse_number
            return j + 1
        elif c_int == 46:  # b"."
            self._curtoken = c_byte
            self._parse1 = self._parse_float
            return j + 1
        # ASCII alphabetic check
        elif (65 <= c_int <= 90) or (97 <= c_int <= 122):
            self._curtoken = c_byte
            self._parse1 = self._parse_keyword
            return j + 1
        elif c_int == 40:  # b"("
            self._curtoken = b""
            self.paren = 1
            self._parse1 = self._parse_string
            return j + 1
        elif c_int == 60:  # b"<"
            self._curtoken = b""
            self._parse1 = self._parse_wopen
            return j + 1
        elif c_int == 62:  # b">"
            self._curtoken = b""
            self._parse1 = self._parse_wclose
            return j + 1
        elif c_int == 0:  # b"\\x00"
            return j + 1
        else:
            self._add_token(KWD(c_byte))
            return j + 1

    def _add_token(self, obj: PSBaseParserToken) -> None:
        self._tokens.append((self._curtokenpos, obj))

    def _parse_comment(self, s: bytes, i: int) -> int:
        m = EOL.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        self._parse1 = self._parse_main
        # We ignore comments.
        # self._tokens.append(self._curtoken)
        return j

    def _parse_literal(self, s: bytes, i: int) -> int:
        m = END_LITERAL.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        c_int = s[j]
        if c_int == 35:  # b"#"
            self.hex = b""
            self._parse1 = self._parse_literal_hex
            return j + 1
        try:
            name: str | bytes = str(self._curtoken, "utf-8")
        except Exception:
            name = self._curtoken
        self._add_token(LIT(name))
        self._parse1 = self._parse_main
        return j

    def _parse_number(self, s: bytes, i: int) -> int:
        m = END_NUMBER.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        c_int = s[j]
        if c_int == 46:  # b"."
            self._curtoken += b"."
            self._parse1 = self._parse_float
            return j + 1
        with contextlib.suppress(ValueError):
            self._add_token(int(self._curtoken))
        self._parse1 = self._parse_main
        return j

    def _parse_float(self, s: bytes, i: int) -> int:
        m = END_NUMBER.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        with contextlib.suppress(ValueError):
            self._add_token(float(self._curtoken))
        self._parse1 = self._parse_main
        return j

    def _parse_keyword(self, s: bytes, i: int) -> int:
        m = END_KEYWORD.search(s, i)
        if m:
            j = m.start(0)
            self._curtoken += s[i:j]
        else:
            self._curtoken += s[i:]
            return len(s)
        if self._curtoken == b"true":
            token: bool | PSKeyword = True
        elif self._curtoken == b"false":
            token = False
        else:
            token = KWD(self._curtoken)
        self._add_token(token)
        self._parse1 = self._parse_main
        return j

    def _parse_string(self, s: bytes, i: int) -> int:
        m = END_STRING.search(s, i)
        if not m:
            self._curtoken += s[i:]
            return len(s)
        j = m.start(0)
        self._curtoken += s[i:j]
        c_int = s[j]
        if c_int == 92:  # b"\\\\"
            self.oct = b""
            self._parse1 = self._parse_string_1
            return j + 1
        if c_int == 40:  # b"("
            self.paren += 1
            # append the literal "(" byte
            self._curtoken += b"("
            return j + 1
        if c_int == 41:  # b")"
            self.paren -= 1
            if self.paren:
                # WTF, they said balanced parens need no special treatment.
                self._curtoken += b")"
                return j + 1
        self._add_token(self._curtoken)
        self._parse1 = self._parse_main
        return j + 1

    def _parse_wopen(self, s: bytes, i: int) -> int:
        c_int = s[i]
        if c_int == 60:  # b"<"
            self._add_token(KEYWORD_DICT_BEGIN)
            self._parse1 = self._parse_main
            i += 1
        else:
            self._parse1 = self._parse_hexstring
        return i

    def _parse_wclose(self, s: bytes, i: int) -> int:
        c_int = s[i]
        if c_int == 62:  # b">"
            self._add_token(KEYWORD_DICT_END)
            i += 1
        self._parse1 = self._parse_main
        return i
"""


@pytest.fixture
def temp_project():
    temp_dir = Path(tempfile.mkdtemp()).resolve()
    source_file = temp_dir / "psparser.py"
    source_file.write_text(ORIGINAL_SOURCE, encoding="utf-8")

    yield temp_dir, source_file

    import shutil

    shutil.rmtree(temp_dir, ignore_errors=True)


def run_replacement(temp_project):
    """Helper: run the full replacement pipeline and return (fto, code_context, final_content)."""
    temp_dir, source_file = temp_project

    function_to_optimize = FunctionToOptimize(
        file_path=source_file,
        function_name="_parse_main",
        parents=[FunctionParent(name="PSBaseParser", type="ClassDef")],
    )

    code_context = get_code_optimization_context(
        function_to_optimize=function_to_optimize,
        project_root=temp_dir,
    )

    original_content = source_file.read_text(encoding="utf-8")
    original_helper_code = {source_file: original_content}
    optimized_code_markdown = CodeStringsMarkdown.parse_markdown_code(
        MOCK_CANDIDATE_MARKDOWN,
    )
    optimized_code_raw = optimized_code_markdown.code_strings[0].code

    # Include the target function and all helpers in the replacement.
    all_names = [function_to_optimize.function_name] + [
        h.only_function_name or h.qualified_name
        for h in code_context.helper_functions
    ]
    updated_source = replace_function_and_helpers(
        source_code=original_content,
        original_function_names=all_names,
        optimized_code=optimized_code_raw,
        preexisting_objects=code_context.preexisting_objects,
        original_helper_code=original_helper_code,
        function_to_optimize=function_to_optimize,
        code_context=code_context,
        optimized_code_markdown=optimized_code_markdown,
        project_root=temp_dir,
    )
    source_file.write_text(updated_source, encoding="utf-8")

    final_content = source_file.read_text(encoding="utf-8")
    return function_to_optimize, code_context, final_content


def test_replace_with_mock_candidate(temp_project):
    """Verify replace_function_and_helpers produces valid output with correct context.

    The code context detects ALL sibling methods as helpers of _parse_main.
    detect_unused_helper_functions correctly recognizes methods referenced via attribute
    assignment (self._parse1 = self._parse_literal) as used, so they are NOT reverted.
    """
    import ast

    _, code_context, final_content = run_replacement(temp_project)

    # Code context correctly detects ALL methods as helpers.
    # The target function itself may also appear in the helpers list.
    helper_names = {h.qualified_name for h in code_context.helper_functions}
    expected_helpers = {
        "PSBaseParser._parse_comment",
        "PSBaseParser._parse_literal",
        "PSBaseParser._parse_number",
        "PSBaseParser._parse_float",
        "PSBaseParser._parse_keyword",
        "PSBaseParser._parse_string",
        "PSBaseParser._parse_wopen",
        "PSBaseParser._parse_wclose",
        "PSBaseParser._add_token",
        "KWD",
    }
    assert expected_helpers.issubset(helper_names)

    # The replacement must produce valid Python.
    ast.parse(final_content)

    # The optimized _parse_main body should be present.
    assert "Use integer byte access" in final_content


def test_detect_unused_helpers_handles_attribute_refs(temp_project):
    """Verify detect_unused_helper_functions recognizes methods referenced via attribute assignment.

    When _parse_main does `self._parse1 = self._parse_literal`, the method is referenced as
    an ast.Attribute value (not an ast.Call). The detection should recognize these as used.
    """
    temp_dir, source_file = temp_project

    function_to_optimize = FunctionToOptimize(
        file_path=source_file,
        function_name="_parse_main",
        parents=[FunctionParent(name="PSBaseParser", type="ClassDef")],
    )

    code_context = get_code_optimization_context(
        function_to_optimize=function_to_optimize,
        project_root=temp_dir,
    )

    optimized_code = CodeStringsMarkdown.parse_markdown_code(
        MOCK_CANDIDATE_MARKDOWN,
    )

    unused_helpers = detect_unused_helper_functions(
        function_to_optimize, code_context, optimized_code
    )
    unused_names = {h.qualified_name for h in unused_helpers}

    # The function being optimized itself may appear as an unused
    # helper due to context extraction including it; exclude it.
    # All real helpers should be detected as used — they are either
    # directly called or referenced via attribute assignment
    # (self._parse1 = self._parse_X).
    unused_names.discard("PSBaseParser._parse_main")
    assert unused_names == set(), (
        f"Expected no unused helpers, got: {unused_names}"
    )


def test_replace_produces_valid_python(temp_project):
    """Verify the final output is valid, parseable Python."""
    _, _, final_content = run_replacement(temp_project)

    import ast

    ast.parse(final_content)
