from __future__ import annotations

import gc
import os
import time
from typing import Any, Dict, List, Optional
from unittest.mock import MagicMock, Mock, patch

import pytest
from inference.core.cache.air_gapped import (get_cached_foundation_models,
                                             has_cached_model_variant,
                                             is_model_cached)


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 _load_blocks() -> list:
    """Load workflow blocks, isolating the heavy import for testability."""
    return []

def _get_block_type_identifier(block) -> str:
    """Extract the canonical ``type`` identifier from a block specification."""
    try:
        schema = block.manifest_class.model_json_schema()
        type_prop = schema.get('properties', {}).get('type', {})
        if 'const' in type_prop:
            return type_prop['const']
        if 'enum' in type_prop and type_prop['enum']:
            return type_prop['enum'][0]
    except Exception:
        pass
    return block.identifier

def test_empty_blocks_list_returns_empty_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that an empty blocks list returns an empty results list.'
    blocks = []
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_empty_blocks_list_returns_empty_list', 'get_cached_foundation_models', '2', codeflash_loop_index, blocks=blocks)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_blocks_with_no_model_variants_returns_empty_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that blocks returning None for model_variants are skipped.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = None
    blocks = [block]
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_blocks_with_no_model_variants_returns_empty_list', 'get_cached_foundation_models', '5', codeflash_loop_index, blocks=blocks)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_block_with_cached_model_variant():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that a single block with a cached model variant is included.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['model_1']
    block.manifest_class.model_json_schema.return_value = {'name': 'TestModel'}
    block.identifier = 'block_1'
    blocks = [block]
    with patch('__main__.is_model_cached', return_value=True):
        with patch('__main__.has_cached_model_variant', return_value=True):
            codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_single_block_with_cached_model_variant', 'get_cached_foundation_models', '7_0_0', codeflash_loop_index, blocks=blocks)
            codeflash_output = codeflash_return_value
            result = codeflash_output

def test_block_with_uncached_model_variant():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that a block without cached model variants is excluded.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['model_1']
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=False):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_block_with_uncached_model_variant', 'get_cached_foundation_models', '5_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_multiple_blocks_mixed_cached_and_uncached():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test filtering with multiple blocks, some cached and some not.'
    block1 = Mock()
    block1.manifest_class = Mock()
    block1.manifest_class.get_supported_model_variants.return_value = ['model_1']
    block1.manifest_class.model_json_schema.return_value = {'name': 'Model1'}
    block1.identifier = 'block_1'
    block2 = Mock()
    block2.manifest_class = Mock()
    block2.manifest_class.get_supported_model_variants.return_value = ['model_2']
    block2.identifier = 'block_2'
    block3 = Mock()
    block3.manifest_class = Mock()
    block3.manifest_class.get_supported_model_variants.return_value = ['model_3']
    block3.manifest_class.model_json_schema.return_value = {'name': 'Model3'}
    block3.identifier = 'block_3'
    blocks = [block1, block2, block3]

    def mock_has_cached(variants):
        if variants and variants[0] in ['model_1', 'model_3']:
            return True
        return False
    with patch('__main__.has_cached_model_variant', side_effect=mock_has_cached):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_multiple_blocks_mixed_cached_and_uncached', 'get_cached_foundation_models', '17_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_block_name_defaults_to_model_id_when_schema_fails():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that block_name falls back to model_id when schema retrieval fails.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['model_1']
    block.manifest_class.model_json_schema.side_effect = Exception('Schema error')
    block.identifier = 'block_1'
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_block_name_defaults_to_model_id_when_schema_fails', 'get_cached_foundation_models', '7_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_none_blocks_parameter_returns_empty_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that None blocks parameter triggers _load_blocks and returns empty list on failure.'
    with patch('__main__._load_blocks', side_effect=Exception('Load failed')):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_none_blocks_parameter_returns_empty_list', 'get_cached_foundation_models', '1_0', codeflash_loop_index, blocks=None)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_block_with_empty_model_variants_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that a block with empty model_variants list is handled.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = []
    blocks = [block]
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_block_with_empty_model_variants_list', 'get_cached_foundation_models', '5', codeflash_loop_index, blocks=blocks)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_model_id_extraction_uses_first_variant():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that model_id is extracted as the first element of variants.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['first_model', 'second_model', 'third_model']
    block.manifest_class.model_json_schema.return_value = {}
    block.identifier = 'block_1'
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_model_id_extraction_uses_first_variant', 'get_cached_foundation_models', '7_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_schema_name_field_overrides_model_id():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that schema name field is used when available.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['model_1']
    block.manifest_class.model_json_schema.return_value = {'name': 'CustomName'}
    block.identifier = 'block_1'
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_schema_name_field_overrides_model_id', 'get_cached_foundation_models', '7_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_schema_without_name_field():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that model_id is used when schema lacks name field.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['model_1']
    block.manifest_class.model_json_schema.return_value = {'other_field': 'value'}
    block.identifier = 'block_1'
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_schema_without_name_field', 'get_cached_foundation_models', '7_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_block_type_identifier_from_schema_const():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test extraction of block_type from schema const field.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['model_1']
    block.manifest_class.model_json_schema.return_value = {'name': 'TestModel', 'properties': {'type': {'const': 'const_type_value'}}}
    block.identifier = 'default_block_id'
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_block_type_identifier_from_schema_const', 'get_cached_foundation_models', '7_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_block_type_identifier_from_schema_enum():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test extraction of block_type from schema enum field.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['model_1']
    block.manifest_class.model_json_schema.return_value = {'name': 'TestModel', 'properties': {'type': {'enum': ['enum_type_1', 'enum_type_2']}}}
    block.identifier = 'default_block_id'
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_block_type_identifier_from_schema_enum', 'get_cached_foundation_models', '7_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_block_type_identifier_defaults_to_block_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that block_type defaults to block.identifier when schema parsing fails.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['model_1']
    block.manifest_class.model_json_schema.return_value = {'name': 'TestModel'}
    block.identifier = 'fallback_block_id'
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_block_type_identifier_defaults_to_block_identifier', 'get_cached_foundation_models', '7_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_result_has_all_required_fields():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that each result entry has all required fields.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['model_1']
    block.manifest_class.model_json_schema.return_value = {'name': 'TestModel'}
    block.identifier = 'block_1'
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_result_has_all_required_fields', 'get_cached_foundation_models', '7_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    required_fields = ['model_id', 'name', 'task_type', 'model_architecture', 'is_foundation', 'block_type']
    for field in required_fields:
        pass

