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 import air_gapped
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}######!')
    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 BlockStub:
    """A minimal container object imitating the attributes used by the code under test."""

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

def test_single_block_const_type_returns_mapping():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    class ManifestA:

        @classmethod
        def get_compatible_task_types(cls):
            return ['task.alpha']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'block.type.alpha'}}}
    block = BlockStub(ManifestA, identifier='fallback.identifier')
    _call__bound__arguments = inspect.signature(get_task_type_to_block_mapping).bind(blocks=[block])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_0', None, 'test_single_block_const_type_returns_mapping', 'get_task_type_to_block_mapping', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    mapping = codeflash_output

def test_multiple_blocks_same_task_type_accumulate_identifiers():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    class ManifestOne:

        @classmethod
        def get_compatible_task_types(cls):
            return ['task.shared']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'enum': ['block.one', 'block.one.extra']}}}

    class ManifestTwo:

        @classmethod
        def get_compatible_task_types(cls):
            return ['task.shared']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'enum': ['block.two']}}}
    b1 = BlockStub(ManifestOne, identifier='id.one')
    b2 = BlockStub(ManifestTwo, identifier='id.two')
    _call__bound__arguments = inspect.signature(get_task_type_to_block_mapping).bind(blocks=[b1, b2])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_0', None, 'test_multiple_blocks_same_task_type_accumulate_identifiers', 'get_task_type_to_block_mapping', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    mapping = codeflash_output

def test_block_with_none_task_types_is_skipped():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    class ManifestSkip:

        @classmethod
        def get_compatible_task_types(cls):
            return None

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'const': 'should.not.be.used'}}}
    block = BlockStub(ManifestSkip, identifier='skip.id')
    _call__bound__arguments = inspect.signature(get_task_type_to_block_mapping).bind(blocks=[block])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_0', None, 'test_block_with_none_task_types_is_skipped', 'get_task_type_to_block_mapping', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    mapping = codeflash_output

def test_model_json_schema_exception_uses_identifier_fallback():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    class ManifestBadSchema:

        @classmethod
        def get_compatible_task_types(cls):
            return ['task.fallback']

        @classmethod
        def model_json_schema(cls):
            raise RuntimeError('schema parse error')
    block = BlockStub(ManifestBadSchema, identifier='fallback.id')
    _call__bound__arguments = inspect.signature(get_task_type_to_block_mapping).bind(blocks=[block])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_0', None, 'test_model_json_schema_exception_uses_identifier_fallback', 'get_task_type_to_block_mapping', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    mapping = codeflash_output

def test_enum_empty_falls_back_to_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    class ManifestEmptyEnum:

        @classmethod
        def get_compatible_task_types(cls):
            return ['task.empty_enum']

        @classmethod
        def model_json_schema(cls):
            return {'properties': {'type': {'enum': []}}}
    block = BlockStub(ManifestEmptyEnum, identifier='empty.enum.id')
    _call__bound__arguments = inspect.signature(get_task_type_to_block_mapping).bind(blocks=[block])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_0', None, 'test_enum_empty_falls_back_to_identifier', 'get_task_type_to_block_mapping', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    mapping = codeflash_output

def test_empty_blocks_list_returns_empty_mapping():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(get_task_type_to_block_mapping).bind(blocks=[])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_0', None, 'test_empty_blocks_list_returns_empty_mapping', 'get_task_type_to_block_mapping', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    mapping = codeflash_output

def test_none_blocks_calls_loader_and_handles_loader_exception(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

    def fake_loader():
        raise RuntimeError('loader failure')
    monkeypatch.setattr(air_gapped, '_load_blocks', fake_loader)
    _call__bound__arguments = inspect.signature(get_task_type_to_block_mapping).bind(blocks=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_0', None, 'test_none_blocks_calls_loader_and_handles_loader_exception', 'get_task_type_to_block_mapping', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    mapping = codeflash_output

def test_large_scale_mapping_performance_and_correctness():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    num_blocks = 1000
    num_task_types = 50
    blocks = []
    for i in range(num_blocks):
        tt_index = i % num_task_types
        task_type = f'task_{tt_index}'
        schema_const = f'block_type_{i}'
        Manifest = type(f'Manifest_{i}', (), {'get_compatible_task_types': classmethod(lambda cls, t=task_type: [t]), 'model_json_schema': classmethod(lambda cls, c=schema_const: {'properties': {'type': {'const': c}}})})
        blocks.append(BlockStub(Manifest, identifier=f'identifier_{i}'))
    _call__bound__arguments = inspect.signature(get_task_type_to_block_mapping).bind(blocks=blocks)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_task_type_to_block_mapping, 'tests.inference.unit_tests.test_get_task_type_to_block_mapping__unit_test_0', None, 'test_large_scale_mapping_performance_and_correctness', 'get_task_type_to_block_mapping', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    mapping = codeflash_output
    expected_size = num_blocks // num_task_types
    for tt_index in range(num_task_types):
        key = f'task_{tt_index}'
        values = mapping[key]
        first_expected = f'block_type_{tt_index}'
        last_expected = f'block_type_{tt_index + (expected_size - 1) * num_task_types}'