from __future__ import annotations

import ast
import os
import tempfile
from pathlib import Path

import pytest

from codeflash_python._model import (
    FunctionParent,
    FunctionToOptimize,
    TestingMode,
)
from codeflash_python.test_discovery.models import CodePosition
from codeflash_python.testing._instrument_core import FunctionImportedAsVisitor
from codeflash_python.testing._instrumentation import (
    inject_profiling_into_existing_test,
)

project_root = Path(__file__).parent.resolve()


@pytest.fixture
def tmp_dir():
    """Create a temporary directory for test results."""
    with tempfile.TemporaryDirectory() as tmpdirname:
        yield Path(tmpdirname)


def _assert_sync_instrumentation_present(source: str) -> None:
    """Assert that the new sync instrumentation markers are present."""
    assert "_codeflash_call_site.set(" in source
    assert "from codeflash_async_wrapper import _codeflash_call_site" in source


def _assert_old_instrumentation_absent(source: str) -> None:
    """Assert that old sync instrumentation artifacts are absent."""
    assert "codeflash_wrap" not in source
    assert "codeflash_con" not in source
    assert "codeflash_cur" not in source
    assert "import sqlite3" not in source


def test_perfinjector_bubble_sort(tmp_dir) -> None:
    """Instrument a unittest bubble sort test with profiling."""
    code = """import unittest

from code_to_optimize.bubble_sort import sorter


class TestPigLatin(unittest.TestCase):
    def test_sort(self):
        input = [5, 4, 3, 2, 1, 0]
        output = sorter(input)
        self.assertEqual(output, [0, 1, 2, 3, 4, 5])

        input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
        output = sorter(input)
        self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0])

        input = list(reversed(range(5000)))
        self.assertEqual(sorter(input), list(range(5000)))
"""

    with (tmp_dir / "test_sort.py").open("w") as f:
        f.write(code)
        f.flush()
        func = FunctionToOptimize(
            function_name="sorter", parents=(), file_path=Path(f.name)
        )
        original_cwd = Path.cwd()
        run_cwd = project_root
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            Path(f.name),
            [CodePosition(9, 17), CodePosition(13, 17), CodePosition(17, 17)],
            func,
            Path(f.name).parent,
        )
        os.chdir(original_cwd)
    assert success
    assert new_test is not None
    _assert_sync_instrumentation_present(new_test)
    _assert_old_instrumentation_absent(new_test)
    assert "sorter" in new_test
    assert "class TestPigLatin" in new_test
    assert "def test_sort" in new_test
    count = new_test.count("_codeflash_call_site.set(")
    assert count == 3


def test_perfinjector_only_replay_test(tmp_dir) -> None:
    """Instrument a replay test with profiling."""
    code = """import dill as pickle
import pytest
from codeflash.tracing.replay_test import get_next_arg_and_return
from codeflash.validation.equivalence import compare_results
from packagename.ml.yolo.image_reshaping_utils import prepare_image_for_yolo as packagename_ml_yolo_image_reshaping_utils_prepare_image_for_yolo
def test_prepare_image_for_yolo():
    for arg_val_pkl, return_val_pkl in get_next_arg_and_return('/home/saurabh/packagename/traces/first.trace', 3):
        args = pickle.loads(arg_val_pkl)
        return_val_1= pickle.loads(return_val_pkl)
        ret = packagename_ml_yolo_image_reshaping_utils_prepare_image_for_yolo(**args)
        assert compare_results(return_val_1, ret)
"""

    with (tmp_dir / "test_return_values.py").open("w") as f:
        f.write(code)
        f.flush()
        func = FunctionToOptimize(
            function_name="prepare_image_for_yolo",
            parents=(),
            file_path=Path("module.py"),
        )
        original_cwd = Path.cwd()
        run_cwd = project_root
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            Path(f.name), [CodePosition(10, 14)], func, Path(f.name).parent
        )
        os.chdir(original_cwd)
    assert success
    assert new_test is not None
    _assert_sync_instrumentation_present(new_test)
    _assert_old_instrumentation_absent(new_test)
    assert (
        "packagename_ml_yolo_image_reshaping_utils_prepare_image_for_yolo"
        in new_test
    )
    assert "def test_prepare_image_for_yolo" in new_test
    assert "_codeflash_call_site.set(" in new_test


