from __future__ import annotations

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

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}:{codeflash_duration}######!')
    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}
    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, batch_masks, dims)
    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}
    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, batch_masks, dims)
    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}
    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, batch_masks, dims)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for d in dims:
        pass
    masks = [{make_db_index([1])}, {make_db_index([1, 1])}]
    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, masks, set())
    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}
    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, batch_masks, dims)
    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):
        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, [mask_with_bad_element], {1})

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)
    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, batch_masks, dims)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected_set = set(common_indices)