from __future__ import annotations

import gc
import os
import time

import pytest
from inference.core.telemetry import _ForceTraceRootSampler


def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, *args, **kwargs):
    test_id = f'{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}'
    if not hasattr(codeflash_wrap, 'index'):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f'{codeflash_line_id}_{codeflash_test_index}'
    test_stdout_tag = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}"
    print(f'!$######{test_stdout_tag}######$!')
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    print(f'!######{test_stdout_tag}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

class SimpleDelegateWithDescription:
    """A simple delegate that implements get_description()."""

    def __init__(self, description: str) -> None:
        self._description = description

    def get_description(self) -> str:
        return self._description

class AnotherDelegate:
    """Another delegate implementation for diversity."""

    def get_description(self) -> str:
        return 'AnotherSampler'

class NestedDelegateDescription:
    """A delegate that returns a complex nested description."""

    def get_description(self) -> str:
        return 'ParentBased(always_on)'

def test_get_description_basic():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test basic get_description with simple delegate.'
    delegate = SimpleDelegateWithDescription('ratio_sampler')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_basic', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_returns_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that get_description returns a string type.'
    delegate = SimpleDelegateWithDescription('test_delegate')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_returns_string', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_with_different_delegate():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description with a different delegate.'
    delegate = AnotherDelegate()
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_different_delegate', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_with_complex_nested_description():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description when delegate has complex nested description.'
    delegate = NestedDelegateDescription()
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_complex_nested_description', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_contains_wrapper_prefix():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that result always contains the ForceTraceRootSampler wrapper.'
    delegate = SimpleDelegateWithDescription('any_sampler')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_contains_wrapper_prefix', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_preserves_delegate_description():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "Test that the delegate's description is preserved exactly."
    test_descriptions = ['simple', 'complex_name_with_underscores', '123numeric', 'MixedCase']
    for desc in test_descriptions:
        delegate = SimpleDelegateWithDescription(desc)
        sampler = _ForceTraceRootSampler(delegate)
        codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_preserves_delegate_description', 'get_description', '2_2', codeflash_loop_index)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_get_description_with_empty_string_delegate():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description when delegate returns empty string.'
    delegate = SimpleDelegateWithDescription('')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_empty_string_delegate', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_with_special_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description when delegate description has special characters.'
    delegate = SimpleDelegateWithDescription('sampler@1.0-beta_v2')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_special_characters', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_with_whitespace_in_delegate():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description when delegate description contains whitespace.'
    delegate = SimpleDelegateWithDescription('sampler with spaces')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_whitespace_in_delegate', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_with_parentheses_in_delegate():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description when delegate description has nested parentheses.'
    delegate = SimpleDelegateWithDescription('ParentBased(sampler)')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_parentheses_in_delegate', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_with_deeply_nested_parentheses():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description with multiple levels of nested parentheses.'
    delegate = SimpleDelegateWithDescription('Outer(Middle(Inner()))')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_deeply_nested_parentheses', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_with_very_long_delegate_description():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description with a very long delegate description.'
    long_description = 'x' * 1000
    delegate = SimpleDelegateWithDescription(long_description)
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_very_long_delegate_description', 'get_description', '4', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = f'ForceTraceRootSampler({long_description})'

def test_get_description_with_numeric_delegate_description():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description when delegate returns numeric-like string.'
    delegate = SimpleDelegateWithDescription('1.5')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_numeric_delegate_description', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_with_unicode_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description with unicode characters in delegate description.'
    delegate = SimpleDelegateWithDescription('sampler_αβγ')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_unicode_characters', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_with_newlines_in_delegate():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description when delegate description contains newlines.'
    delegate = SimpleDelegateWithDescription('line1\nline2')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_newlines_in_delegate', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_consistency():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that multiple calls to get_description return same result.'
    delegate = SimpleDelegateWithDescription('consistent_sampler')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_consistency', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result1 = codeflash_output
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_consistency', 'get_description', '5', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result2 = codeflash_output
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_consistency', 'get_description', '7', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result3 = codeflash_output

def test_get_description_different_instances_same_delegate():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that different sampler instances with same delegate work correctly.'
    delegate1 = SimpleDelegateWithDescription('test')
    delegate2 = SimpleDelegateWithDescription('test')
    sampler1 = _ForceTraceRootSampler(delegate1)
    sampler2 = _ForceTraceRootSampler(delegate2)
    codeflash_return_value = codeflash_wrap(sampler1.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_different_instances_same_delegate', 'get_description', '5', codeflash_loop_index)
    codeflash_output = codeflash_return_value

def test_get_description_format_structure():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that get_description follows the expected format structure.'
    delegate = SimpleDelegateWithDescription('delegate_desc')
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_format_structure', 'get_description', '3', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    content = result[len('ForceTraceRootSampler('):-1]

def test_get_description_with_many_sequential_calls():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test get_description performance with many sequential calls.'
    delegate = SimpleDelegateWithDescription('perf_test_sampler')
    sampler = _ForceTraceRootSampler(delegate)
    results = []
    for _ in range(1000):
        codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_many_sequential_calls', 'get_description', '4_0', codeflash_loop_index)

def test_get_description_with_many_sampler_instances():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test creating many sampler instances and calling get_description.'
    samplers = []
    for i in range(1000):
        delegate = SimpleDelegateWithDescription(f'sampler_{i}')
        samplers.append(_ForceTraceRootSampler(delegate))
    for i, sampler in enumerate(samplers):
        expected = f'ForceTraceRootSampler(sampler_{i})'
        codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_many_sampler_instances', 'get_description', '3_1', codeflash_loop_index)
        codeflash_output = codeflash_return_value

def test_get_description_with_large_description_strings():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test performance with increasingly large delegate descriptions.'
    size_multipliers = [1, 10, 100, 1000]
    base_desc = 'sampler_'
    for multiplier in size_multipliers:
        desc = base_desc * multiplier
        delegate = SimpleDelegateWithDescription(desc)
        sampler = _ForceTraceRootSampler(delegate)
        codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_large_description_strings', 'get_description', '3_3', codeflash_loop_index)
        codeflash_output = codeflash_return_value
        result = codeflash_output
        expected = f'ForceTraceRootSampler({desc})'

def test_get_description_with_deeply_nested_delegates():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with chain of delegating samplers.'
    innermost = SimpleDelegateWithDescription('base')
    samplers = [_ForceTraceRootSampler(innermost)]

    class SamplerAsDelegate:

        def __init__(self, sampler):
            self._sampler = sampler

        def get_description(self):
            return self._sampler.get_description()
    for _ in range(100):
        wrapped = SamplerAsDelegate(samplers[-1])
        samplers.append(_ForceTraceRootSampler(wrapped))
    codeflash_return_value = codeflash_wrap(samplers[-1].get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_deeply_nested_delegates', 'get_description', '5', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_description_memory_efficiency():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "Test that get_description doesn't leak memory with large operations."
    delegates = [SimpleDelegateWithDescription(f'sampler_{i}') for i in range(500)]
    samplers = [_ForceTraceRootSampler(d) for d in delegates]
    for _ in range(10):
        descriptions = [s.get_description() for s in samplers]

def test_get_description_with_special_character_stress():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test handling many special characters at scale.'
    special_chars = "!@#$%^&*()-_=+[]{}|;:',.<>?/~`"
    delegate = SimpleDelegateWithDescription(special_chars * 10)
    sampler = _ForceTraceRootSampler(delegate)
    codeflash_return_value = codeflash_wrap(sampler.get_description, 'tests.inference.unit_tests.test_get_description__unit_test_1', None, 'test_get_description_with_special_character_stress', 'get_description', '4', codeflash_loop_index)
    codeflash_output = codeflash_return_value
    result = codeflash_output