def test_perfinjector_bubble_sort_results() -> None:
    """Instrument bubble sort and verify output structure."""
    code = """from code_to_optimize.bubble_sort import sorter
import datetime


def test_sort():
    input = [5, 4, 3, 2, 1, 0]
    print(datetime.datetime.now().isoformat())
    output = sorter(input)
    assert output == [0, 1, 2, 3, 4, 5]

    input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
    output = sorter(input)
    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]"""

    test_path = (
        project_root
        / "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_results_temp.py"
    ).resolve()
    try:
        with test_path.open("w") as f:
            f.write(code)
        code_path = (
            project_root / "code_to_optimize/bubble_sort.py"
        ).resolve()
        project_root_path = project_root
        original_cwd = Path.cwd()
        run_cwd = project_root
        func = FunctionToOptimize(
            function_name="sorter", parents=(), file_path=code_path
        )
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(8, 14), CodePosition(12, 14)],
            func,
            project_root_path,
            mode=TestingMode.BEHAVIOR,
        )
        os.chdir(original_cwd)
        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "def test_sort" in new_test
        assert "sorter" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 2

        os.chdir(run_cwd)
        success_perf, new_perf_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(8, 14), CodePosition(12, 14)],
            func,
            project_root_path,
            mode=TestingMode.PERFORMANCE,
        )
        os.chdir(original_cwd)
        assert success_perf
        assert new_perf_test is not None
        _assert_sync_instrumentation_present(new_perf_test)
        _assert_old_instrumentation_absent(new_perf_test)
    finally:
        test_path.unlink(missing_ok=True)


def test_perfinjector_bubble_sort_parametrized_results() -> None:
    """Instrument parametrized bubble sort and verify output structure."""
    code = """from code_to_optimize.bubble_sort import sorter
import pytest


@pytest.mark.parametrize(
    "input, expected_output",
    [
        ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
        ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
        (list(reversed(range(50))), list(range(50))),
    ],
)
def test_sort_parametrized(input, expected_output):
    output = sorter(input)
    assert output == expected_output
"""
    code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
    test_path = (
        project_root
        / "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_parametrized_results_temp.py"
    ).resolve()
    try:
        with test_path.open("w") as f:
            f.write(code)

        project_root_path = project_root
        original_cwd = Path.cwd()
        run_cwd = project_root

        func = FunctionToOptimize(
            function_name="sorter", parents=(), file_path=code_path
        )
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(14, 13)],
            func,
            project_root_path,
            mode=TestingMode.BEHAVIOR,
        )
        assert success
        success_perf, new_test_perf = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(14, 13)],
            func,
            project_root_path,
            mode=TestingMode.PERFORMANCE,
        )
        os.chdir(original_cwd)

        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "def test_sort_parametrized" in new_test
        assert "sorter" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 1

        assert success_perf
        assert new_test_perf is not None
        _assert_sync_instrumentation_present(new_test_perf)
        _assert_old_instrumentation_absent(new_test_perf)
    finally:
        test_path.unlink(missing_ok=True)


def test_perfinjector_bubble_sort_parametrized_loop_results() -> None:
    """Instrument parametrized loop bubble sort and verify output structure."""
    code = """from code_to_optimize.bubble_sort import sorter
import pytest


@pytest.mark.parametrize(
    "input, expected_output",
    [
        ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
        ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
        (list(reversed(range(50))), list(range(50))),
    ],
)
def test_sort_parametrized_loop(input, expected_output):
    for i in range(2):
        output = sorter(input)
        assert output == expected_output
"""
    code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
    test_path = (
        project_root
        / "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_parametrized_loop_results_temp.py"
    ).resolve()
    try:
        with test_path.open("w") as f:
            f.write(code)

        project_root_path = project_root
        original_cwd = Path.cwd()
        run_cwd = project_root

        func = FunctionToOptimize(
            function_name="sorter", parents=(), file_path=code_path
        )
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(15, 17)],
            func,
            project_root_path,
            mode=TestingMode.BEHAVIOR,
        )
        assert success
        success_perf, new_test_perf = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(15, 17)],
            func,
            project_root_path,
            mode=TestingMode.PERFORMANCE,
        )
        os.chdir(original_cwd)

        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "def test_sort_parametrized_loop" in new_test
        assert "sorter" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 1

        assert success_perf
        assert new_test_perf is not None
        _assert_sync_instrumentation_present(new_test_perf)
        _assert_old_instrumentation_absent(new_test_perf)
    finally:
        test_path.unlink(missing_ok=True)


