from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence, Set

import dill as pickle
import pytest
from inference.core.workflows.execution_engine.v1.executor.execution_data_manager.dynamic_batches_manager import \
    DynamicBatchIndex
from inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler import \
    intersect_masks_per_dimension


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 make_db_index(seq: Iterable[int]) -> Any:
    seq_list = list(seq)
    try:
        instance = DynamicBatchIndex(seq_list)
        return instance
    except Exception:
        return tuple(seq_list)

def test_basic_single_dimension_intersection():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    m1 = {make_db_index([1]), make_db_index([2])}
    m2 = {make_db_index([2]), make_db_index([3])}
    m3 = {make_db_index([2])}
    batch_masks: List[Set[DynamicBatchIndex]] = [m1, m2, m3]
    dims = {1}
    _call__bound__arguments = inspect.signature(intersect_masks_per_dimension).bind(batch_masks, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(intersect_masks_per_dimension, 'tests.inference.unit_tests.test_intersect_masks_per_dimension__unit_test_0', None, 'test_basic_single_dimension_intersection', 'intersect_masks_per_dimension', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = {make_db_index([2])}

def test_basic_multiple_dimensions_with_ignored_masks():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask1 = {make_db_index([1]), make_db_index([5, 5]), make_db_index([9, 9])}
    mask2 = {make_db_index([1]), make_db_index([5, 5]), make_db_index([2, 2])}
    mask3 = {make_db_index([1]), make_db_index([3, 3])}
    batch_masks: List[Set[DynamicBatchIndex]] = [mask1, mask2, mask3]
    dims = {1, 2}
    _call__bound__arguments = inspect.signature(intersect_masks_per_dimension).bind(batch_masks, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(intersect_masks_per_dimension, 'tests.inference.unit_tests.test_intersect_masks_per_dimension__unit_test_0', None, 'test_basic_multiple_dimensions_with_ignored_masks', 'intersect_masks_per_dimension', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_inputs_and_empty_dimensions_behavior():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    batch_masks: List[Set[DynamicBatchIndex]] = []
    dims = {1, 2, 3}
    _call__bound__arguments = inspect.signature(intersect_masks_per_dimension).bind(batch_masks, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(intersect_masks_per_dimension, 'tests.inference.unit_tests.test_intersect_masks_per_dimension__unit_test_0', None, 'test_empty_inputs_and_empty_dimensions_behavior', 'intersect_masks_per_dimension', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for d in dims:
        pass
    masks = [{make_db_index([1])}, {make_db_index([1, 1])}]
    _call__bound__arguments = inspect.signature(intersect_masks_per_dimension).bind(masks, set())
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(intersect_masks_per_dimension, 'tests.inference.unit_tests.test_intersect_masks_per_dimension__unit_test_0', None, 'test_empty_inputs_and_empty_dimensions_behavior', 'intersect_masks_per_dimension', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result2 = codeflash_output

def test_masks_with_some_masks_lacking_dimension_entries_are_ignored():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask1 = {make_db_index([1, 2]), make_db_index([2, 3])}
    mask2 = {make_db_index([1, 2]), make_db_index([4, 5])}
    mask3 = {make_db_index([99])}
    batch_masks: List[Set[DynamicBatchIndex]] = [mask1, mask2, mask3]
    dims = {2}
    _call__bound__arguments = inspect.signature(intersect_masks_per_dimension).bind(batch_masks, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(intersect_masks_per_dimension, 'tests.inference.unit_tests.test_intersect_masks_per_dimension__unit_test_0', None, 'test_masks_with_some_masks_lacking_dimension_entries_are_ignored', 'intersect_masks_per_dimension', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_type_error_when_elements_have_no_len():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask_with_bad_element = {42}
    with pytest.raises(TypeError):
        _call__bound__arguments = inspect.signature(intersect_masks_per_dimension).bind([mask_with_bad_element], {1})
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(intersect_masks_per_dimension, 'tests.inference.unit_tests.test_intersect_masks_per_dimension__unit_test_0', None, 'test_type_error_when_elements_have_no_len', 'intersect_masks_per_dimension', '1_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_scale_intersection_performance_and_correctness():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n_masks = 1000
    common_size = 30
    dims = {3}
    common_indices = [make_db_index([i, i + 1, i + 2]) for i in range(common_size)]
    batch_masks: List[Set[DynamicBatchIndex]] = []
    for m in range(n_masks):
        unique1 = make_db_index([1000000 + m, 1000000 + m + 1, 1000000 + m + 2])
        unique2 = make_db_index([2000000 + m, 2000000 + m + 1, 2000000 + m + 2])
        mask_set = set(common_indices)
        mask_set.add(unique1)
        mask_set.add(unique2)
        batch_masks.append(mask_set)
    _call__bound__arguments = inspect.signature(intersect_masks_per_dimension).bind(batch_masks, dims)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(intersect_masks_per_dimension, 'tests.inference.unit_tests.test_intersect_masks_per_dimension__unit_test_0', None, 'test_large_scale_intersection_performance_and_correctness', 'intersect_masks_per_dimension', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected_set = set(common_indices)