- test_instrument_codeflash_trace.py

from __future__ import annotations

from dataclasses import dataclass, field

import libcst as cst
# imports
import pytest  # used for our unit tests
from codeflash.benchmarking.instrument_codeflash_trace import \
    add_codeflash_decorator_to_code


# --- Minimal stub for FunctionToOptimize for testing ---
@dataclass
class Parent:
    type: str
    name: str

@dataclass
class FunctionToOptimize:
    function_name: str
    parents: list[Parent] = field(default_factory=list)

# --- Minimal AddDecoratorTransformer for testing ---
class AddDecoratorTransformer(cst.CSTTransformer):
    def __init__(self, target_functions):
        self.target_functions = target_functions

    def leave_FunctionDef(self, original_node, updated_node):
        # Get class name if inside a class
        class_name = ""
        parent = getattr(original_node, "parent", None)
        # libcst doesn't link parents by default, so we rely on the transformer input
        # Instead, we use the function name and class name from the target_functions set
        # So we just match by name
        fn_name = original_node.name.value
        # Try to get class name from the decorator context, but since we don't have it,
        # we just check for both (class_name, fn_name) and ("", fn_name)
        for (cls, name) in self.target_functions:
            if name == fn_name:
                if cls == "":
                    # Top-level function
                    break
                # If a class, check if the parent is a classdef
                parent = getattr(original_node, "parent", None)
                while parent:
                    if isinstance(parent, cst.ClassDef) and parent.name.value == cls:
                        break
                    parent = getattr(parent, "parent", None)
                else:
                    continue
                break
        else:
            return updated_node

        # Check if already decorated with codeflash_trace
        for dec in updated_node.decorators:
            if isinstance(dec.decorator, cst.Name) and dec.decorator.value == "codeflash_trace":
                return updated_node
            if (
                isinstance(dec.decorator, cst.Attribute)
                and dec.decorator.attr.value == "codeflash_trace"
            ):
                return updated_node
        # Add the decorator
        new_decorator = cst.Decorator(decorator=cst.Name("codeflash_trace"))
        return updated_node.with_changes(decorators=[new_decorator] + list(updated_node.decorators))
from codeflash.benchmarking.instrument_codeflash_trace import \
    add_codeflash_decorator_to_code

# unit tests

# -------------------------
# 1. BASIC TEST CASES
# -------------------------

def test_add_decorator_simple_function():
    # Single top-level function, no decorators
    code = "def foo():\n    pass"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 331μs -> 335μs

def test_add_decorator_already_present():
    # Function already has the decorator, should not duplicate
    code = "@codeflash_trace\ndef foo():\n    pass"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 379μs -> 384μs

