from __future__ import annotations

import gc
import os
import time
import types
from typing import Set, Tuple

import pytest
from inference.core.workflows.execution_engine.v1.executor.execution_data_manager.branching_manager import \
    BranchingManager
from inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler import \
    collect_inputs_dimensionalities as real_collect_inputs_dimensionalities
from inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler import \
    construct_mask_for_all_inputs_dimensionalities


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 test_unregistered_branch_returns_empty_sets_and_true(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr('inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler.collect_inputs_dimensionalities', lambda step_node: {'input_a': 2})
    step_node = types.SimpleNamespace(execution_branches_impacting_inputs=['branch_missing'])
    bm = BranchingManager(masks={})
    codeflash_return_value = codeflash_wrap(construct_mask_for_all_inputs_dimensionalities, 'tests.inference.unit_tests.test_construct_mask_for_all_inputs_dimensionalities__unit_test_0', None, 'test_unregistered_branch_returns_empty_sets_and_true', 'construct_mask_for_all_inputs_dimensionalities', '3', codeflash_loop_index, step_node=step_node, branching_manager=bm)
    masks_by_dim, scalar_false = codeflash_return_value

def test_non_batch_true_returns_none_for_each_dimension(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr('inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler.collect_inputs_dimensionalities', lambda step_node: {'p1': 1, 'p2': 2})
    step_node = types.SimpleNamespace(execution_branches_impacting_inputs=['b1', 'b2'])
    bm = BranchingManager(masks={'b1': True, 'b2': True})
    codeflash_return_value = codeflash_wrap(construct_mask_for_all_inputs_dimensionalities, 'tests.inference.unit_tests.test_construct_mask_for_all_inputs_dimensionalities__unit_test_0', None, 'test_non_batch_true_returns_none_for_each_dimension', 'construct_mask_for_all_inputs_dimensionalities', '3', codeflash_loop_index, step_node=step_node, branching_manager=bm)
    masks_by_dim, scalar_false = codeflash_return_value

def test_batch_masks_intersection_basic(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr('inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler.collect_inputs_dimensionalities', lambda step_node: {'a': 1, 'b': 2})
    mask_common = {(1,), (1, 1)}
    step_node = types.SimpleNamespace(execution_branches_impacting_inputs=['batch1', 'batch2'])
    bm = BranchingManager(masks={'batch1': mask_common, 'batch2': set(mask_common)})
    codeflash_return_value = codeflash_wrap(construct_mask_for_all_inputs_dimensionalities, 'tests.inference.unit_tests.test_construct_mask_for_all_inputs_dimensionalities__unit_test_0', None, 'test_batch_masks_intersection_basic', 'construct_mask_for_all_inputs_dimensionalities', '4', codeflash_loop_index, step_node=step_node, branching_manager=bm)
    masks_by_dim, scalar_false = codeflash_return_value

def test_disjoint_children_produces_empty_valid_sets(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr('inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler.collect_inputs_dimensionalities', lambda step_node: {'x': 1, 'y': 2})
    mask1 = {(1,), (2, 2)}
    step_node = types.SimpleNamespace(execution_branches_impacting_inputs=['batch_only'])
    bm = BranchingManager(masks={'batch_only': mask1})
    codeflash_return_value = codeflash_wrap(construct_mask_for_all_inputs_dimensionalities, 'tests.inference.unit_tests.test_construct_mask_for_all_inputs_dimensionalities__unit_test_0', None, 'test_disjoint_children_produces_empty_valid_sets', 'construct_mask_for_all_inputs_dimensionalities', '4', codeflash_loop_index, step_node=step_node, branching_manager=bm)
    masks_by_dim, scalar_false = codeflash_return_value

def test_all_zero_dimensionalities_returns_empty_mapping(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr('inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler.collect_inputs_dimensionalities', lambda step_node: {'a': 0, 'b': 0})
    step_node = types.SimpleNamespace(execution_branches_impacting_inputs=[])
    bm = BranchingManager(masks={})
    codeflash_return_value = codeflash_wrap(construct_mask_for_all_inputs_dimensionalities, 'tests.inference.unit_tests.test_construct_mask_for_all_inputs_dimensionalities__unit_test_0', None, 'test_all_zero_dimensionalities_returns_empty_mapping', 'construct_mask_for_all_inputs_dimensionalities', '3', codeflash_loop_index, step_node=step_node, branching_manager=bm)
    masks_by_dim, scalar_false = codeflash_return_value

def test_non_batch_false_among_registered_branch_short_circuits_to_empty_sets(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr('inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler.collect_inputs_dimensionalities', lambda step_node: {'param': 3})
    step_node = types.SimpleNamespace(execution_branches_impacting_inputs=['nb_false', 'batch_ok'])
    bm = BranchingManager(masks={'nb_false': False, 'batch_ok': {(0, 0, 0)}})
    codeflash_return_value = codeflash_wrap(construct_mask_for_all_inputs_dimensionalities, 'tests.inference.unit_tests.test_construct_mask_for_all_inputs_dimensionalities__unit_test_0', None, 'test_non_batch_false_among_registered_branch_short_circuits_to_empty_sets', 'construct_mask_for_all_inputs_dimensionalities', '3', codeflash_loop_index, step_node=step_node, branching_manager=bm)
    masks_by_dim, scalar_false = codeflash_return_value

def test_large_scale_many_identical_batch_masks(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    DIM = 3
    monkeypatch.setattr('inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler.collect_inputs_dimensionalities', lambda step_node: {'bulk': DIM})
    common_mask: Set[Tuple[int, int, int]] = {(0, 0, 0), (0, 1, 0), (1, 0, 0)}
    N = 1000
    masks = {f'b{i}': set(common_mask) for i in range(N)}
    step_node = types.SimpleNamespace(execution_branches_impacting_inputs=list(masks.keys()))
    bm = BranchingManager(masks=masks)
    codeflash_return_value = codeflash_wrap(construct_mask_for_all_inputs_dimensionalities, 'tests.inference.unit_tests.test_construct_mask_for_all_inputs_dimensionalities__unit_test_0', None, 'test_large_scale_many_identical_batch_masks', 'construct_mask_for_all_inputs_dimensionalities', '7', codeflash_loop_index, step_node=step_node, branching_manager=bm)
    masks_by_dim, scalar_false = codeflash_return_value