from __future__ import annotations

import gc
import inspect
import os
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Set

import dill as pickle
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}######!')
    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

@dataclass
class MockStepNode:
    """Real StepNode representation for testing."""
    name: str
    child_execution_branches: Dict[str, str] = field(default_factory=dict)

def create_execution_data_manager(child_execution_branches: Optional[Dict[str, str]]=None) -> tuple[ExecutionDataManager, BranchingManager]:
    """
    Helper function to create real ExecutionDataManager instances for testing.
    Returns tuple of (manager, branching_manager) for inspection.
    """
    graph = DiGraph()
    runtime_parameters = {}
    dynamic_batches_manager = DynamicBatchesManager(lineage2indices={})
    branching_manager = BranchingManager(masks={})
    execution_cache = ExecutionCache(cache_content={}, batches_compatibility={}, step_outputs_registered=set())
    manager = ExecutionDataManager(execution_graph=graph, runtime_parameters=runtime_parameters, execution_cache=execution_cache, dynamic_batches_manager=dynamic_batches_manager, branching_manager=branching_manager)
    return (manager, branching_manager)

def test_basic_single_branch_selection():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test basic functionality: registering control flow output with a single\n    selected target step that maps to one branch.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1'})
    flow_output = FlowControl(context='step_a')
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_basic_single_branch_selection', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_basic_multiple_branches_one_selected():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with multiple branches where only one is selected.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2', 'step_c': 'branch_3'})
    flow_output = FlowControl(context='step_b')
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_basic_multiple_branches_one_selected', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_basic_multiple_branches_multiple_selected_as_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with multiple branches where several are selected via list.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2', 'step_c': 'branch_3'})
    flow_output = FlowControl(context=['step_a', 'step_c'])
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_basic_multiple_branches_multiple_selected_as_list', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_no_steps_selected():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test when control flow output selects no target steps (all branches False).\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2'})
    flow_output = FlowControl(context='step_nonexistent')
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_no_steps_selected', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_all_steps_selected():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test when all available steps are selected in the control flow output.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2', 'step_c': 'branch_3'})
    flow_output = FlowControl(context=['step_a', 'step_b', 'step_c'])
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_all_steps_selected', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_error_when_no_child_branches():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that ExecutionEngineRuntimeError is raised when step_node has no\n    child_execution_branches (not a flow control step).\n    '
    manager, _ = create_execution_data_manager()
    step_node = MockStepNode(name='non_flow_control_step', child_execution_branches={})
    flow_output = FlowControl(context='any_step')
    with pytest.raises(ExecutionEngineRuntimeError) as exc_info:
        _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_error_when_no_child_branches', '_register_control_flow_output_for_non_simd_step', '4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_case_single_branch_only():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with exactly one branch in the child_execution_branches.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1'})
    flow_output = FlowControl(context='step_a')
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_edge_case_single_branch_only', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_case_many_branches():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with many branches (100) to check scaling.\n    '
    manager, branching_manager = create_execution_data_manager()
    branches = {f'step_{i}': f'branch_{i}' for i in range(100)}
    step_node = MockStepNode(name='flow_control_step', child_execution_branches=branches)
    selected_steps = [f'step_{i}' for i in range(0, 100, 5)]
    flow_output = FlowControl(context=selected_steps)
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_edge_case_many_branches', '_register_control_flow_output_for_non_simd_step', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    for i in range(100):
        branch_name = f'branch_{i}'
        if i % 5 == 0:
            pass
        else:
            pass