def test_add_decorator_multiple_functions():
    # Multiple functions, only one should get the decorator
    code = "def foo():\n    pass\ndef bar():\n    pass"
    functions = [FunctionToOptimize(function_name="bar")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 469μs -> 478μs

def test_add_decorator_class_method():
    # Add to a method inside a class
    code = "class A:\n    def foo(self):\n        pass"
    functions = [FunctionToOptimize(function_name="foo", parents=[Parent(type="ClassDef", name="A")])]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 459μs -> 467μs

def test_add_decorator_staticmethod():
    # Add to a static method with existing @staticmethod decorator
    code = "class A:\n    @staticmethod\n    def foo():\n        pass"
    functions = [FunctionToOptimize(function_name="foo", parents=[Parent(type="ClassDef", name="A")])]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 467μs -> 470μs

def test_add_decorator_no_functions():
    # No functions to optimize, code should remain unchanged
    code = "def foo():\n    pass"
    codeflash_output = add_codeflash_decorator_to_code(code, []); result = codeflash_output # 241μs -> 107μs

# -------------------------
# 2. EDGE TEST CASES
# -------------------------

def test_add_decorator_function_with_existing_other_decorator():
    # Function has a different decorator, codeflash_trace should be added above it
    code = "@other_decorator\ndef foo():\n    pass"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 370μs -> 382μs

def test_add_decorator_nested_class_method():
    # Method inside nested class, only add if parents match
    code = "class A:\n    class B:\n        def foo(self):\n            pass"
    # Only add if class is B
    functions = [FunctionToOptimize(function_name="foo", parents=[Parent(type="ClassDef", name="B")])]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 317μs -> 321μs

def test_add_decorator_function_with_attribute_decorator():
    # Function has a decorator like @module.codeflash_trace, should not add another
    code = "@module.codeflash_trace\ndef foo():\n    pass"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 431μs -> 438μs

def test_add_decorator_function_with_args_in_decorator():
    # Function has a decorator with arguments, codeflash_trace should still be added
    code = "@decorator(arg=1)\ndef foo():\n    pass"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 525μs -> 537μs

def test_add_decorator_function_with_multiline_signature():
    # Function with multiline signature
    code = "def foo(\n    a,\n    b\n):\n    pass"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 496μs -> 507μs

def test_add_decorator_unicode_function_name():
    # Function name with unicode
    code = "def föö():\n    pass"
    functions = [FunctionToOptimize(function_name="föö")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 326μs -> 324μs

def test_add_decorator_function_with_docstring():
    # Function with a docstring
    code = 'def foo():\n    """Docstring"""\n    pass'
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 394μs -> 401μs

def test_add_decorator_empty_code():
    # Empty code string should not fail
    code = ""
    codeflash_output = add_codeflash_decorator_to_code(code, []); result = codeflash_output # 56.7μs -> 41.6μs

def test_add_decorator_function_with_type_annotations():
    # Function with type annotations
    code = "def foo(a: int, b: str) -> bool:\n    return True"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 550μs -> 567μs

def test_add_decorator_function_with_complex_signature():
    # Function with *args, **kwargs, defaults, and type hints
    code = "def foo(a: int = 1, *args, b=None, **kwargs) -> None:\n    pass"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 663μs -> 672μs

def test_add_decorator_function_with_blank_lines():
    # Function with blank lines before/after
    code = "\n\ndef foo():\n    pass\n\n"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 378μs -> 373μs

def test_add_decorator_function_with_comment_above():
    # Function with a comment above
    code = "# comment\ndef foo():\n    pass"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 349μs -> 352μs

def test_add_decorator_function_with_multiple_decorators():
    # Function with multiple decorators
    code = "@a\n@b\n@c\ndef foo():\n    pass"
    functions = [FunctionToOptimize(function_name="foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 461μs -> 459μs

def test_add_decorator_function_name_collision():
    # Two functions with same name in different classes
    code = (
        "class A:\n    def foo(self):\n        pass\n"
        "class B:\n    def foo(self):\n        pass"
    )
    # Only add to class B
    functions = [FunctionToOptimize(function_name="foo", parents=[Parent(type="ClassDef", name="B")])]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 733μs -> 741μs

# -------------------------
# 3. LARGE SCALE TEST CASES
# -------------------------

def test_add_decorator_many_functions():
    # Large number of top-level functions
    code = "\n".join([f"def foo{i}():\n    pass" for i in range(100)])
    functions = [FunctionToOptimize(function_name=f"foo{i}") for i in range(100)]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 14.2ms -> 14.5ms

def test_add_decorator_many_classes_and_methods():
    # Large number of classes, each with one method
    code = "\n".join(
        [f"class C{i}:\n    def m(self):\n        pass" for i in range(50)]
    )
    functions = [
        FunctionToOptimize(function_name="m", parents=[Parent(type="ClassDef", name=f"C{i}")])
        for i in range(50)
    ]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 13.1ms -> 13.4ms

def test_add_decorator_large_mixed_code():
    # Mix of classes, functions, and methods, some to decorate, some not
    code = ""
    functions = []
    for i in range(20):
        code += f"def f{i}():\n    pass\n"
        if i % 2 == 0:
            functions.append(FunctionToOptimize(function_name=f"f{i}"))
    for i in range(10):
        code += f"class X{i}:\n"
        for j in range(5):
            code += f"    def m{j}(self):\n        pass\n"
            if j % 2 == 1:
                functions.append(
                    FunctionToOptimize(
                        function_name=f"m{j}",
                        parents=[Parent(type="ClassDef", name=f"X{i}")]
                    )
                )
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 12.2ms -> 12.4ms

def test_add_decorator_performance():
    # Stress test with 500 functions
    code = "\n".join([f"def foo{i}():\n    pass" for i in range(500)])
    functions = [FunctionToOptimize(function_name=f"foo{i}") for i in range(500)]
    codeflash_output = add_codeflash_decorator_to_code(code, functions); result = codeflash_output # 70.6ms -> 72.9ms
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Tuple

import libcst as cst
# imports
import pytest  # used for our unit tests
from codeflash.benchmarking.instrument_codeflash_trace import \
    add_codeflash_decorator_to_code

# --- Minimal stub for FunctionToOptimize and its parent structure ---

@dataclass
class Parent:
    type: str
    name: str

@dataclass
class FunctionToOptimize:
    function_name: str
    parents: List[Parent]

# --- Transformer for adding the decorator ---

class AddDecoratorTransformer(cst.CSTTransformer):
    def __init__(self, target_functions: set[Tuple[str, str]]):
        self.target_functions = target_functions

    def leave_FunctionDef(
        self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
    ) -> cst.FunctionDef:
        # Determine class context, if any
        class_name = ""
        parent = self.get_metadata(cst.metadata.ParentNodeProvider, original_node, None)
        while parent:
            if isinstance(parent, cst.ClassDef):
                class_name = parent.name.value
                break
            parent = self.get_metadata(cst.metadata.ParentNodeProvider, parent, None)
        # Check if this function is a target
        if (class_name, original_node.name.value) in self.target_functions:
            # Check if decorator already present
            for dec in updated_node.decorators:
                if (
                    isinstance(dec.decorator, cst.Name)
                    and dec.decorator.value == "codeflash_trace"
                ):
                    return updated_node  # Already decorated
                if (
                    isinstance(dec.decorator, cst.Attribute)
                    and dec.decorator.attr.value == "codeflash_trace"
                ):
                    return updated_node  # Already decorated
            # Add the decorator
            return updated_node.with_changes(
                decorators=[
                    cst.Decorator(decorator=cst.Name("codeflash_trace"))
                ] + list(updated_node.decorators)
            )
        return updated_node

    @staticmethod
    def get_metadata(provider, node, default=None):
        # Dummy implementation for test, as metadata is not used in this context
        return None
from codeflash.benchmarking.instrument_codeflash_trace import \
    add_codeflash_decorator_to_code

# --- Unit Tests ---

# Helper to create FunctionToOptimize with/without class context
def make_fopt(func_name, class_name=None):
    if class_name:
        parents = [Parent(type="ClassDef", name=class_name)]
    else:
        parents = []
    return FunctionToOptimize(function_name=func_name, parents=parents)

# 1. BASIC TEST CASES

def test_add_decorator_to_simple_function():
    # Basic: Add to top-level function
    code = "def foo():\n    pass\n"
    fopts = [make_fopt("foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 319μs -> 333μs

def test_add_decorator_to_method_in_class():
    # Basic: Add to method in class
    code = "class Bar:\n    def baz(self):\n        return 1\n"
    fopts = [make_fopt("baz", "Bar")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 499μs -> 515μs

def test_add_decorator_to_multiple_functions():
    # Basic: Add to multiple functions
    code = (
        "def foo():\n    pass\n"
        "def bar():\n    pass\n"
    )
    fopts = [make_fopt("foo"), make_fopt("bar")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 479μs -> 493μs

def test_no_decorator_if_function_not_in_list():
    # Basic: Don't add if not in list
    code = "def foo():\n    pass\n"
    fopts = [make_fopt("bar")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 239μs -> 243μs

def test_add_decorator_to_static_and_class_methods():
    # Basic: Add decorator to static and class methods as well
    code = (
        "class C:\n"
        "    @staticmethod\n"
        "    def foo():\n"
        "        pass\n"
        "    @classmethod\n"
        "    def bar(cls):\n"
        "        pass\n"
    )
    fopts = [make_fopt("foo", "C"), make_fopt("bar", "C")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 706μs -> 738μs

# 2. EDGE TEST CASES

def test_already_decorated_function():
    # Edge: Should not add duplicate decorator
    code = "@codeflash_trace\ndef foo():\n    pass\n"
    fopts = [make_fopt("foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 368μs -> 376μs

def test_function_with_other_decorators():
    # Edge: Should add codeflash_trace before other decorators
    code = (
        "@other_decorator\n"
        "def foo():\n"
        "    pass\n"
    )
    fopts = [make_fopt("foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 369μs -> 376μs

def test_function_with_attribute_decorator():
    # Edge: Should not add if decorator is codeflash_trace as attribute
    code = "@some.codeflash_trace\ndef foo():\n    pass\n"
    fopts = [make_fopt("foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 419μs -> 433μs

def test_nested_functions():
    # Edge: Should only decorate top-level or class methods, not inner functions
    code = (
        "def outer():\n"
        "    def inner():\n"
        "        pass\n"
        "    return inner\n"
    )
    fopts = [make_fopt("outer"), make_fopt("inner")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 431μs -> 446μs

def test_function_in_inherited_class():
    # Edge: Should only decorate if class name matches
    code = (
        "class Base:\n"
        "    def foo(self):\n"
        "        pass\n"
        "class Child(Base):\n"
        "    def foo(self):\n"
        "        pass\n"
    )
    fopts = [make_fopt("foo", "Child")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 827μs -> 854μs

def test_function_with_multiline_signature():
    # Edge: Handles multiline signatures
    code = (
        "def foo(\n"
        "    a,\n"
        "    b\n"
        "):\n"
        "    pass\n"
    )
    fopts = [make_fopt("foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 488μs -> 502μs

def test_function_with_unicode_name():
    # Edge: Unicode function name
    code = "def f_áéíóú():\n    pass\n"
    fopts = [make_fopt("f_áéíóú")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 323μs -> 326μs


def test_no_functions_in_code():
    # Edge: Code with no functions
    code = "x = 42\nprint(x)\n"
    fopts = [make_fopt("foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 347μs -> 408μs

def test_function_with_type_hints_and_annotations():
    # Edge: Function with type hints and return annotation
    code = "def foo(a: int, b: str) -> bool:\n    return True\n"
    fopts = [make_fopt("foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 547μs -> 596μs

def test_function_with_docstring():
    # Edge: Function with a docstring
    code = (
        "def foo():\n"
        "    '''This is a docstring.'''\n"
        "    pass\n"
    )
    fopts = [make_fopt("foo")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 390μs -> 415μs

def test_class_method_with_same_name_in_two_classes():
    # Edge: Only decorate the method in the specified class
    code = (
        "class A:\n"
        "    def foo(self):\n"
        "        pass\n"
        "class B:\n"
        "    def foo(self):\n"
        "        pass\n"
    )
    fopts = [make_fopt("foo", "B")]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 722μs -> 754μs
    # Only B.foo should be decorated
    a_section = result.split("class A:")[1].split("class B:")[0]
    b_section = result.split("class B:")[1]

# 3. LARGE SCALE TEST CASES

def test_large_number_of_functions():
    # Large: 500 top-level functions, decorate 100 of them
    N = 500
    code = "\n".join([f"def func_{i}():\n    return {i}" for i in range(N)])
    fopts = [make_fopt(f"func_{i}") for i in range(0, 100)]
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 78.0ms -> 81.3ms
    # All decorated functions should have the decorator right above
    for i in range(0, 100):
        func_def = f"def func_{i}():"
        idx = result.index(func_def)
        before = result[:idx].rstrip().splitlines()[-1]

def test_large_number_of_classes_and_methods():
    # Large: 50 classes, each with 10 methods, decorate 5 methods per class
    N_CLASSES = 50
    N_METHODS = 10
    code = ""
    fopts = []
    for c in range(N_CLASSES):
        code += f"class C{c}:\n"
        for m in range(N_METHODS):
            code += f"    def m{m}(self):\n        return {m}\n"
            if m < 5:
                fopts.append(make_fopt(f"m{m}", f"C{c}"))
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 104ms -> 108ms
    # Check that for each class, the first 5 methods have the decorator
    for c in range(N_CLASSES):
        class_code = result.split(f"class C{c}:")[1]
        for m in range(5):
            method_def = f"def m{m}(self):"
            idx = class_code.index(method_def)
            before = class_code[:idx].rstrip().splitlines()[-1]

def test_large_code_with_no_targets():
    # Large: 1000 functions, none to decorate
    N = 1000
    code = "\n".join([f"def f{i}():\n    return {i}" for i in range(N)])
    fopts = []
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 153ms -> 51.5ms

def test_performance_large_mix():
    # Large: Mix of classes, functions, some to decorate, some not
    N_FUNCS = 300
    N_CLASSES = 20
    N_METHODS = 20
    code = ""
    fopts = []
    # Top-level functions
    for i in range(N_FUNCS):
        code += f"def top_{i}():\n    return {i}\n"
        if i % 10 == 0:
            fopts.append(make_fopt(f"top_{i}"))
    # Classes and methods
    for c in range(N_CLASSES):
        code += f"class Cls{c}:\n"
        for m in range(N_METHODS):
            code += f"    def meth_{m}(self):\n        return {m}\n"
            if m % 5 == 0:
                fopts.append(make_fopt(f"meth_{m}", f"Cls{c}"))
    codeflash_output = add_codeflash_decorator_to_code(code, fopts); result = codeflash_output # 126ms -> 131ms
    # Should have correct number of decorators
    expected = (N_FUNCS // 10) + (N_CLASSES * (N_METHODS // 5))
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.