def test_perfinjector_bubble_sort_loop_results() -> None:
    """Instrument loop bubble sort and verify output structure."""
    code = """from code_to_optimize.bubble_sort import sorter


def test_sort():
    inputs = [[5, 4, 3, 2, 1, 0], [5.0, 4.0, 3.0, 2.0, 1.0, 0.0], list(reversed(range(50)))]
    expected_outputs = [[0, 1, 2, 3, 4, 5], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0], list(range(50))]

    for i in range(3):
        input = inputs[i]
        expected_output = expected_outputs[i]
        output = sorter(input)
        assert output == expected_output"""

    code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
    test_path = (
        project_root
        / "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_loop_results_temp.py"
    ).resolve()
    try:
        with test_path.open("w") as f:
            f.write(code)

        project_root_path = project_root
        original_cwd = Path.cwd()
        run_cwd = project_root

        func = FunctionToOptimize(
            function_name="sorter", parents=(), file_path=code_path
        )
        os.chdir(str(run_cwd))
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(11, 17)],
            func,
            project_root_path,
            mode=TestingMode.BEHAVIOR,
        )
        assert success
        success_perf, new_test_perf = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(11, 17)],
            func,
            project_root_path,
            mode=TestingMode.PERFORMANCE,
        )
        os.chdir(original_cwd)

        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "def test_sort" in new_test
        assert "sorter" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 1

        assert success_perf
        assert new_test_perf is not None
        _assert_sync_instrumentation_present(new_test_perf)
        _assert_old_instrumentation_absent(new_test_perf)
    finally:
        test_path.unlink(missing_ok=True)


def test_perfinjector_bubble_sort_unittest_results() -> None:
    """Instrument unittest bubble sort and verify output structure."""
    code = """import unittest

from code_to_optimize.bubble_sort import sorter


class TestPigLatin(unittest.TestCase):
    def test_sort(self):
        input = [5, 4, 3, 2, 1, 0]
        output = sorter(input)
        self.assertEqual(output, [0, 1, 2, 3, 4, 5])

        input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
        output = sorter(input)
        self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0])

        input = list(reversed(range(50)))
        output = sorter(input)
        self.assertEqual(output, list(range(50)))
"""

    code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
    test_path = (
        project_root
        / "code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_results_temp.py"
    ).resolve()
    try:
        with test_path.open("w") as f:
            f.write(code)

        project_root_path = project_root
        run_cwd = project_root
        original_cwd = Path.cwd()

        func = FunctionToOptimize(
            function_name="sorter", parents=(), file_path=code_path
        )
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(9, 17), CodePosition(13, 17), CodePosition(17, 17)],
            func,
            project_root_path,
            mode=TestingMode.BEHAVIOR,
        )
        assert success
        success_perf, new_test_perf = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(9, 17), CodePosition(13, 17), CodePosition(17, 17)],
            func,
            project_root_path,
            mode=TestingMode.PERFORMANCE,
        )
        os.chdir(original_cwd)

        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "class TestPigLatin" in new_test
        assert "def test_sort" in new_test
        assert "sorter" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 3

        assert success_perf
        assert new_test_perf is not None
        _assert_sync_instrumentation_present(new_test_perf)
        _assert_old_instrumentation_absent(new_test_perf)
        assert new_test_perf.count("_codeflash_call_site.set(") == 3
    finally:
        test_path.unlink(missing_ok=True)