def test_edge_case_empty_list_context():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test FlowControl with an empty list as context (selects no steps).\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2'})
    flow_output = FlowControl(context=[])
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_edge_case_empty_list_context', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_case_none_context():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test FlowControl with None as context (selects no steps).\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2'})
    flow_output = FlowControl(context=None)
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_edge_case_none_context', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_case_duplicate_steps_in_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test FlowControl with duplicate step names in the context list.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2'})
    flow_output = FlowControl(context=['step_a', 'step_a', 'step_a'])
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_edge_case_duplicate_steps_in_list', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_case_special_characters_in_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with special characters in step and branch names.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step-a_1': 'branch:1', 'step.b.2': 'branch|2'})
    flow_output = FlowControl(context='step-a_1')
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_edge_case_special_characters_in_names', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_case_numeric_string_contexts():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with numeric strings in contexts.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'1': 'branch_1', '2': 'branch_2', '3': 'branch_3'})
    flow_output = FlowControl(context=['1', '3'])
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_edge_case_numeric_string_contexts', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_scale_1000_branches():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with 1000 branches to assess scalability and performance.\n    '
    manager, branching_manager = create_execution_data_manager()
    num_branches = 1000
    branches = {f'step_{i}': f'branch_{i}' for i in range(num_branches)}
    step_node = MockStepNode(name='flow_control_step', child_execution_branches=branches)
    selected_steps = [f'step_{i}' for i in range(0, num_branches, 10)]
    flow_output = FlowControl(context=selected_steps)
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_large_scale_1000_branches', '_register_control_flow_output_for_non_simd_step', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_scale_all_selected_1000_branches():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test selecting all branches from 1000 branches.\n    '
    manager, branching_manager = create_execution_data_manager()
    num_branches = 1000
    branches = {f'step_{i}': f'branch_{i}' for i in range(num_branches)}
    step_node = MockStepNode(name='flow_control_step', child_execution_branches=branches)
    selected_steps = [f'step_{i}' for i in range(num_branches)]
    flow_output = FlowControl(context=selected_steps)
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_large_scale_all_selected_1000_branches', '_register_control_flow_output_for_non_simd_step', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    for i in range(num_branches):
        pass

def test_large_scale_none_selected_1000_branches():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test selecting no branches from 1000 branches.\n    '
    manager, branching_manager = create_execution_data_manager()
    num_branches = 1000
    branches = {f'step_{i}': f'branch_{i}' for i in range(num_branches)}
    step_node = MockStepNode(name='flow_control_step', child_execution_branches=branches)
    flow_output = FlowControl(context='nonexistent')
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_large_scale_none_selected_1000_branches', '_register_control_flow_output_for_non_simd_step', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    for i in range(0, num_branches, 100):
        pass

def test_large_scale_many_selected_steps():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with 500 selected steps from 1000 available branches.\n    '
    manager, branching_manager = create_execution_data_manager()
    num_branches = 1000
    branches = {f'step_{i}': f'branch_{i}' for i in range(num_branches)}
    step_node = MockStepNode(name='flow_control_step', child_execution_branches=branches)
    selected_steps = [f'step_{i}' for i in range(0, num_branches, 2)]
    flow_output = FlowControl(context=selected_steps)
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_large_scale_many_selected_steps', '_register_control_flow_output_for_non_simd_step', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    true_count = sum((1 for mask in branching_manager._masks.values() if mask))

