from __future__ import annotations

import gc
import os
import time

import pytest
from inference.core.workflows.errors import ExecutionEngineRuntimeError
from inference.core.workflows.execution_engine.v1.compiler.entities import \
    StepNode
from inference.core.workflows.execution_engine.v1.entities import FlowControl
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.dynamic_batches_manager import \
    DynamicBatchesManager
from inference.core.workflows.execution_engine.v1.executor.execution_data_manager.execution_cache import \
    ExecutionCache
from inference.core.workflows.execution_engine.v1.executor.execution_data_manager.manager import \
    ExecutionDataManager
from networkx import DiGraph


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_raises_when_step_has_no_child_execution_branches():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    exec_graph = DiGraph()
    exec_cache = ExecutionCache({}, {}, set())
    dyn_batches_mgr = DynamicBatchesManager({})
    branching_mgr = BranchingManager({})
    edm = ExecutionDataManager(execution_graph=exec_graph, runtime_parameters={}, execution_cache=exec_cache, dynamic_batches_manager=dyn_batches_mgr, branching_manager=branching_mgr)
    step = StepNode.__new__(StepNode)
    step.name = 'step_without_children'
    step.child_execution_branches = {}
    fc = FlowControl(context=None)
    with pytest.raises(ExecutionEngineRuntimeError):
        codeflash_return_value = codeflash_wrap(edm._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_0', None, 'test_raises_when_step_has_no_child_execution_branches', '_register_control_flow_output_for_non_simd_step', '9_0', codeflash_loop_index, step_node=step, output=fc)

def test_registers_masks_for_single_string_context_selection():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    exec_graph = DiGraph()
    exec_cache = ExecutionCache({}, {}, set())
    dyn_batches_mgr = DynamicBatchesManager({})
    branching_mgr = BranchingManager({})
    edm = ExecutionDataManager(execution_graph=exec_graph, runtime_parameters={}, execution_cache=exec_cache, dynamic_batches_manager=dyn_batches_mgr, branching_manager=branching_mgr)
    step = StepNode.__new__(StepNode)
    step.name = 'flow_step'
    step.child_execution_branches = {'target_A': 'branchA', 'target_B': 'branchB'}
    fc = FlowControl(context='target_A')
    codeflash_return_value = codeflash_wrap(edm._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_0', None, 'test_registers_masks_for_single_string_context_selection', '_register_control_flow_output_for_non_simd_step', '9', codeflash_loop_index, step_node=step, output=fc)

def test_registers_masks_for_list_context_selection_with_duplicates_and_special_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    exec_graph = DiGraph()
    exec_cache = ExecutionCache({}, {}, set())
    dyn_batches_mgr = DynamicBatchesManager({})
    branching_mgr = BranchingManager({})
    edm = ExecutionDataManager(execution_graph=exec_graph, runtime_parameters={}, execution_cache=exec_cache, dynamic_batches_manager=dyn_batches_mgr, branching_manager=branching_mgr)
    step = StepNode.__new__(StepNode)
    step.name = 'flow_special'
    step.child_execution_branches = {'t$1': 'b$1', 't-2': 'b-2', 't 3': 'b 3', 'duplicate': 'bdup'}
    fc = FlowControl(context=['t$1', 'duplicate', 'duplicate'])
    codeflash_return_value = codeflash_wrap(edm._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_0', None, 'test_registers_masks_for_list_context_selection_with_duplicates_and_special_names', '_register_control_flow_output_for_non_simd_step', '9', codeflash_loop_index, step_node=step, output=fc)

def test_none_context_matches_none_child_key():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    exec_graph = DiGraph()
    exec_cache = ExecutionCache({}, {}, set())
    dyn_batches_mgr = DynamicBatchesManager({})
    branching_mgr = BranchingManager({})
    edm = ExecutionDataManager(execution_graph=exec_graph, runtime_parameters={}, execution_cache=exec_cache, dynamic_batches_manager=dyn_batches_mgr, branching_manager=branching_mgr)
    step = StepNode.__new__(StepNode)
    step.name = 'flow_none'
    step.child_execution_branches = {None: 'branch_none', 's': 'branch_s'}
    fc = FlowControl(context=None)
    codeflash_return_value = codeflash_wrap(edm._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_0', None, 'test_none_context_matches_none_child_key', '_register_control_flow_output_for_non_simd_step', '9', codeflash_loop_index, step_node=step, output=fc)

def test_raises_when_same_execution_branch_registered_twice_by_child_mapping():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    exec_graph = DiGraph()
    exec_cache = ExecutionCache({}, {}, set())
    dyn_batches_mgr = DynamicBatchesManager({})
    branching_mgr = BranchingManager({})
    edm = ExecutionDataManager(execution_graph=exec_graph, runtime_parameters={}, execution_cache=exec_cache, dynamic_batches_manager=dyn_batches_mgr, branching_manager=branching_mgr)
    step = StepNode.__new__(StepNode)
    step.name = 'flow_dup_branch'
    step.child_execution_branches = {'t1': 'shared_branch', 't2': 'shared_branch'}
    fc = FlowControl(context='t1')
    with pytest.raises(ExecutionEngineRuntimeError):
        codeflash_return_value = codeflash_wrap(edm._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_0', None, 'test_raises_when_same_execution_branch_registered_twice_by_child_mapping', '_register_control_flow_output_for_non_simd_step', '9_0', codeflash_loop_index, step_node=step, output=fc)

def test_large_scale_many_child_branches_half_selected():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    exec_graph = DiGraph()
    exec_cache = ExecutionCache({}, {}, set())
    dyn_batches_mgr = DynamicBatchesManager({})
    branching_mgr = BranchingManager({})
    edm = ExecutionDataManager(execution_graph=exec_graph, runtime_parameters={}, execution_cache=exec_cache, dynamic_batches_manager=dyn_batches_mgr, branching_manager=branching_mgr)
    total = 1000
    step = StepNode.__new__(StepNode)
    step.name = 'flow_large'
    child_map = {}
    selected_targets = []
    for i in range(total):
        target = f't{i}'
        branch = f'b{i}'
        child_map[target] = branch
        if i % 2 == 0:
            selected_targets.append(target)
    step.child_execution_branches = child_map
    fc = FlowControl(context=selected_targets)
    codeflash_return_value = codeflash_wrap(edm._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_0', None, 'test_large_scale_many_child_branches_half_selected', '_register_control_flow_output_for_non_simd_step', '13', codeflash_loop_index, step_node=step, output=fc)
    for i in range(total):
        branch = f'b{i}'
        expected = i % 2 == 0