from __future__ import annotations

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

import pytest
from inference.core.cache.air_gapped import get_task_type_to_block_mapping


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 _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():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that an empty blocks list returns an empty mapping.'
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_empty_blocks_list', 'get_task_type_to_block_mapping', '1', codeflash_loop_index, blocks=[])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_block_single_task_type():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test mapping with a single block and single task type.'

    class SingleTaskManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'detection_block'}}}
    block = MockBlockSpecification('test_block', SingleTaskManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_single_block_single_task_type', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_block_multiple_task_types():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test mapping with a single block supporting multiple task types.'

    class MultiTaskManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection', 'segmentation', 'classification']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'multi_task_block'}}}
    block = MockBlockSpecification('test_block', MultiTaskManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_single_block_multiple_task_types', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_multiple_blocks_same_task_type():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test mapping with multiple blocks supporting the same task type.'

    class DetectionManifest1:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'detector_v1'}}}

    class DetectionManifest2:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'detector_v2'}}}
    blocks = [MockBlockSpecification('block1', DetectionManifest1), MockBlockSpecification('block2', DetectionManifest2)]
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_multiple_blocks_same_task_type', 'get_task_type_to_block_mapping', '4', codeflash_loop_index, blocks=blocks)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_with_none_task_types():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that blocks returning None from get_compatible_task_types are skipped.'

    class NoTaskManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return None

        @classmethod
        def model_json_schema(cls):
            return {}
    block = MockBlockSpecification('skipped_block', NoTaskManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_block_with_none_task_types', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_mixed_blocks_with_and_without_task_types():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "Test mapping with both blocks that have and don't have task types."

    class TaskManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'detector'}}}

    class NoTaskManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return None

        @classmethod
        def model_json_schema(cls):
            return {}
    blocks = [MockBlockSpecification('with_tasks', TaskManifest), MockBlockSpecification('without_tasks', NoTaskManifest)]
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_mixed_blocks_with_and_without_task_types', 'get_task_type_to_block_mapping', '4', codeflash_loop_index, blocks=blocks)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_type_identifier_fallback_to_block_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that block identifier is used when schema extraction fails.'

    class BrokenSchemaManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            return {}
    block = MockBlockSpecification('fallback_identifier', BrokenSchemaManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_block_type_identifier_fallback_to_block_identifier', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_block_type_identifier_from_enum():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test extracting block type from enum field in schema.'

    class EnumSchemaManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'enum': ['detector_enum', 'detector_alt']}}}
    block = MockBlockSpecification('block_id', EnumSchemaManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_block_type_identifier_from_enum', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_task_types_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test block with empty task types list.'

    class EmptyTasksManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return []

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'block_type'}}}
    block = MockBlockSpecification('block', EmptyTasksManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_empty_task_types_list', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_task_type_names_with_special_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that task type names with special characters are handled correctly.'

    class SpecialCharTaskManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection-v2', 'segmentation_3d', 'classification.binary']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'special_block'}}}
    block = MockBlockSpecification('block', SpecialCharTaskManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_task_type_names_with_special_characters', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_duplicate_task_types_in_single_block():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test block with duplicate task types in its list.'

    class DuplicateTaskManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection', 'detection']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'detector'}}}
    block = MockBlockSpecification('block', DuplicateTaskManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_duplicate_task_types_in_single_block', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_schema_with_empty_enum():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test schema with enum property but empty enum list.'

    class EmptyEnumManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'enum': []}}}
    block = MockBlockSpecification('fallback_id', EmptyEnumManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_schema_with_empty_enum', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_schema_with_neither_const_nor_enum():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test schema with type property but no const or enum.'

    class MinimalSchemaManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'type': 'string'}}}
    block = MockBlockSpecification('minimal_id', MinimalSchemaManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_schema_with_neither_const_nor_enum', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_schema_missing_type_property():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test schema without the type property.'

    class NoTypePropertyManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'other_field': {'const': 'value'}}}
    block = MockBlockSpecification('no_type_id', NoTypePropertyManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_schema_missing_type_property', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_schema_missing_properties():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test schema without properties key.'

    class NoPropertiesManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            return {'type': 'object'}
    block = MockBlockSpecification('no_props_id', NoPropertiesManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_schema_missing_properties', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_exception_in_model_json_schema():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that exceptions in model_json_schema are handled gracefully.'

    class ExceptionManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['detection']

        @classmethod
        def model_json_schema(cls):
            raise RuntimeError('Schema generation failed')
    block = MockBlockSpecification('exception_id', ExceptionManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_exception_in_model_json_schema', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_uppercase_task_type_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that uppercase task type names are preserved as-is.'

    class UppercaseManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return ['DETECTION', 'Detection', 'DeTeCTiOn']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'detector'}}}
    block = MockBlockSpecification('block', UppercaseManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_uppercase_task_type_names', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_number_of_blocks():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test mapping with a large number of blocks (1000).'
    blocks = []
    for i in range(1000):

        class LargeScaleManifest:
            task_id = i

            @classmethod
            def get_compatible_task_types(cls):
                return [f'task_type_{cls.task_id % 10}']

            @classmethod
            def model_json_schema(cls):
                return {'properties': {'type': {'const': f'block_type_{cls.task_id}'}}}
        LargeScaleManifest.task_id = i
        block = MockBlockSpecification(f'block_{i}', LargeScaleManifest)
        blocks.append(block)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_large_number_of_blocks', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=blocks)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for i in range(10):
        pass
    for i in range(10):
        pass

def test_block_with_many_task_types():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test mapping with a single block supporting many task types.'
    task_types = [f'task_type_{i}' for i in range(500)]

    class ManyTasksManifest:

        @classmethod
        def get_compatible_task_types(cls):
            return task_types

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'versatile_block'}}}
    block = MockBlockSpecification('block', ManyTasksManifest)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_block_with_many_task_types', 'get_task_type_to_block_mapping', '4', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for i in range(500):
        pass

