from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path

import dill as pickle
import pytest
from inference.core.cache.air_gapped import _get_block_type_identifier


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}######!')
    iteration = os.environ['CODEFLASH_TEST_ITERATION']
    with Path('{codeflash_run_tmp_dir_client_side}', f'test_return_values_{iteration}.bin').open('ab') as f:
        pickled_values = pickle.dumps((args, kwargs, exception)) if exception else pickle.dumps((args, kwargs, return_value))
        _test_name = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_line_id}".encode('ascii')
        f.write(len(_test_name).to_bytes(4, byteorder='big'))
        f.write(_test_name)
        f.write(codeflash_duration.to_bytes(8, byteorder='big'))
        f.write(len(pickled_values).to_bytes(4, byteorder='big'))
        f.write(pickled_values)
        f.write(codeflash_loop_index.to_bytes(8, byteorder='big'))
        f.write(len(invocation_id).to_bytes(4, byteorder='big'))
        f.write(invocation_id.encode('ascii'))
    if exception:
        raise exception
    return return_value

class MockManifestClass:
    """Mock manifest class with configurable schema."""

    def __init__(self, schema):
        self._schema = schema

    def model_json_schema(self):
        return self._schema

class MockBlock:
    """Real mock block object with manifest_class and identifier."""

    def __init__(self, manifest_class, identifier):
        self.manifest_class = manifest_class
        self.identifier = identifier

def test_block_type_with_const_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test extraction of type identifier from schema with const value.'
    schema = {'properties': {'type': {'const': 'my_block_type'}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback_identifier')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_type_with_const_identifier', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_type_with_enum_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test extraction of type identifier from schema with enum values.'
    schema = {'properties': {'type': {'enum': ['enum_block_type', 'other_type']}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback_identifier')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_type_with_enum_identifier', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_type_fallback_to_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test fallback to block identifier when no type in schema.'
    schema = {'properties': {}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'my_fallback_id')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_type_fallback_to_identifier', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_type_const_takes_precedence_over_enum():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that const takes precedence when both const and enum exist.'
    schema = {'properties': {'type': {'const': 'const_type', 'enum': ['enum_type_1', 'enum_type_2']}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback_identifier')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_type_const_takes_precedence_over_enum', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_exception_in_manifest():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test fallback when manifest_class raises exception.'

    class BrokenManifest:

        def model_json_schema(self):
            raise RuntimeError('Manifest error')
    manifest = BrokenManifest()
    block = MockBlock(manifest, 'safe_fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_exception_in_manifest', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_empty_enum_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test behavior when enum list is empty.'
    schema = {'properties': {'type': {'enum': []}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback_for_empty_enum')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_empty_enum_list', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_missing_properties_key():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test when schema has no properties key.'
    schema = {'title': 'SomeBlock'}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback_no_properties')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_missing_properties_key', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_missing_type_key():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test when properties exists but has no type key.'
    schema = {'properties': {'name': {'type': 'string'}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback_no_type')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_missing_type_key', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_empty_schema():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with completely empty schema.'
    schema = {}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback_empty_schema')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_empty_schema', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_single_character_const_type():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with single character const type value.'
    schema = {'properties': {'type': {'const': 'A'}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_single_character_const_type', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_special_characters_in_type():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with special characters in type identifier.'
    schema = {'properties': {'type': {'const': 'block-type_v2.0'}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_special_characters_in_type', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_numeric_string_type():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with numeric string as type identifier.'
    schema = {'properties': {'type': {'const': '12345'}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_numeric_string_type', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_identifier_with_special_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test fallback identifier with special characters.'
    schema = {'properties': {}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'special-id_v1.0')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_identifier_with_special_characters', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_none_manifest():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test when manifest_class is None but has identifier.'

    class BlockWithNoneManifest:

        def __init__(self):
            self.manifest_class = None
            self.identifier = 'fallback_when_none'
    block = BlockWithNoneManifest()
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_none_manifest', '_get_block_type_identifier', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_missing_identifier_attribute():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test robustness when identifier attribute might be missing.'

    class PartialBlock:

        def __init__(self):
            self.manifest_class = MockManifestClass({'properties': {}})
    block = PartialBlock()
    try:
        _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_missing_identifier_attribute', '_get_block_type_identifier', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    except AttributeError:
        pass

def test_block_with_enum_single_value():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test enum with exactly one value.'
    schema = {'properties': {'type': {'enum': ['only_one_type']}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_enum_single_value', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_very_long_type_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with very long type identifier string.'
    long_type = 'a' * 1000
    schema = {'properties': {'type': {'const': long_type}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_very_long_type_identifier', '_get_block_type_identifier', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_unicode_type_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with unicode characters in type identifier.'
    schema = {'properties': {'type': {'const': 'блок_тип'}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_unicode_type_identifier', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_whitespace_in_type():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test type identifier with whitespace.'
    schema = {'properties': {'type': {'const': 'block type with spaces'}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_whitespace_in_type', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_large_enum_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test performance with large enum list.'
    large_enum = [f'type_{i}' for i in range(1000)]
    schema = {'properties': {'type': {'enum': large_enum}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_large_enum_list', '_get_block_type_identifier', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_large_schema():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test performance with large schema containing many properties.'
    large_properties = {f'prop_{i}': {'type': 'string'} for i in range(1000)}
    large_properties['type'] = {'const': 'large_schema_type'}
    schema = {'properties': large_properties}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_large_schema', '_get_block_type_identifier', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_repeated_calls_consistency():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that repeated calls return consistent results.'
    schema = {'properties': {'type': {'const': 'consistent_type'}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    results = [_get_block_type_identifier(block) for _ in range(1000)]

def test_block_with_deeply_nested_schema():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with deeply nested schema structure.'
    schema = {'properties': {'type': {'const': 'nested_schema_type'}, 'nested': {'type': 'object', 'properties': {'deep': {'type': 'object'}}}}}
    manifest = MockManifestClass(schema)
    block = MockBlock(manifest, 'fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_get_block_type_identifier, 'tests.inference.unit_tests.test__get_block_type_identifier__unit_test_1', None, 'test_block_with_deeply_nested_schema', '_get_block_type_identifier', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_bulk_block_processing():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test processing many blocks in succession.'
    blocks = []
    for i in range(100):
        if i % 3 == 0:
            schema = {'properties': {'type': {'const': f'const_type_{i}'}}}
        elif i % 3 == 1:
            schema = {'properties': {'type': {'enum': [f'enum_type_{i}', f'other_{i}']}}}
        else:
            schema = {'properties': {}}
        manifest = MockManifestClass(schema)
        block = MockBlock(manifest, f'fallback_{i}')
        blocks.append(block)
    results = [_get_block_type_identifier(block) for block in blocks]