from __future__ import annotations

import gc
import os
import time

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

def _make_allowed_reference_with_selected_element(value):
    obj = Exception('allowed-ref')
    setattr(obj, 'selected_element', value)
    return obj

def test_returns_true_when_single_allowed_reference_matches():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    definition_holder = Exception('definition-holder')
    definition_holder.allowed_references = [_make_allowed_reference_with_selected_element(STEP_AS_SELECTED_ELEMENT)]
    parsed = ParsedSelector(definition=definition_holder, step_name='s', value='v', 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_0', None, 'test_returns_true_when_single_allowed_reference_matches', 'step_definition_allows_control_flow_references', '3', codeflash_loop_index, parsed)
    codeflash_output = codeflash_return_value

def test_returns_false_when_no_allowed_references_match():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    definition_holder = Exception('definition-holder')
    definition_holder.allowed_references = [_make_allowed_reference_with_selected_element('not-step'), _make_allowed_reference_with_selected_element('another')]
    parsed = ParsedSelector(definition=definition_holder, step_name='s', value='v', index=0, 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_0', None, 'test_returns_false_when_no_allowed_references_match', 'step_definition_allows_control_flow_references', '3', codeflash_loop_index, parsed)
    codeflash_output = codeflash_return_value

def test_returns_false_with_empty_allowed_references_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    definition_holder = Exception('definition-holder')
    definition_holder.allowed_references = []
    parsed = ParsedSelector(definition=definition_holder, step_name='s', value='v', index=1, key='k')
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_0', None, 'test_returns_false_with_empty_allowed_references_list', 'step_definition_allows_control_flow_references', '3', codeflash_loop_index, parsed)
    codeflash_output = codeflash_return_value

def test_raises_attribute_error_when_an_allowed_reference_lacks_selected_element():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    definition_holder = Exception('definition-holder')
    missing_attr_obj = Exception('no-selected-element')
    definition_holder.allowed_references = [missing_attr_obj, _make_allowed_reference_with_selected_element(STEP_AS_SELECTED_ELEMENT)]
    parsed = ParsedSelector(definition=definition_holder, step_name='s', value='v', index=None, key=None)
    with pytest.raises(AttributeError):
        codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_0', None, 'test_raises_attribute_error_when_an_allowed_reference_lacks_selected_element', 'step_definition_allows_control_flow_references', '4_0', codeflash_loop_index, parsed)

def test_raises_type_error_when_allowed_references_is_none():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    definition_holder = Exception('definition-holder')
    definition_holder.allowed_references = None
    parsed = ParsedSelector(definition=definition_holder, step_name='s', value='v', index=None, key=None)
    with pytest.raises(TypeError):
        codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_0', None, 'test_raises_type_error_when_allowed_references_is_none', 'step_definition_allows_control_flow_references', '3_0', codeflash_loop_index, parsed)

def test_large_scale_one_match_in_middle_of_1000():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    definition_holder = Exception('definition-holder')
    refs = []
    for i in range(499):
        refs.append(_make_allowed_reference_with_selected_element(f'not-step-{i}'))
    refs.append(_make_allowed_reference_with_selected_element(STEP_AS_SELECTED_ELEMENT))
    for i in range(500, 1000):
        refs.append(_make_allowed_reference_with_selected_element(f'not-step-{i}'))
    definition_holder.allowed_references = refs
    parsed = ParsedSelector(definition=definition_holder, step_name='s', value='v', index=42, key='k')
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_0', None, 'test_large_scale_one_match_in_middle_of_1000', 'step_definition_allows_control_flow_references', '7', codeflash_loop_index, parsed)
    codeflash_output = codeflash_return_value

def test_large_scale_no_match_1000_items_returns_false():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    definition_holder = Exception('definition-holder')
    refs = [_make_allowed_reference_with_selected_element(f'not-step-{i}') for i in range(1000)]
    definition_holder.allowed_references = refs
    parsed = ParsedSelector(definition=definition_holder, step_name='s', value='v', index=0, 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_0', None, 'test_large_scale_no_match_1000_items_returns_false', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed)
    codeflash_output = codeflash_return_value

def test_multiple_matches_still_returns_true():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    definition_holder = Exception('definition-holder')
    refs = [_make_allowed_reference_with_selected_element(STEP_AS_SELECTED_ELEMENT), _make_allowed_reference_with_selected_element(STEP_AS_SELECTED_ELEMENT), _make_allowed_reference_with_selected_element('not-step')]
    definition_holder.allowed_references = refs
    parsed = ParsedSelector(definition=definition_holder, step_name='multi', value='v', index=None, key='kk')
    codeflash_return_value = codeflash_wrap(step_definition_allows_control_flow_references, 'tests.inference.unit_tests.test_step_definition_allows_control_flow_references__unit_test_0', None, 'test_multiple_matches_still_returns_true', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed)
    codeflash_output = codeflash_return_value

def test_non_string_selected_element_that_matches_constant_by_identity_or_equality():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    definition_holder = Exception('definition-holder')
    refs = [_make_allowed_reference_with_selected_element(STEP_AS_SELECTED_ELEMENT)]
    definition_holder.allowed_references = refs
    parsed = ParsedSelector(definition=definition_holder, step_name='s', value='v', 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_0', None, 'test_non_string_selected_element_that_matches_constant_by_identity_or_equality', 'step_definition_allows_control_flow_references', '4', codeflash_loop_index, parsed)
    codeflash_output = codeflash_return_value