def test_perfinjector_bubble_sort_unittest_parametrized_results() -> None:
    """Instrument unittest parametrized bubble sort and verify output structure."""
    code = """import unittest
from parameterized import parameterized

from code_to_optimize.bubble_sort import sorter


class TestPigLatin(unittest.TestCase):
    @parameterized.expand(
        [
            ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
            ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
            (list(reversed(range(50))), list(range(50))),
        ]
    )
    def test_sort(self, input, expected_output):
        output = sorter(input)
        self.assertEqual(output, expected_output)
"""

    code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
    test_path = (
        project_root
        / "code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_parametrized_results_temp.py"
    ).resolve()
    try:
        with test_path.open("w") as f:
            f.write(code)
        project_root_path = project_root
        run_cwd = project_root
        original_cwd = Path.cwd()

        func = FunctionToOptimize(
            function_name="sorter", parents=(), file_path=code_path
        )
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(16, 17)],
            func,
            project_root_path,
            mode=TestingMode.BEHAVIOR,
        )
        assert success
        success_perf, new_test_perf = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(16, 17)],
            func,
            project_root_path,
            mode=TestingMode.PERFORMANCE,
        )
        os.chdir(original_cwd)

        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "class TestPigLatin" in new_test
        assert "def test_sort" in new_test
        assert "sorter" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 1

        assert success_perf
        assert new_test_perf is not None
        _assert_sync_instrumentation_present(new_test_perf)
        _assert_old_instrumentation_absent(new_test_perf)
    finally:
        test_path.unlink(missing_ok=True)


def test_perfinjector_bubble_sort_unittest_loop_results() -> None:
    """Instrument unittest loop bubble sort and verify output structure."""
    code = """import unittest

from code_to_optimize.bubble_sort import sorter


class TestPigLatin(unittest.TestCase):
    def test_sort(self):
        inputs = [[5, 4, 3, 2, 1, 0], [5.0, 4.0, 3.0, 2.0, 1.0, 0.0], list(reversed(range(50)))]
        expected_outputs = [[0, 1, 2, 3, 4, 5], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0], list(range(50))]

        for i in range(3):
            input = inputs[i]
            expected_output = expected_outputs[i]
            output = sorter(input)
            self.assertEqual(output, expected_output)"""

    code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
    test_path = (
        project_root
        / "code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_loop_results_temp.py"
    ).resolve()
    try:
        with test_path.open("w") as f:
            f.write(code)

        project_root_path = project_root
        run_cwd = project_root
        original_cwd = Path.cwd()

        func = FunctionToOptimize(
            function_name="sorter", parents=(), file_path=code_path
        )
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(14, 21)],
            func,
            project_root_path,
            mode=TestingMode.BEHAVIOR,
        )
        assert success
        success_perf, new_test_perf = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(14, 21)],
            func,
            project_root_path,
            mode=TestingMode.PERFORMANCE,
        )
        os.chdir(original_cwd)

        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "class TestPigLatin" in new_test
        assert "def test_sort" in new_test
        assert "sorter" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 1

        assert success_perf
        assert new_test_perf is not None
        _assert_sync_instrumentation_present(new_test_perf)
        _assert_old_instrumentation_absent(new_test_perf)
    finally:
        test_path.unlink(missing_ok=True)


def test_perfinjector_bubble_sort_unittest_parametrized_loop_results() -> None:
    """Instrument unittest parametrized loop bubble sort and verify output structure."""
    code = """import unittest
from parameterized import parameterized

from code_to_optimize.bubble_sort import sorter


class TestPigLatin(unittest.TestCase):
    @parameterized.expand(
        [
            ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
            ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
            (list(reversed(range(50))), list(range(50))),
        ]
    )
    def test_sort(self, input, expected_output):
        for i in range(2):
            output = sorter(input)
            self.assertEqual(output, expected_output)
"""

    code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
    test_path = (
        project_root
        / "code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_parametrized_loop_results_temp.py"
    ).resolve()
    try:
        with test_path.open("w") as f:
            f.write(code)
        project_root_path = project_root
        run_cwd = project_root
        original_cwd = Path.cwd()

        func = FunctionToOptimize(
            function_name="sorter", file_path=code_path, parents=()
        )
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(17, 21)],
            func,
            project_root_path,
            mode=TestingMode.BEHAVIOR,
        )
        success_perf, new_test_perf = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(17, 21)],
            func,
            project_root_path,
            mode=TestingMode.PERFORMANCE,
        )
        os.chdir(original_cwd)

        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "class TestPigLatin" in new_test
        assert "def test_sort" in new_test
        assert "sorter" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 1

        assert success_perf
        assert new_test_perf is not None
        _assert_sync_instrumentation_present(new_test_perf)
        _assert_old_instrumentation_absent(new_test_perf)
    finally:
        test_path.unlink(missing_ok=True)


