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 ManifestProvider:
    """A simple provider object exposing model_json_schema() method."""

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

    def model_json_schema(self):
        return self._schema

class RaisingManifestProvider:
    """A manifest provider whose model_json_schema() raises an exception."""

    def model_json_schema(self):
        raise RuntimeError('schema generation failed')

class Block:
    """A simple block-like object with manifest_class and identifier attributes."""

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

def test_const_field_is_returned_when_present():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'properties': {'type': {'const': 'my.block.Type'}}}
    manifest = ManifestProvider(schema)
    block = Block(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_0', None, 'test_const_field_is_returned_when_present', '_get_block_type_identifier', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_enum_field_is_returned_when_const_missing():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'properties': {'type': {'enum': ['enum.first', 'enum.second']}}}
    manifest = ManifestProvider(schema)
    block = Block(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_0', None, 'test_enum_field_is_returned_when_const_missing', '_get_block_type_identifier', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_empty_enum_falls_back_to_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'properties': {'type': {'enum': []}}}
    manifest = ManifestProvider(schema)
    block = Block(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_0', None, 'test_empty_enum_falls_back_to_identifier', '_get_block_type_identifier', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_missing_type_property_falls_back_to_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'properties': {'other': {'const': 'nope'}}}
    manifest = ManifestProvider(schema)
    block = Block(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_0', None, 'test_missing_type_property_falls_back_to_identifier', '_get_block_type_identifier', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    schema2 = {'not_properties': {}}
    manifest2 = ManifestProvider(schema2)
    block2 = Block(manifest2, 'other.fallback')
    _call__bound__arguments = inspect.signature(_get_block_type_identifier).bind(block2)
    _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_0', None, 'test_missing_type_property_falls_back_to_identifier', '_get_block_type_identifier', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_non_dict_type_property_triggers_fallback():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'properties': {'type': 'just-a-string'}}
    manifest = ManifestProvider(schema)
    block = Block(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_0', None, 'test_non_dict_type_property_triggers_fallback', '_get_block_type_identifier', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_model_json_schema_raising_exception_returns_identifier():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    manifest = RaisingManifestProvider()
    block = Block(manifest, 'raised.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_0', None, 'test_model_json_schema_raising_exception_returns_identifier', '_get_block_type_identifier', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_const_takes_precedence_over_enum_when_both_present():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'properties': {'type': {'const': 'const.value', 'enum': ['enum.value']}}}
    manifest = ManifestProvider(schema)
    block = Block(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_0', None, 'test_const_takes_precedence_over_enum_when_both_present', '_get_block_type_identifier', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_identifier_can_be_non_string_and_is_returned_as_is():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    schema = {'properties': {}}
    manifest = ManifestProvider(schema)
    block = Block(manifest, 12345)
    _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_0', None, 'test_identifier_can_be_non_string_and_is_returned_as_is', '_get_block_type_identifier', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_large_scale_varied_schemas_1000_blocks():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    blocks = []
    expected = []
    for i in range(1000):
        if i % 4 == 0:
            schema = {'properties': {'type': {'const': f'const.{i}'}}}
            exp = f'const.{i}'
        elif i % 4 == 1:
            schema = {'properties': {'type': {'enum': [f'enum.{i}.a', f'enum.{i}.b']}}}
            exp = f'enum.{i}.a'
        elif i % 4 == 2:
            schema = {'properties': {'type': {'enum': []}}}
            exp = f'fallback.{i}'
        else:
            schema = None
            exp = f'raised.{i}'
        if i % 4 == 3:
            manifest = RaisingManifestProvider()
            block = Block(manifest, exp)
        else:
            manifest = ManifestProvider(schema)
            block = Block(manifest, exp)
        blocks.append(block)
        expected.append(exp)
    results = [_get_block_type_identifier(b) for b in blocks]
    for res, exp in zip(results, expected):
        pass