#!/usr/bin/env python3
"""
Test sanitize function with duplicate function definitions
"""

import sys
sys.path.append('/home/ubuntu/RLVR/TestTime-RLVR-v2/evaluation/code_eval/coding/evalplus')

from evalplus.sanitize import sanitize

# Case 1: LLM이 함수 전체를 생성한 경우 (중복 정의)
duplicate_code = '''from typing import List


def has_close_elements(numbers: List[float], threshold: float) -> bool:
    """ Check if in given list of numbers, are any two numbers closer to each other than
    given threshold.
    >>> has_close_elements([1.0, 2.0, 3.0], 0.5)
    False
    >>> has_close_elements([1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)
    True
    """
    def has_close_elements(numbers: List[float], threshold: float) -> bool:
        for i in range(len(numbers) - 1):
            for j in range(i + 1, len(numbers)):
                if abs(numbers[i] - numbers[j]) < threshold:
                    return True
        return False'''

# Case 2: LLM이 함수 본문만 생성한 경우 (정상)
normal_code = '''from typing import List


def has_close_elements(numbers: List[float], threshold: float) -> bool:
    """ Check if in given list of numbers, are any two numbers closer to each other than
    given threshold.
    >>> has_close_elements([1.0, 2.0, 3.0], 0.5)
    False
    >>> has_close_elements([1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)
    True
    """
    for i in range(len(numbers) - 1):
        for j in range(i + 1, len(numbers)):
            if abs(numbers[i] - numbers[j]) < threshold:
                return True
    return False'''

print("=== Test 1: Duplicate function definition ===")
print("Original code:")
print(duplicate_code)
print("\n" + "="*50 + "\n")

sanitized_duplicate = sanitize(duplicate_code, entrypoint="has_close_elements")
print("Sanitized code:")
print(sanitized_duplicate)
print("\n" + "="*50 + "\n")

print("=== Test 2: Normal function definition ===")
print("Original code:")
print(normal_code)
print("\n" + "="*50 + "\n")

sanitized_normal = sanitize(normal_code, entrypoint="has_close_elements")
print("Sanitized code:")
print(sanitized_normal)

# 실제 실행 테스트
print("\n" + "="*50 + "\n")
print("=== Execution Test ===")

try:
    exec(duplicate_code)
    print("Duplicate code executed successfully (unexpected!)")
except Exception as e:
    print(f"Duplicate code failed: {type(e).__name__}: {e}")

try:
    exec(sanitized_duplicate)
    print("Sanitized duplicate code executed successfully")
except Exception as e:
    print(f"Sanitized duplicate code failed: {type(e).__name__}: {e}")

try:
    exec(normal_code)
    print("Normal code executed successfully")
except Exception as e:
    print(f"Normal code failed: {type(e).__name__}: {e}")