def test_is_foundation_always_true():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that is_foundation is always True in results.'
    block = Mock()
    block.manifest_class = Mock()
    block.manifest_class.get_supported_model_variants.return_value = ['model_1']
    block.manifest_class.model_json_schema.return_value = {}
    block.identifier = 'block_1'
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_is_foundation_always_true', 'get_cached_foundation_models', '7_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_large_number_of_blocks_no_cached():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test performance with 100 blocks, none cached.'
    blocks = []
    for i in range(100):
        block = Mock()
        block.manifest_class = Mock()
        block.manifest_class.get_supported_model_variants.return_value = [f'model_{i}']
        blocks.append(block)
    with patch('__main__.has_cached_model_variant', return_value=False):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_large_number_of_blocks_no_cached', 'get_cached_foundation_models', '3_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_large_number_of_blocks_all_cached():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test performance with 100 blocks, all cached.'
    blocks = []
    for i in range(100):
        block = Mock()
        block.manifest_class = Mock()
        block.manifest_class.get_supported_model_variants.return_value = [f'model_{i}']
        block.manifest_class.model_json_schema.return_value = {'name': f'Model_{i}'}
        block.identifier = f'block_{i}'
        blocks.append(block)
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_large_number_of_blocks_all_cached', 'get_cached_foundation_models', '3_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    for i, entry in enumerate(result):
        pass

def test_large_number_of_model_variants_per_block():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test handling of blocks with many model variants (only first used).'
    block = Mock()
    block.manifest_class = Mock()
    variants = [f'model_{i}' for i in range(50)]
    block.manifest_class.get_supported_model_variants.return_value = variants
    block.manifest_class.model_json_schema.return_value = {}
    block.identifier = 'block_1'
    blocks = [block]
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_large_number_of_model_variants_per_block', 'get_cached_foundation_models', '8_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_500_blocks_with_alternating_cached_status():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test performance with 500 blocks, alternating cached/uncached.'
    blocks = []
    for i in range(500):
        block = Mock()
        block.manifest_class = Mock()
        block.manifest_class.get_supported_model_variants.return_value = [f'model_{i}']
        block.manifest_class.model_json_schema.return_value = {'name': f'Name_{i}'}
        block.identifier = f'id_{i}'
        blocks.append(block)

    def mock_has_cached_alternating(variants):
        if variants and int(variants[0].split('_')[1]) % 2 == 0:
            return True
        return False
    with patch('__main__.has_cached_model_variant', side_effect=mock_has_cached_alternating):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_500_blocks_with_alternating_cached_status', 'get_cached_foundation_models', '4_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_1000_blocks_with_none_variants():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test performance with 1000 blocks that return None for variants.'
    blocks = []
    for i in range(1000):
        block = Mock()
        block.manifest_class = Mock()
        block.manifest_class.get_supported_model_variants.return_value = None
        blocks.append(block)
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_1000_blocks_with_none_variants', 'get_cached_foundation_models', '3', codeflash_loop_index, blocks=blocks)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_blocks_with_schema_exception_handling():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test performance with blocks that raise exceptions during schema retrieval.'
    blocks = []
    for i in range(50):
        block = Mock()
        block.manifest_class = Mock()
        block.manifest_class.get_supported_model_variants.return_value = [f'model_{i}']
        block.identifier = f'block_{i}'
        if i % 2 == 0:
            block.manifest_class.model_json_schema.side_effect = Exception('Schema error')
        else:
            block.manifest_class.model_json_schema.return_value = {'name': f'Name_{i}'}
        blocks.append(block)
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_blocks_with_schema_exception_handling', 'get_cached_foundation_models', '3_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    for i, entry in enumerate(result):
        if i % 2 == 0:
            pass
        else:
            pass

def test_duplicate_model_ids_in_results():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test handling when multiple blocks have same model_id.'
    model_id = 'shared_model'
    blocks = []
    for i in range(5):
        block = Mock()
        block.manifest_class = Mock()
        block.manifest_class.get_supported_model_variants.return_value = [model_id]
        block.manifest_class.model_json_schema.return_value = {'name': f'Variant_{i}'}
        block.identifier = f'block_{i}'
        blocks.append(block)
    with patch('__main__.has_cached_model_variant', return_value=True):
        codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_1', None, 'test_duplicate_model_ids_in_results', 'get_cached_foundation_models', '4_0', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    for entry in result:
        pass