from __future__ import annotations

import gc
import os
import time

import pytest
from inference.core.workflows.core_steps.common.query_language.introspection.core import \
    _ref_to_def_name


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 test_basic_ref_extraction_default_key():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'$ref': '#/$defs/Foo'}
    codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_basic_ref_extraction_default_key', '_ref_to_def_name', '1', codeflash_loop_index, schema)
    codeflash_output = codeflash_return_value

def test_ref_with_nested_path_segments():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'$ref': '#/components/schemas/MyType'}
    codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_ref_with_nested_path_segments', '_ref_to_def_name', '1', codeflash_loop_index, schema)
    codeflash_output = codeflash_return_value

def test_custom_ref_key():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'customRef': '#/$defs/Bar'}
    codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_custom_ref_key', '_ref_to_def_name', '1', codeflash_loop_index, schema, ref_key='customRef')
    codeflash_output = codeflash_return_value

def test_missing_ref_key_returns_none():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {}
    codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_missing_ref_key_returns_none', '_ref_to_def_name', '1', codeflash_loop_index, schema)
    codeflash_output = codeflash_return_value

def test_ref_value_none_returns_none():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'$ref': None}
    codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_ref_value_none_returns_none', '_ref_to_def_name', '1', codeflash_loop_index, schema)
    codeflash_output = codeflash_return_value

def test_empty_string_ref_returns_empty_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'$ref': ''}
    codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_empty_string_ref_returns_empty_string', '_ref_to_def_name', '1', codeflash_loop_index, schema)
    codeflash_output = codeflash_return_value

def test_ref_without_slashes_returns_whole_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'$ref': 'PlainName'}
    codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_ref_without_slashes_returns_whole_string', '_ref_to_def_name', '1', codeflash_loop_index, schema)
    codeflash_output = codeflash_return_value

def test_trailing_slash_returns_empty_segment():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'$ref': '#/$defs/'}
    codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_trailing_slash_returns_empty_segment', '_ref_to_def_name', '1', codeflash_loop_index, schema)
    codeflash_output = codeflash_return_value

def test_non_string_ref_raises_attribute_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'$ref': 123}
    with pytest.raises(AttributeError):
        codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_non_string_ref_raises_attribute_error', '_ref_to_def_name', '1_0', codeflash_loop_index, schema)

def test_input_dictionary_not_mutated():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'$ref': '#/$defs/Immutable'}
    original_copy = schema.copy()
    codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_input_dictionary_not_mutated', '_ref_to_def_name', '2', codeflash_loop_index, schema)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_many_schemas():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    count = 1000
    schemas = [{'$ref': f'#/$defs/Name{i}'} for i in range(count)]
    results = [_ref_to_def_name(s) for s in schemas]
    expected = [f'Name{i}' for i in range(count)]

def test_large_scale_repeated_calls_consistency():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'$ref': '#/$defs/Repeatable'}
    outputs = []
    for _ in range(1000):
        codeflash_return_value = codeflash_wrap(_ref_to_def_name, 'tests.inference.unit_tests.test__ref_to_def_name__unit_test_0', None, 'test_large_scale_repeated_calls_consistency', '_ref_to_def_name', '2_0', codeflash_loop_index, schema)