from __future__ import annotations

import gc
import os
import time

import inference.core.workflows.execution_engine.v1.compiler.graph_constructor as gc_module
import pytest
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    separate_control_flow_predecessors_from_data_providers
from networkx import DiGraph
from networkx.exception import NetworkXError


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_single_control_flow_predecessor(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    G = DiGraph()
    G.add_node('a')
    G.add_node('b')
    G.add_edge('a', 'b')
    monkeypatch.setattr(gc_module, 'is_control_flow_step', lambda execution_graph, node: node == 'a', raising=True)
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_0', None, 'test_single_control_flow_predecessor', 'separate_control_flow_predecessors_from_data_providers', '5', codeflash_loop_index, execution_graph=G, node='b')
    control, non_control = codeflash_return_value

def test_multiple_mixed_predecessors_order_preserved(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    G = DiGraph()
    predecessors = ['p1', 'p2', 'p3']
    target = 'target_node'
    G.add_node(target)
    for p in predecessors:
        G.add_node(p)
        G.add_edge(p, target)
    monkeypatch.setattr(gc_module, 'is_control_flow_step', lambda execution_graph, node: node == 'p2', raising=True)
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_0', None, 'test_multiple_mixed_predecessors_order_preserved', 'separate_control_flow_predecessors_from_data_providers', '6', codeflash_loop_index, execution_graph=G, node=target)
    control, non_control = codeflash_return_value

def test_no_predecessors_returns_empty_lists(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    G = DiGraph()
    G.add_node('isolated')
    monkeypatch.setattr(gc_module, 'is_control_flow_step', lambda execution_graph, node: True, raising=True)
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_0', None, 'test_no_predecessors_returns_empty_lists', 'separate_control_flow_predecessors_from_data_providers', '3', codeflash_loop_index, execution_graph=G, node='isolated')
    control, non_control = codeflash_return_value

def test_self_loop_predecessor_handled(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    G = DiGraph()
    G.add_node('loop')
    G.add_edge('loop', 'loop')
    monkeypatch.setattr(gc_module, 'is_control_flow_step', lambda execution_graph, node: node == 'loop', raising=True)
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_0', None, 'test_self_loop_predecessor_handled', 'separate_control_flow_predecessors_from_data_providers', '4', codeflash_loop_index, execution_graph=G, node='loop')
    control, non_control = codeflash_return_value

def test_special_character_node_names(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    G = DiGraph()
    special_nodes = ['', '!!', '正常']
    target = 'T'
    G.add_node(target)
    for name in special_nodes:
        G.add_node(name)
        G.add_edge(name, target)
    monkeypatch.setattr(gc_module, 'is_control_flow_step', lambda execution_graph, node: node == '', raising=True)
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_0', None, 'test_special_character_node_names', 'separate_control_flow_predecessors_from_data_providers', '6', codeflash_loop_index, execution_graph=G, node=target)
    control, non_control = codeflash_return_value

def test_missing_node_raises_networkx_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    G = DiGraph()
    with pytest.raises(NetworkXError):
        codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_0', None, 'test_missing_node_raises_networkx_error', 'separate_control_flow_predecessors_from_data_providers', '1_0', codeflash_loop_index, execution_graph=G, node='nope')

def test_large_scale_partitioning_even_odd(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    G = DiGraph()
    target = 'big_target'
    G.add_node(target)
    total = 1000
    for i in range(total):
        name = f'n{i}'
        G.add_node(name)
        G.add_edge(name, target)

    def is_control_flow_sim(execution_graph, node):
        if not node.startswith('n'):
            return False
        try:
            idx = int(node[1:])
        except ValueError:
            return False
        return idx % 2 == 0
    monkeypatch.setattr(gc_module, 'is_control_flow_step', is_control_flow_sim, raising=True)
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_0', None, 'test_large_scale_partitioning_even_odd', 'separate_control_flow_predecessors_from_data_providers', '7', codeflash_loop_index, execution_graph=G, node=target)
    control, non_control = codeflash_return_value
    expected_control_prefix = [f'n{i}' for i in range(0, 20, 2)]
    expected_non_control_prefix = [f'n{i}' for i in range(1, 21, 2)]