def test_large_scale_long_list_duplicates():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with a large list containing many duplicates.\n    '
    manager, branching_manager = create_execution_data_manager()
    num_branches = 100
    branches = {f'step_{i}': f'branch_{i}' for i in range(num_branches)}
    step_node = MockStepNode(name='flow_control_step', child_execution_branches=branches)
    selected_steps = ['step_0', 'step_1', 'step_2'] * 333 + ['step_0']
    flow_output = FlowControl(context=selected_steps)
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_large_scale_long_list_duplicates', '_register_control_flow_output_for_non_simd_step', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_scale_complex_branch_mapping():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with complex mappings where multiple branches go to different names.\n    '
    manager, branching_manager = create_execution_data_manager()
    branches = {}
    for i in range(500):
        if i % 3 == 0:
            step_name = f'step_alpha_{i}'
            branch_name = f'br_alpha_{i}'
        elif i % 3 == 1:
            step_name = f'step_beta_{i}'
            branch_name = f'br_beta_{i}'
        else:
            step_name = f'step_gamma_{i}'
            branch_name = f'br_gamma_{i}'
        branches[step_name] = branch_name
    step_node = MockStepNode(name='flow_control_step', child_execution_branches=branches)
    selected_steps = [f'step_alpha_{i}' for i in range(0, 500, 3) if i % 3 == 0]
    flow_output = FlowControl(context=selected_steps)
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_large_scale_complex_branch_mapping', '_register_control_flow_output_for_non_simd_step', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_mutation_all_branches_must_be_registered():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that verifies ALL branches are registered, not just selected ones.\n    This catches mutations that might skip unselected branches.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2', 'step_c': 'branch_3'})
    flow_output = FlowControl(context='step_a')
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_mutation_all_branches_must_be_registered', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_mutation_boolean_values_must_be_exact():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that boolean values are exactly True or False, not truthy/falsy.\n    This catches mutations that might use wrong boolean values.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2'})
    flow_output = FlowControl(context='step_a')
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_mutation_boolean_values_must_be_exact', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_mutation_return_must_be_none():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that the function returns exactly None, not False or other falsy value.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1'})
    flow_output = FlowControl(context='step_a')
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_mutation_return_must_be_none', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_mutation_selection_must_be_inverted_correctly():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that checks selected vs unselected branches are inverted correctly.\n    This catches mutations that might have inverted logic.\n    '
    manager, branching_manager = create_execution_data_manager()
    step_node = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2', 'step_c': 'branch_3'})
    flow_output = FlowControl(context=['step_a', 'step_c'])
    _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_mutation_selection_must_be_inverted_correctly', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    selected_branches = [('step_a', 'branch_1'), ('step_c', 'branch_3')]
    unselected_branches = [('step_b', 'branch_2')]
    for step, branch in selected_branches:
        pass
    for step, branch in unselected_branches:
        pass

def test_mutation_error_message_contains_step_name():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that the error message for non-flow-control steps includes the step name.\n    This ensures the error is informative and catches mutations of error messages.\n    '
    manager, _ = create_execution_data_manager()
    step_name = 'my_special_step_name'
    step_node = MockStepNode(name=step_name, child_execution_branches={})
    flow_output = FlowControl(context='any')
    with pytest.raises(ExecutionEngineRuntimeError) as exc_info:
        _call__bound__arguments = inspect.signature(manager._register_control_flow_output_for_non_simd_step).bind(step_node=step_node, output=flow_output)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(manager._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_mutation_error_message_contains_step_name', '_register_control_flow_output_for_non_simd_step', '5_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    error_msg = str(exc_info.value.public_message)

def test_mutation_list_vs_string_context_handling():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that both string and list contexts are handled equivalently where applicable.\n    This catches mutations in the context type checking logic.\n    '
    manager1, branching_manager1 = create_execution_data_manager()
    step_node1 = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2'})
    flow_output1 = FlowControl(context='step_a')
    _call__bound__arguments = inspect.signature(manager1._register_control_flow_output_for_non_simd_step).bind(step_node1, flow_output1)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager1._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_mutation_list_vs_string_context_handling', '_register_control_flow_output_for_non_simd_step', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    manager2, branching_manager2 = create_execution_data_manager()
    step_node2 = MockStepNode(name='flow_control_step', child_execution_branches={'step_a': 'branch_1', 'step_b': 'branch_2'})
    flow_output2 = FlowControl(context=['step_a'])
    _call__bound__arguments = inspect.signature(manager2._register_control_flow_output_for_non_simd_step).bind(step_node2, flow_output2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(manager2._register_control_flow_output_for_non_simd_step, 'tests.inference.unit_tests.test__register_control_flow_output_for_non_simd_step__unit_test_1', None, 'test_mutation_list_vs_string_context_handling', '_register_control_flow_output_for_non_simd_step', '8', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)