from __future__ import annotations

import gc
import os
import time
from dataclasses import dataclass
from typing import List

import pytest
from inference.core.workflows.execution_engine.entities.types import \
    STEP_AS_SELECTED_ELEMENT
from inference.core.workflows.execution_engine.introspection.entities import \
    ParsedSelector
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    step_definition_allows_control_flow_references


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

@dataclass(frozen=True)
class AllowedReference:
    """Represents an allowed reference in a selector definition."""
    selected_element: str
    reference_type: str = 'default'

@dataclass(frozen=True)
class SelectorDefinition:
    """Represents a selector definition with allowed references."""
    allowed_references: List[AllowedReference]
    name: str = 'test_selector'

def test_step_definition_allows_control_flow_references_with_matching_element():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns True when an allowed reference has STEP_AS_SELECTED_ELEMENT.'
    allowed_ref = AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT)
    selector_def = SelectorDefinition(allowed_references=[allowed_ref])
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_matching_element', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_without_matching_element():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns False when no allowed reference has STEP_AS_SELECTED_ELEMENT.'
    allowed_ref = AllowedReference(selected_element='other_element')
    selector_def = SelectorDefinition(allowed_references=[allowed_ref])
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_without_matching_element', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_multiple_references_one_matching():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns True when one of multiple references is STEP_AS_SELECTED_ELEMENT.'
    allowed_refs = [AllowedReference(selected_element='element_1'), AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT), AllowedReference(selected_element='element_2')]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_multiple_references_one_matching', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_multiple_references_none_matching():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns False when multiple references exist but none match.'
    allowed_refs = [AllowedReference(selected_element='element_1'), AllowedReference(selected_element='element_2'), AllowedReference(selected_element='element_3')]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_multiple_references_none_matching', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_empty_references():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns False when allowed_references list is empty.'
    selector_def = SelectorDefinition(allowed_references=[])
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_empty_references', 'step_definition_allows_control_flow_references', '3', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_control_flow_element_first():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns True when STEP_AS_SELECTED_ELEMENT is the first element.'
    allowed_refs = [AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT), AllowedReference(selected_element='other_element')]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_control_flow_element_first', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_control_flow_element_last():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns True when STEP_AS_SELECTED_ELEMENT is the last element.'
    allowed_refs = [AllowedReference(selected_element='other_element'), AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT)]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_control_flow_element_last', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_multiple_control_flow_elements():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns True when multiple references are STEP_AS_SELECTED_ELEMENT.'
    allowed_refs = [AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT), AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT), AllowedReference(selected_element='other_element')]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_multiple_control_flow_elements', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_special_characters_in_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test function with special characters in step names and values.'
    allowed_refs = [AllowedReference(selected_element='element-1'), AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT)]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step@#$', value='test_value!@#', index=0, key='special_key')
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_special_characters_in_names', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_none_index_and_key():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function works correctly when index and key are None.'
    allowed_refs = [AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT)]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_none_index_and_key', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_numeric_index_and_string_key():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function works with various index and key values.'
    allowed_refs = [AllowedReference(selected_element='other')]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=42, key='my_key')
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_numeric_index_and_string_key', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_case_sensitive_comparison():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that the comparison is case-sensitive for the selected_element.'
    allowed_refs = [AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT.lower() if isinstance(STEP_AS_SELECTED_ELEMENT, str) else STEP_AS_SELECTED_ELEMENT)]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_case_sensitive_comparison', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    if isinstance(STEP_AS_SELECTED_ELEMENT, str):
        pass
    else:
        pass

def test_step_definition_allows_control_flow_references_with_large_number_of_references():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function performs efficiently with many allowed references.'
    allowed_refs = [AllowedReference(selected_element=f'element_{i}') for i in range(999)]
    allowed_refs.append(AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT))
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_large_number_of_references', 'step_definition_allows_control_flow_references', '5', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_large_number_of_references_at_start():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function performs efficiently when control flow element is first.'
    allowed_refs = [AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT)]
    allowed_refs.extend([AllowedReference(selected_element=f'element_{i}') for i in range(999)])
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_large_number_of_references_at_start', 'step_definition_allows_control_flow_references', '5', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_many_references_no_match():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function handles large lists when no control flow element exists.'
    allowed_refs = [AllowedReference(selected_element=f'element_{i}') for i in range(1000)]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_many_references_no_match', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_many_control_flow_elements():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function handles many matching elements efficiently.'
    allowed_refs = [AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT) for _ in range(1000)]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_many_control_flow_elements', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_with_mixed_large_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test function with a large list containing the control flow element in the middle.'
    allowed_refs = [AllowedReference(selected_element=f'element_{i}') for i in range(500)]
    allowed_refs.append(AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT))
    allowed_refs.extend([AllowedReference(selected_element=f'element_{i}') for i in range(500, 999)])
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_1', None, 'test_step_definition_allows_control_flow_references_with_mixed_large_list', 'step_definition_allows_control_flow_references', '6', codeflash_loop_index, parsed_selector)
    codeflash_output = codeflash_return_value

def test_step_definition_allows_control_flow_references_repeated_calls_consistency():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that the function returns consistent results across multiple calls.'
    allowed_refs = [AllowedReference(selected_element='element_1'), AllowedReference(selected_element=STEP_AS_SELECTED_ELEMENT), AllowedReference(selected_element='element_2')]
    selector_def = SelectorDefinition(allowed_references=allowed_refs)
    parsed_selector = ParsedSelector(definition=selector_def, step_name='test_step', value='test_value', index=None, key=None)
    results = [step_definition_allows_control_flow_references(parsed_selector) for _ in range(100)]