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.workflows.core_steps.composition.workflow_caller.v1 import \
    _convert_output_descriptions_to_kinds
from inference.core.workflows.core_steps.loader import load_kinds
from inference.core.workflows.execution_engine.entities.types import (
    DICTIONARY_KIND, WILDCARD_KIND, Kind)


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

def test_single_known_kind_name_maps_to_kind_object():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    kinds = load_kinds()
    known_kind_name = kinds[0].name
    outputs_description = {'output_a': [known_kind_name]}
    _call__bound__arguments = inspect.signature(_convert_output_descriptions_to_kinds).bind(outputs_description)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_convert_output_descriptions_to_kinds, 'tests.inference.unit_tests.test__convert_output_descriptions_to_kinds__unit_test_0', None, 'test_single_known_kind_name_maps_to_kind_object', '_convert_output_descriptions_to_kinds', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    returned_kind = result['output_a'][0]

def test_multiple_known_kind_names_preserve_order_and_map_correctly():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    kinds = load_kinds()
    name1 = kinds[1].name
    name2 = kinds[2].name
    outputs_description = {'multi': [name2, name1]}
    _call__bound__arguments = inspect.signature(_convert_output_descriptions_to_kinds).bind(outputs_description)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_convert_output_descriptions_to_kinds, 'tests.inference.unit_tests.test__convert_output_descriptions_to_kinds__unit_test_0', None, 'test_multiple_known_kind_names_preserve_order_and_map_correctly', '_convert_output_descriptions_to_kinds', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    returned_list = result['multi']

def test_nested_dict_value_results_in_dictionary_kind():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    outputs_description = {'nested': {'steps': ['step.*']}}
    _call__bound__arguments = inspect.signature(_convert_output_descriptions_to_kinds).bind(outputs_description)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_convert_output_descriptions_to_kinds, 'tests.inference.unit_tests.test__convert_output_descriptions_to_kinds__unit_test_0', None, 'test_nested_dict_value_results_in_dictionary_kind', '_convert_output_descriptions_to_kinds', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_unknown_kind_names_yield_wildcard_when_no_knowns_present():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    outputs_description = {'unknown_only': ['I_DO_NOT_EXIST', 'ANOTHER_UNKNOWN']}
    _call__bound__arguments = inspect.signature(_convert_output_descriptions_to_kinds).bind(outputs_description)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_convert_output_descriptions_to_kinds, 'tests.inference.unit_tests.test__convert_output_descriptions_to_kinds__unit_test_0', None, 'test_unknown_kind_names_yield_wildcard_when_no_knowns_present', '_convert_output_descriptions_to_kinds', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_mixed_known_and_unknown_names_filters_out_unknowns_not_replacing_with_wildcard():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    kinds = load_kinds()
    known_name = kinds[0].name
    outputs_description = {'mixed': ['UNKNOWN_NAME', known_name, 'ANOTHER_UNKNOWN']}
    _call__bound__arguments = inspect.signature(_convert_output_descriptions_to_kinds).bind(outputs_description)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_convert_output_descriptions_to_kinds, 'tests.inference.unit_tests.test__convert_output_descriptions_to_kinds__unit_test_0', None, 'test_mixed_known_and_unknown_names_filters_out_unknowns_not_replacing_with_wildcard', '_convert_output_descriptions_to_kinds', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_list_for_output_becomes_wildcard():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    outputs_description = {'empty_list': []}
    _call__bound__arguments = inspect.signature(_convert_output_descriptions_to_kinds).bind(outputs_description)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_convert_output_descriptions_to_kinds, 'tests.inference.unit_tests.test__convert_output_descriptions_to_kinds__unit_test_0', None, 'test_empty_list_for_output_becomes_wildcard', '_convert_output_descriptions_to_kinds', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_outputs_description_returns_empty_dict():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(_convert_output_descriptions_to_kinds).bind({})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_convert_output_descriptions_to_kinds, 'tests.inference.unit_tests.test__convert_output_descriptions_to_kinds__unit_test_0', None, 'test_empty_outputs_description_returns_empty_dict', '_convert_output_descriptions_to_kinds', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_type_error_for_none_value_in_outputs_description():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    with pytest.raises(TypeError):
        _call__bound__arguments = inspect.signature(_convert_output_descriptions_to_kinds).bind({'bad': None})
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(_convert_output_descriptions_to_kinds, 'tests.inference.unit_tests.test__convert_output_descriptions_to_kinds__unit_test_0', None, 'test_type_error_for_none_value_in_outputs_description', '_convert_output_descriptions_to_kinds', '0_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_scale_mapping_with_many_entries():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    kinds = load_kinds()
    name_a = kinds[0].name
    name_b = kinds[1].name if len(kinds) > 1 else kinds[0].name
    n = 1000
    big_inputs = {}
    for i in range(n):
        key = f'key_{i}'
        big_inputs[key] = [name_a] if i % 2 == 0 else [name_b]
    _call__bound__arguments = inspect.signature(_convert_output_descriptions_to_kinds).bind(big_inputs)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_convert_output_descriptions_to_kinds, 'tests.inference.unit_tests.test__convert_output_descriptions_to_kinds__unit_test_0', None, 'test_large_scale_mapping_with_many_entries', '_convert_output_descriptions_to_kinds', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for i in range(n):
        key = f'key_{i}'
        expected_name = name_a if i % 2 == 0 else name_b