def test_class_method_imported_as() -> None:
    """Detect import aliases for functions and class methods."""
    code = """import functionA
import moduleB as module_B
from module import functionB as function_B
import class_name_B
from nuitka.nodes.ImportNodes import ExpressionBuiltinImport as nuitka_nodes_ImportNodes_ExpressionBuiltinImport
"""
    f = FunctionToOptimize(
        function_name="functionA", file_path=Path("module.py"), parents=()
    )
    tree = ast.parse(code)
    visitor = FunctionImportedAsVisitor(f)
    visitor.visit(tree)
    assert visitor.imported_as.function_name == "functionA"

    f = FunctionToOptimize(
        function_name="functionB", file_path=Path("module.py"), parents=()
    )
    visitor = FunctionImportedAsVisitor(f)
    visitor.visit(tree)
    assert visitor.imported_as.function_name == "function_B"

    f = FunctionToOptimize(
        function_name="method_name",
        file_path=Path("module.py"),
        parents=(FunctionParent("ExpressionBuiltinImport", "ClassDef"),),
    )
    visitor = FunctionImportedAsVisitor(f)
    visitor.visit(tree)
    assert (
        visitor.imported_as.qualified_name
        == "nuitka_nodes_ImportNodes_ExpressionBuiltinImport.method_name"
    )

    f = FunctionToOptimize(
        function_name="class_name_B", file_path=Path("module.py"), parents=()
    )
    visitor = FunctionImportedAsVisitor(f)
    visitor.visit(tree)
    assert visitor.imported_as.qualified_name == "class_name_B"


def test_class_function_instrumentation() -> None:
    """Instrument a class method call in a test function."""
    code = """from module import class_name as class_name_A

def test_class_name_A_function_name():
    ret = class_name_A.function_name(**args)
"""

    test_path = (
        project_root
        / "code_to_optimize/tests/pytest/test_class_function_instrumentation_temp.py"
    )
    try:
        with open(test_path, "w") as f:
            f.write(code)

        project_root_path = project_root / "code_to_optimize/"
        run_cwd = project_root
        original_cwd = Path.cwd()
        func = FunctionToOptimize(
            function_name="function_name",
            file_path=project_root_path / "module.py",
            parents=(FunctionParent("class_name", "ClassDef"),),
        )
        os.chdir(str(run_cwd))
        success, new_test = inject_profiling_into_existing_test(
            test_path, [CodePosition(4, 23)], func, project_root_path
        )
        os.chdir(original_cwd)
    finally:
        test_path.unlink(missing_ok=True)
    assert success
    assert new_test is not None
    _assert_sync_instrumentation_present(new_test)
    _assert_old_instrumentation_absent(new_test)
    assert "class_name_A.function_name" in new_test
    assert "def test_class_name_A_function_name" in new_test
    assert new_test.count("_codeflash_call_site.set(") == 1


