from __future__ import annotations

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

import dill as pickle
import networkx as nx
import pytest
from inference.core.workflows.errors import AssumptionError
from inference.core.workflows.execution_engine.constants import \
    NODE_COMPILATION_OUTPUT_PROPERTY
from inference.core.workflows.execution_engine.v1.compiler.entities import \
    StepNode
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    get_lineage_derived_from_control_flow
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

def _make_graph_with_node(selector, node_obj=None, extra_attrs=None):
    """
    Helper to create a DiGraph with a single node 'selector'.
    node_obj is placed under NODE_COMPILATION_OUTPUT_PROPERTY if provided.
    extra_attrs can add other attributes on the node mapping.
    """
    g = nx.DiGraph()
    g.add_node(selector)
    if node_obj is not None:
        g.nodes[selector][NODE_COMPILATION_OUTPUT_PROPERTY] = node_obj
    if extra_attrs:
        for k, v in extra_attrs.items():
            g.nodes[selector][k] = v
    return g

def test_empty_control_flow_selectors_returns_empty_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = nx.DiGraph()
    _call__bound__arguments = inspect.signature(get_lineage_derived_from_control_flow).bind([], g)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_0', None, 'test_empty_control_flow_selectors_returns_empty_list', 'get_lineage_derived_from_control_flow', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_step_returns_its_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    selector = 'step1'
    node = StepNode(node_category=None, name='step1-name', selector=selector, data_lineage=['a', 'b'])
    g = _make_graph_with_node(selector, node_obj=node)
    _call__bound__arguments = inspect.signature(get_lineage_derived_from_control_flow).bind([selector], g)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_0', None, 'test_single_step_returns_its_lineage', 'get_lineage_derived_from_control_flow', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_multiple_distinct_lineages_preserve_order():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = nx.DiGraph()
    s1 = 's1'
    s2 = 's2'
    node1 = StepNode(node_category=None, name='n1', selector=s1, data_lineage=['x'])
    node2 = StepNode(node_category=None, name='n2', selector=s2, data_lineage=['y'])
    g.add_node(s1)
    g.add_node(s2)
    g.nodes[s1][NODE_COMPILATION_OUTPUT_PROPERTY] = node1
    g.nodes[s2][NODE_COMPILATION_OUTPUT_PROPERTY] = node2
    _call__bound__arguments = inspect.signature(get_lineage_derived_from_control_flow).bind([s1, s2], g)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_0', None, 'test_multiple_distinct_lineages_preserve_order', 'get_lineage_derived_from_control_flow', '9', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_lineages_with_same_elements_different_order_are_deduplicated():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    s1 = 'step_a'
    s2 = 'step_b'
    node1 = StepNode(node_category=None, name='n1', selector=s1, data_lineage=['a', 'b'])
    node2 = StepNode(node_category=None, name='n2', selector=s2, data_lineage=['b', 'a'])
    g = nx.DiGraph()
    g.add_node(s1)
    g.add_node(s2)
    g.nodes[s1][NODE_COMPILATION_OUTPUT_PROPERTY] = node1
    g.nodes[s2][NODE_COMPILATION_OUTPUT_PROPERTY] = node2
    _call__bound__arguments = inspect.signature(get_lineage_derived_from_control_flow).bind([s1, s2], g)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_0', None, 'test_lineages_with_same_elements_different_order_are_deduplicated', 'get_lineage_derived_from_control_flow', '9', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_lineages_are_omitted_from_result():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    s_empty = 'empty'
    s_nonempty = 'nonempty'
    node_empty = StepNode(node_category=None, name='ne', selector=s_empty, data_lineage=[])
    node_nonempty = StepNode(node_category=None, name='n', selector=s_nonempty, data_lineage=['only'])
    g = nx.DiGraph()
    g.add_node(s_empty)
    g.add_node(s_nonempty)
    g.nodes[s_empty][NODE_COMPILATION_OUTPUT_PROPERTY] = node_empty
    g.nodes[s_nonempty][NODE_COMPILATION_OUTPUT_PROPERTY] = node_nonempty
    _call__bound__arguments = inspect.signature(get_lineage_derived_from_control_flow).bind([s_empty, s_nonempty], g)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_0', None, 'test_empty_lineages_are_omitted_from_result', 'get_lineage_derived_from_control_flow', '9', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_missing_node_in_graph_raises_assumption_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = nx.DiGraph()
    missing_selector = 'i_do_not_exist'
    with pytest.raises(AssumptionError):
        _call__bound__arguments = inspect.signature(get_lineage_derived_from_control_flow).bind([missing_selector], g)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_0', None, 'test_missing_node_in_graph_raises_assumption_error', 'get_lineage_derived_from_control_flow', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_missing_compilation_output_property_raises_assumption_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    selector = 'node_no_prop'
    g = nx.DiGraph()
    g.add_node(selector)
    with pytest.raises(AssumptionError):
        _call__bound__arguments = inspect.signature(get_lineage_derived_from_control_flow).bind([selector], g)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_0', None, 'test_missing_compilation_output_property_raises_assumption_error', 'get_lineage_derived_from_control_flow', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_incorrect_node_type_under_property_raises_assumption_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    selector = 'bad_type'
    g = nx.DiGraph()
    g.add_node(selector)
    g.nodes[selector][NODE_COMPILATION_OUTPUT_PROPERTY] = 12345
    with pytest.raises(AssumptionError):
        _call__bound__arguments = inspect.signature(get_lineage_derived_from_control_flow).bind([selector], g)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_0', None, 'test_incorrect_node_type_under_property_raises_assumption_error', 'get_lineage_derived_from_control_flow', '4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_scale_many_nodes_deduped_and_order_preserved():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    total_nodes = 1000
    unique_count = 50
    g = nx.DiGraph()
    selectors = []
    for i in range(total_nodes):
        sel = f'n{i}'
        selectors.append(sel)
        g.add_node(sel)
        if i % 60 >= unique_count:
            lineage = []
        else:
            u = i % unique_count
            lineage = [f'v{u}']
        node = StepNode(node_category=None, name=f'name{i}', selector=sel, data_lineage=lineage)
        g.nodes[sel][NODE_COMPILATION_OUTPUT_PROPERTY] = node
    _call__bound__arguments = inspect.signature(get_lineage_derived_from_control_flow).bind(selectors, g)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_0', None, 'test_large_scale_many_nodes_deduped_and_order_preserved', 'get_lineage_derived_from_control_flow', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = [[f'v{i}'] for i in range(unique_count)]