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.execution_engine.v1.executor.execution_data_manager.step_input_assembler import \
    filter_to_valid_prefix_chains


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_basic_two_level_chain_keeps_matching_parent_child():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    per_dim_sets = {0: {(1,)}, 1: {(1, 2), (1,), (3, 4)}}
    dims = {0, 1}
    _call__bound__arguments = inspect.signature(filter_to_valid_prefix_chains).bind(per_dim_sets, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(filter_to_valid_prefix_chains, 'tests.inference.unit_tests.test_filter_to_valid_prefix_chains__unit_test_0', None, 'test_basic_two_level_chain_keeps_matching_parent_child', 'filter_to_valid_prefix_chains', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_three_level_chain_with_extras_filters_non_chains():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    per_dim_sets = {0: {(10,), (20,), (30,)}, 1: {(10, 1), (20, 2), (99, 9)}, 2: {(10, 1, 7), (20, 2, 8), (100, 0, 0)}}
    dims = {2, 0, 1}
    _call__bound__arguments = inspect.signature(filter_to_valid_prefix_chains).bind(per_dim_sets, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(filter_to_valid_prefix_chains, 'tests.inference.unit_tests.test_filter_to_valid_prefix_chains__unit_test_0', None, 'test_three_level_chain_with_extras_filters_non_chains', 'filter_to_valid_prefix_chains', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_dimension_returns_copy_of_input_sets_and_handles_empty():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    per_dim_sets = {0: {(1,), (2,)}}
    dims = {0}
    _call__bound__arguments = inspect.signature(filter_to_valid_prefix_chains).bind(per_dim_sets, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(filter_to_valid_prefix_chains, 'tests.inference.unit_tests.test_filter_to_valid_prefix_chains__unit_test_0', None, 'test_single_dimension_returns_copy_of_input_sets_and_handles_empty', 'filter_to_valid_prefix_chains', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    _call__bound__arguments = inspect.signature(filter_to_valid_prefix_chains).bind({}, {5})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(filter_to_valid_prefix_chains, 'tests.inference.unit_tests.test_filter_to_valid_prefix_chains__unit_test_0', None, 'test_single_dimension_returns_copy_of_input_sets_and_handles_empty', 'filter_to_valid_prefix_chains', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result_missing = codeflash_output

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

def test_orphan_children_are_removed_when_parent_not_valid():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    per_dim_sets = {0: {(42,)}, 1: {(99, 1), (100, 2)}}
    dims = {0, 1}
    _call__bound__arguments = inspect.signature(filter_to_valid_prefix_chains).bind(per_dim_sets, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(filter_to_valid_prefix_chains, 'tests.inference.unit_tests.test_filter_to_valid_prefix_chains__unit_test_0', None, 'test_orphan_children_are_removed_when_parent_not_valid', 'filter_to_valid_prefix_chains', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_tuple_indices_are_handled_gracefully():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    per_dim_sets = {0: {()}, 1: {('',), (1,)}}
    dims = {0, 1}
    _call__bound__arguments = inspect.signature(filter_to_valid_prefix_chains).bind(per_dim_sets, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(filter_to_valid_prefix_chains, 'tests.inference.unit_tests.test_filter_to_valid_prefix_chains__unit_test_0', None, 'test_empty_tuple_indices_are_handled_gracefully', 'filter_to_valid_prefix_chains', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_correctness_and_performance():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    N = 1000
    top = {(i,) for i in range(N)}
    middle = {(i, 0) for i in range(N)}
    bottom = {(i, 0, 0) for i in range(0, N, 2)}
    per_dim_sets = {0: top, 1: middle, 2: bottom}
    dims = {0, 1, 2}
    _call__bound__arguments = inspect.signature(filter_to_valid_prefix_chains).bind(per_dim_sets, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(filter_to_valid_prefix_chains, 'tests.inference.unit_tests.test_filter_to_valid_prefix_chains__unit_test_0', None, 'test_large_scale_correctness_and_performance', 'filter_to_valid_prefix_chains', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected_top = {(i,) for i in range(0, N, 2)}
    expected_middle = {(i, 0) for i in range(0, N, 2)}
    expected_bottom = {(i, 0, 0) for i in range(0, N, 2)}