def test_wrong_function_instrumentation() -> None:
    """Instrument multiple calls to find_common_tags correctly."""
    code = """from codeflash.result.common_tags import find_common_tags


def test_common_tags_1():
    articles_1 = [1, 2, 3]

    assert find_common_tags(articles_1) == set(1, 2)

    articles_2 = [1, 2]

    assert find_common_tags(articles_2) == set(1)
"""

    test_path = (
        project_root
        / "code_to_optimize/tests/pytest/test_wrong_function_instrumentation_temp.py"
    )
    try:
        with test_path.open("w") as f:
            f.write(code)

        project_root_path = project_root / "code_to_optimize/"
        run_cwd = project_root
        original_cwd = Path.cwd()
        func = FunctionToOptimize(
            function_name="find_common_tags",
            file_path=project_root_path / "module.py",
            parents=(),
        )

        os.chdir(str(run_cwd))
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(7, 11), CodePosition(11, 11)],
            func,
            project_root_path,
        )
        os.chdir(original_cwd)
        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "find_common_tags" in new_test
        assert "def test_common_tags_1" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 2
    finally:
        test_path.unlink(missing_ok=True)


def test_conditional_instrumentation() -> None:
    """Instrument a function call inside an if block."""
    code = """from code_to_optimize.bubble_sort import sorter


def test_sort():
    input = [5, 4, 3, 2, 1, 0]
    if len(input) > 0:
        assert sorter(input) == [0, 1, 2, 3, 4, 5]"""

    test_path = (
        project_root
        / "code_to_optimize/tests/pytest/test_conditional_instrumentation_temp.py"
    )
    try:
        with open(test_path, "w") as f:
            f.write(code)

        project_root_path = project_root / "code_to_optimize/"
        run_cwd = project_root
        original_cwd = Path.cwd()
        func = FunctionToOptimize(
            function_name="sorter",
            file_path=project_root_path / "module.py",
            parents=(),
        )

        os.chdir(str(run_cwd))
        success, new_test = inject_profiling_into_existing_test(
            test_path, [CodePosition(7, 15)], func, project_root_path
        )
        os.chdir(original_cwd)
        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "sorter" in new_test
        assert "def test_sort" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 1
    finally:
        test_path.unlink(missing_ok=True)


def test_static_method_instrumentation():
    """Instrument a static method call (BubbleSorter.sorter)."""
    code = """from code_to_optimize.bubble_sort import BubbleSorter


def test_sort():
    input = [5, 4, 3, 2, 1, 0]
    output = BubbleSorter.sorter(input)
    assert output == [0, 1, 2, 3, 4, 5]

    input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
    output = BubbleSorter.sorter(input)
    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]"""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path=Path(
            "/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py"
        ),
        parents=(FunctionParent("BubbleSorter", "ClassDef"),),
        starting_line=None,
        ending_line=None,
    )

    test_path = (
        project_root
        / "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_results_temp.py"
    )
    try:
        with test_path.open("w") as f:
            f.write(code)
        project_root_path = project_root / "code_to_optimize/"
        run_cwd = project_root
        original_cwd = Path.cwd()

        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(6, 26), CodePosition(10, 26)],
            function_to_optimize,
            project_root_path,
        )
        os.chdir(original_cwd)
        assert success
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "BubbleSorter.sorter" in new_test
        assert "def test_sort" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 2
    finally:
        test_path.unlink(missing_ok=True)


def test_class_method_instrumentation(tmp_path: Path) -> None:
    """Instrument a class method call (Optimizer.get_code_optimization_context)."""
    code = """from codeflash.optimization.optimizer import Optimizer
def test_code_replacement10() -> None:
    get_code_output = '''random code'''
    file_path = Path(__file__).resolve()
    opt = Optimizer(
        Namespace(
            project_root=str(file_path.parent.resolve()),
            disable_telemetry=True,
            tests_root="tests",
            test_framework="pytest",
            pytest_cmd="pytest",
            experiment_id=None,
        ),
    )
    func_top_optimize = FunctionToOptimize(
        function_name="main_method",
        file_path=str(file_path),
        parents=[FunctionParent("MainClass", "ClassDef")],
    )
    with open(file_path) as f:
        original_code = f.read()
        code_context = opt.get_code_optimization_context(
            function_to_optimize=func_top_optimize,
            project_root=str(file_path.parent),
            original_source_code=original_code,
        ).unwrap()
        assert code_context.testgen_context_code == get_code_output
        code_context = opt.get_code_optimization_context(
            function_to_optimize=func_top_optimize,
            project_root=str(file_path.parent),
            original_source_code=original_code,
        )
        assert code_context.testgen_context_code == get_code_output
    """

    test_file_path = tmp_path / "test_class_method_instrumentation.py"
    test_file_path.write_text(code, encoding="utf-8")

    func = FunctionToOptimize(
        function_name="get_code_optimization_context",
        parents=(FunctionParent("Optimizer", "ClassDef"),),
        file_path=test_file_path,
    )
    original_cwd = Path.cwd()
    run_cwd = project_root
    os.chdir(run_cwd)
    success, new_test = inject_profiling_into_existing_test(
        test_file_path,
        [CodePosition(22, 28), CodePosition(28, 28)],
        func,
        test_file_path.parent,
    )
    os.chdir(original_cwd)
    assert success
    assert new_test is not None
    _assert_sync_instrumentation_present(new_test)
    _assert_old_instrumentation_absent(new_test)
    assert "opt.get_code_optimization_context" in new_test
    assert "def test_code_replacement10" in new_test
    assert new_test.count("_codeflash_call_site.set(") >= 1