def test_many_blocks_same_task_type():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test mapping with many blocks all supporting the same task type.'
    blocks = []
    for i in range(500):

        class SameTaskManifest:
            block_id = i

            @classmethod
            def get_compatible_task_types(cls):
                return ['universal_task']

            @classmethod
            def model_json_schema(cls):
                return {'properties': {'type': {'const': f'block_variant_{cls.block_id}'}}}
        SameTaskManifest.block_id = i
        block = MockBlockSpecification(f'block_{i}', SameTaskManifest)
        blocks.append(block)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_many_blocks_same_task_type', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=blocks)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for i in range(500):
        pass

def test_complex_large_scale_scenario():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test complex scenario with many blocks and many task type combinations.'
    blocks = []
    for block_idx in range(100):
        num_tasks = block_idx % 11 + 1
        task_types = [f'task_type_{(block_idx + j) % 50}' for j in range(num_tasks)]

        class ComplexManifest:
            manifest_task_types = task_types
            manifest_block_id = block_idx

            @classmethod
            def get_compatible_task_types(cls):
                return cls.manifest_task_types

            @classmethod
            def model_json_schema(cls):
                return {'properties': {'type': {'const': f'complex_block_{cls.manifest_block_id}'}}}
        ComplexManifest.manifest_task_types = task_types
        ComplexManifest.manifest_block_id = block_idx
        block = MockBlockSpecification(f'block_{block_idx}', ComplexManifest)
        blocks.append(block)
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_complex_large_scale_scenario', 'get_task_type_to_block_mapping', '3', codeflash_loop_index, blocks=blocks)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    total_mappings = sum((len(v) for v in result.values()))

def test_return_type_and_structure_consistency():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that the return type is always Dict[str, List[str]] regardless of scale.'
    for block_count in [0, 1, 10, 100]:
        blocks = []
        for i in range(block_count):

            class TestManifest:
                idx = i

                @classmethod
                def get_compatible_task_types(cls):
                    return [f'task_{cls.idx % 3}']

                @classmethod
                def model_json_schema(cls):
                    return {'properties': {'type': {'const': f'block_{cls.idx}'}}}
            TestManifest.idx = i
            block = MockBlockSpecification(f'b_{i}', TestManifest)
            blocks.append(block)
        codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_1', None, 'test_return_type_and_structure_consistency', 'get_task_type_to_block_mapping', '1_2', codeflash_loop_index, blocks=blocks)
        codeflash_output = codeflash_return_value
        result = codeflash_output
        for key, value in result.items():
            pass