from __future__ import annotations

import gc
import inspect
import os
import time
from collections import namedtuple
from pathlib import Path

import dill as pickle
import pytest
from inference.core.workflows.execution_engine.constants import \
    NODE_COMPILATION_OUTPUT_PROPERTY
from inference.core.workflows.execution_engine.v1.compiler.entities import \
    NodeCategory
from inference.core.workflows.execution_engine.v1.compiler.utils import \
    is_control_flow_step
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
NodeCompilationOutput = namedtuple('NodeCompilationOutput', ['node_category', 'controls_flow'])

def test_step_node_controls_flow_true_basic():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    comp_out = NodeCompilationOutput(node_category=NodeCategory.STEP_NODE, controls_flow=lambda: True)
    g.add_node('n1', **{NODE_COMPILATION_OUTPUT_PROPERTY: comp_out})
    _call__bound__arguments = inspect.signature(is_control_flow_step).bind(g, 'n1')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(is_control_flow_step, 'tests.inference.unit_tests.test_is_control_flow_step__unit_test_0', None, 'test_step_node_controls_flow_true_basic', 'is_control_flow_step', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_step_node_controls_flow_false_basic():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    comp_out = NodeCompilationOutput(node_category=NodeCategory.STEP_NODE, controls_flow=lambda: False)
    g.add_node('n2', **{NODE_COMPILATION_OUTPUT_PROPERTY: comp_out})
    _call__bound__arguments = inspect.signature(is_control_flow_step).bind(g, 'n2')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(is_control_flow_step, 'tests.inference.unit_tests.test_is_control_flow_step__unit_test_0', None, 'test_step_node_controls_flow_false_basic', 'is_control_flow_step', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_non_step_node_ignored_even_if_controls_true():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    comp_out = NodeCompilationOutput(node_category=NodeCategory.INPUT_NODE, controls_flow=lambda: True)
    g.add_node('n3', **{NODE_COMPILATION_OUTPUT_PROPERTY: comp_out})
    _call__bound__arguments = inspect.signature(is_control_flow_step).bind(g, 'n3')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(is_control_flow_step, 'tests.inference.unit_tests.test_is_control_flow_step__unit_test_0', None, 'test_non_step_node_ignored_even_if_controls_true', 'is_control_flow_step', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_missing_node_raises_keyerror():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    with pytest.raises(KeyError):
        _call__bound__arguments = inspect.signature(is_control_flow_step).bind(g, 'missing_node')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(is_control_flow_step, 'tests.inference.unit_tests.test_is_control_flow_step__unit_test_0', None, 'test_missing_node_raises_keyerror', 'is_control_flow_step', '1_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_missing_compilation_output_property_raises_keyerror():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    g.add_node('n_no_prop')
    with pytest.raises(KeyError):
        _call__bound__arguments = inspect.signature(is_control_flow_step).bind(g, 'n_no_prop')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(is_control_flow_step, 'tests.inference.unit_tests.test_is_control_flow_step__unit_test_0', None, 'test_missing_compilation_output_property_raises_keyerror', 'is_control_flow_step', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_controls_flow_not_callable_raises_typeerror():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    comp_out = NodeCompilationOutput(node_category=NodeCategory.STEP_NODE, controls_flow=True)
    g.add_node('n_bad_callable', **{NODE_COMPILATION_OUTPUT_PROPERTY: comp_out})
    with pytest.raises(TypeError):
        _call__bound__arguments = inspect.signature(is_control_flow_step).bind(g, 'n_bad_callable')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(is_control_flow_step, 'tests.inference.unit_tests.test_is_control_flow_step__unit_test_0', None, 'test_controls_flow_not_callable_raises_typeerror', 'is_control_flow_step', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_controls_flow_exception_propagated():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()

    def explode():
        raise ValueError('controls_flow failure')
    comp_out = NodeCompilationOutput(node_category=NodeCategory.STEP_NODE, controls_flow=explode)
    g.add_node('n_explode', **{NODE_COMPILATION_OUTPUT_PROPERTY: comp_out})
    with pytest.raises(ValueError):
        _call__bound__arguments = inspect.signature(is_control_flow_step).bind(g, 'n_explode')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(is_control_flow_step, 'tests.inference.unit_tests.test_is_control_flow_step__unit_test_0', None, 'test_controls_flow_exception_propagated', 'is_control_flow_step', '4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_scale_many_nodes_counts():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    total_nodes = 1000
    expected_true_count = 0
    for i in range(total_nodes):
        node_name = f'node_{i}'
        if i % 5 == 0:
            comp_out = NodeCompilationOutput(node_category=NodeCategory.STEP_NODE, controls_flow=lambda i=i: True)
            expected_true_count += 1
        elif i % 3 == 0:
            comp_out = NodeCompilationOutput(node_category=NodeCategory.STEP_NODE, controls_flow=lambda i=i: False)
        else:
            comp_out = NodeCompilationOutput(node_category=NodeCategory.INPUT_NODE, controls_flow=lambda i=i: True)
        g.add_node(node_name, **{NODE_COMPILATION_OUTPUT_PROPERTY: comp_out})
    found_true = 0
    for i in range(total_nodes):
        node_name = f'node_{i}'
        if is_control_flow_step(g, node_name):
            found_true += 1

def test_repeated_calls_stability_on_single_node():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    comp_out = NodeCompilationOutput(node_category=NodeCategory.STEP_NODE, controls_flow=lambda: True)
    g.add_node('stable_node', **{NODE_COMPILATION_OUTPUT_PROPERTY: comp_out})
    for _ in range(1000):
        _call__bound__arguments = inspect.signature(is_control_flow_step).bind(g, 'stable_node')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(is_control_flow_step, 'tests.inference.unit_tests.test_is_control_flow_step__unit_test_0', None, 'test_repeated_calls_stability_on_single_node', 'is_control_flow_step', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value