def test_time_correction_instrumentation() -> None:
    """Instrument parametrized sleep test for performance timing."""
    code = """from code_to_optimize.sleeptime import accurate_sleepfunc
import pytest
@pytest.mark.parametrize("n, expected_total_sleep_time", [
    (0.01, 0.010),
    (0.02, 0.020),
])
def test_sleepfunc_sequence_short(n, expected_total_sleep_time):
    output = accurate_sleepfunc(n)
    assert output == expected_total_sleep_time

"""

    code_path = (project_root / "code_to_optimize/sleeptime.py").resolve()
    test_path = (
        project_root
        / "code_to_optimize/tests/pytest/test_time_correction_instrumentation_temp.py"
    ).resolve()
    try:
        with test_path.open("w") as f:
            f.write(code)

        project_root_path = project_root
        original_cwd = Path.cwd()
        run_cwd = project_root
        func = FunctionToOptimize(
            function_name="accurate_sleepfunc", parents=(), file_path=code_path
        )
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(8, 13)],
            func,
            project_root_path,
            mode=TestingMode.PERFORMANCE,
        )
        os.chdir(original_cwd)

        assert success, "Test instrumentation failed"
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "accurate_sleepfunc" in new_test
        assert "def test_sleepfunc_sequence_short" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 1
    finally:
        test_path.unlink(missing_ok=True)


def test_time_correction_instrumentation_unittest() -> None:
    """Instrument parametrized unittest sleep test for performance timing."""
    code = """import unittest
from parameterized import parameterized

from code_to_optimize.sleeptime import accurate_sleepfunc


class TestPigLatin(unittest.TestCase):
    @parameterized.expand([
        (0.01, 0.010),
        (0.02, 0.020),
    ])
    def test_sleepfunc_sequence_short(self, n, expected_total_sleep_time):
        output = accurate_sleepfunc(n)
"""

    code_path = (project_root / "code_to_optimize/sleeptime.py").resolve()
    test_path = (
        project_root
        / "code_to_optimize/tests/unittest/test_time_correction_instrumentation_unittest_temp.py"
    ).resolve()
    try:
        with test_path.open("w") as f:
            f.write(code)

        project_root_path = project_root
        original_cwd = Path.cwd()
        run_cwd = project_root
        func = FunctionToOptimize(
            function_name="accurate_sleepfunc", parents=(), file_path=code_path
        )
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(13, 17)],
            func,
            project_root_path,
            mode=TestingMode.PERFORMANCE,
        )
        os.chdir(original_cwd)

        assert success, "Test instrumentation failed"
        assert new_test is not None
        _assert_sync_instrumentation_present(new_test)
        _assert_old_instrumentation_absent(new_test)
        assert "accurate_sleepfunc" in new_test
        assert "class TestPigLatin" in new_test
        assert "def test_sleepfunc_sequence_short" in new_test
        assert new_test.count("_codeflash_call_site.set(") == 1
    finally:
        test_path.unlink(missing_ok=True)