from __future__ import annotations

import gc
import os
import time
from collections import defaultdict
from typing import Dict, List

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}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

def create_step_node(name: str, selector: str, data_lineage: List[str], node_category: str='step') -> StepNode:
    """Create a real StepNode instance for testing."""
    return StepNode(node_category=node_category, name=name, selector=selector, data_lineage=data_lineage)

def add_step_to_graph(graph: nx.DiGraph, step_selector: str, step_node: StepNode) -> None:
    """Add a step node to the execution graph with proper structure."""
    graph.add_node(step_selector)
    graph.nodes[step_selector][NODE_COMPILATION_OUTPUT_PROPERTY] = step_node

def test_single_step_with_single_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with one step containing one lineage item.'
    graph = nx.DiGraph()
    step_node = create_step_node(name='step1', selector='step1', data_lineage=['input.x'])
    add_step_to_graph(graph, 'step1', step_node)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_single_step_with_single_lineage', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=['step1'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_step_with_multiple_lineage_items():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with one step containing multiple lineage items.'
    graph = nx.DiGraph()
    step_node = create_step_node(name='step1', selector='step1', data_lineage=['input.x', 'input.y', 'input.z'])
    add_step_to_graph(graph, 'step1', step_node)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_single_step_with_multiple_lineage_items', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=['step1'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_multiple_steps_same_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with multiple steps that have identical lineages - should deduplicate.'
    graph = nx.DiGraph()
    step_node1 = create_step_node(name='step1', selector='step1', data_lineage=['input.x'])
    step_node2 = create_step_node(name='step2', selector='step2', data_lineage=['input.x'])
    add_step_to_graph(graph, 'step1', step_node1)
    add_step_to_graph(graph, 'step2', step_node2)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_multiple_steps_same_lineage', 'get_lineage_derived_from_control_flow', '6', codeflash_loop_index, control_flow_steps_selectors=['step1', 'step2'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_multiple_steps_different_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with multiple steps that have different lineages.'
    graph = nx.DiGraph()
    step_node1 = create_step_node(name='step1', selector='step1', data_lineage=['input.x'])
    step_node2 = create_step_node(name='step2', selector='step2', data_lineage=['input.y'])
    add_step_to_graph(graph, 'step1', step_node1)
    add_step_to_graph(graph, 'step2', step_node2)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_multiple_steps_different_lineages', 'get_lineage_derived_from_control_flow', '6', codeflash_loop_index, control_flow_steps_selectors=['step1', 'step2'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_three_steps_with_mixed_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with three steps where two share a lineage and one is different.'
    graph = nx.DiGraph()
    step_node1 = create_step_node(name='step1', selector='step1', data_lineage=['input.a'])
    step_node2 = create_step_node(name='step2', selector='step2', data_lineage=['input.b'])
    step_node3 = create_step_node(name='step3', selector='step3', data_lineage=['input.a'])
    add_step_to_graph(graph, 'step1', step_node1)
    add_step_to_graph(graph, 'step2', step_node2)
    add_step_to_graph(graph, 'step3', step_node3)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_three_steps_with_mixed_lineages', 'get_lineage_derived_from_control_flow', '8', codeflash_loop_index, control_flow_steps_selectors=['step1', 'step2', 'step3'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_control_flow_steps_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with empty list of control flow steps.'
    graph = nx.DiGraph()
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_empty_control_flow_steps_list', 'get_lineage_derived_from_control_flow', '2', codeflash_loop_index, control_flow_steps_selectors=[], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_step_with_empty_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a step that has an empty lineage list.'
    graph = nx.DiGraph()
    step_node = create_step_node(name='step1', selector='step1', data_lineage=[])
    add_step_to_graph(graph, 'step1', step_node)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_single_step_with_empty_lineage', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=['step1'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_multiple_steps_all_with_empty_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with multiple steps all having empty lineages.'
    graph = nx.DiGraph()
    for i in range(1, 4):
        step_node = create_step_node(name=f'step{i}', selector=f'step{i}', data_lineage=[])
        add_step_to_graph(graph, f'step{i}', step_node)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_multiple_steps_all_with_empty_lineages', 'get_lineage_derived_from_control_flow', '3', codeflash_loop_index, control_flow_steps_selectors=['step1', 'step2', 'step3'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_mixed_empty_and_nonempty_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a mix of empty and non-empty lineages.'
    graph = nx.DiGraph()
    step_node1 = create_step_node(name='step1', selector='step1', data_lineage=[])
    step_node2 = create_step_node(name='step2', selector='step2', data_lineage=['input.x'])
    step_node3 = create_step_node(name='step3', selector='step3', data_lineage=[])
    step_node4 = create_step_node(name='step4', selector='step4', data_lineage=['input.y'])
    add_step_to_graph(graph, 'step1', step_node1)
    add_step_to_graph(graph, 'step2', step_node2)
    add_step_to_graph(graph, 'step3', step_node3)
    add_step_to_graph(graph, 'step4', step_node4)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_mixed_empty_and_nonempty_lineages', 'get_lineage_derived_from_control_flow', '10', codeflash_loop_index, control_flow_steps_selectors=['step1', 'step2', 'step3', 'step4'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_step_with_single_character_lineage_items():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with single character lineage items.'
    graph = nx.DiGraph()
    step_node = create_step_node(name='step1', selector='step1', data_lineage=['a', 'b', 'c'])
    add_step_to_graph(graph, 'step1', step_node)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_step_with_single_character_lineage_items', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=['step1'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_step_with_long_lineage_items():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with long lineage item strings.'
    graph = nx.DiGraph()
    long_lineage_item = 'very_long_lineage_item_name_' * 10
    step_node = create_step_node(name='step1', selector='step1', data_lineage=[long_lineage_item, long_lineage_item + '_2'])
    add_step_to_graph(graph, 'step1', step_node)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_step_with_long_lineage_items', 'get_lineage_derived_from_control_flow', '5', codeflash_loop_index, control_flow_steps_selectors=['step1'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_step_with_special_characters_in_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with special characters in lineage items.'
    graph = nx.DiGraph()
    step_node = create_step_node(name='step1', selector='step1', data_lineage=['input.x[0]', 'output:data', 'model@v1'])
    add_step_to_graph(graph, 'step1', step_node)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_step_with_special_characters_in_lineage', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=['step1'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_missing_node_in_graph_raises_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that referencing a non-existent node raises AssumptionError.'
    graph = nx.DiGraph()
    with pytest.raises(AssumptionError) as exc_info:
        codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_missing_node_in_graph_raises_error', 'get_lineage_derived_from_control_flow', '2_0', codeflash_loop_index, control_flow_steps_selectors=['nonexistent_step'], execution_graph=graph)

def test_node_without_compilation_output_raises_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that a node without the compilation output property raises error.'
    graph = nx.DiGraph()
    graph.add_node('step1')
    with pytest.raises(AssumptionError) as exc_info:
        codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_node_without_compilation_output_raises_error', 'get_lineage_derived_from_control_flow', '3_0', codeflash_loop_index, control_flow_steps_selectors=['step1'], execution_graph=graph)

def test_order_preservation_in_results():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that the order of results is consistent.'
    graph = nx.DiGraph()
    step_node1 = create_step_node(name='step1', selector='step1', data_lineage=['lineage_a'])
    step_node2 = create_step_node(name='step2', selector='step2', data_lineage=['lineage_b'])
    step_node3 = create_step_node(name='step3', selector='step3', data_lineage=['lineage_c'])
    add_step_to_graph(graph, 'step1', step_node1)
    add_step_to_graph(graph, 'step2', step_node2)
    add_step_to_graph(graph, 'step3', step_node3)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_order_preservation_in_results', 'get_lineage_derived_from_control_flow', '8', codeflash_loop_index, control_flow_steps_selectors=['step1', 'step2', 'step3'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result1 = codeflash_output
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_order_preservation_in_results', 'get_lineage_derived_from_control_flow', '10', codeflash_loop_index, control_flow_steps_selectors=['step1', 'step2', 'step3'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result2 = codeflash_output

def test_unicode_characters_in_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with unicode characters in lineage items.'
    graph = nx.DiGraph()
    step_node = create_step_node(name='step1', selector='step1', data_lineage=['input_alpha', 'input_beta', 'input_gamma'])
    add_step_to_graph(graph, 'step1', step_node)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_unicode_characters_in_lineage', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=['step1'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_numeric_step_selectors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with numeric step selectors (as strings).'
    graph = nx.DiGraph()
    step_node = create_step_node(name='step1', selector='1', data_lineage=['input.x'])
    add_step_to_graph(graph, '1', step_node)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_numeric_step_selectors', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=['1'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_hundred_steps_all_different_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with 100 steps each having different lineages.'
    graph = nx.DiGraph()
    step_selectors = []
    for i in range(100):
        selector = f'step_{i}'
        step_node = create_step_node(name=f'step_{i}', selector=selector, data_lineage=[f'input_{i}'])
        add_step_to_graph(graph, selector, step_node)
        step_selectors.append(selector)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_hundred_steps_all_different_lineages', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=step_selectors, execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_hundred_steps_single_shared_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with 100 steps all sharing the same lineage.'
    graph = nx.DiGraph()
    step_selectors = []
    for i in range(100):
        selector = f'step_{i}'
        step_node = create_step_node(name=f'step_{i}', selector=selector, data_lineage=['shared_input'])
        add_step_to_graph(graph, selector, step_node)
        step_selectors.append(selector)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_hundred_steps_single_shared_lineage', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=step_selectors, execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_hundred_steps_with_multiple_items_per_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with 100 steps each having 50 items in their lineage.'
    graph = nx.DiGraph()
    step_selectors = []
    for i in range(100):
        selector = f'step_{i}'
        lineage = [f'input_{i}_item_{j}' for j in range(50)]
        step_node = create_step_node(name=f'step_{i}', selector=selector, data_lineage=lineage)
        add_step_to_graph(graph, selector, step_node)
        step_selectors.append(selector)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_hundred_steps_with_multiple_items_per_lineage', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=step_selectors, execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_thousand_steps_with_pattern_deduplication():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with 1000 steps using a pattern that creates exactly 10 unique lineages.'
    graph = nx.DiGraph()
    step_selectors = []
    for i in range(1000):
        selector = f'step_{i}'
        lineage_id = i // 100
        step_node = create_step_node(name=f'step_{i}', selector=selector, data_lineage=[f'lineage_{lineage_id}'])
        add_step_to_graph(graph, selector, step_node)
        step_selectors.append(selector)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_thousand_steps_with_pattern_deduplication', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=step_selectors, execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_mixed_empty_and_full_lineages_large_scale():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with 500 steps mixing empty and full lineages.'
    graph = nx.DiGraph()
    step_selectors = []
    for i in range(500):
        selector = f'step_{i}'
        if i % 2 == 0:
            lineage = []
        else:
            lineage = [f'input_{i}']
        step_node = create_step_node(name=f'step_{i}', selector=selector, data_lineage=lineage)
        add_step_to_graph(graph, selector, step_node)
        step_selectors.append(selector)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_mixed_empty_and_full_lineages_large_scale', 'get_lineage_derived_from_control_flow', '4', codeflash_loop_index, control_flow_steps_selectors=step_selectors, execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_return_type_is_list_of_lists():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that the return type is always a list of lists of strings.'
    graph = nx.DiGraph()
    step_node1 = create_step_node(name='step1', selector='step1', data_lineage=['a'])
    step_node2 = create_step_node(name='step2', selector='step2', data_lineage=['b', 'c'])
    add_step_to_graph(graph, 'step1', step_node1)
    add_step_to_graph(graph, 'step2', step_node2)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_return_type_is_list_of_lists', 'get_lineage_derived_from_control_flow', '6', codeflash_loop_index, control_flow_steps_selectors=['step1', 'step2'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for lineage in result:
        for item in lineage:
            pass

def test_result_contains_original_lineage_data():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that results contain the original lineage data, not copies/transformations.'
    graph = nx.DiGraph()
    original_lineage = ['original_a', 'original_b', 'original_c']
    step_node = create_step_node(name='step1', selector='step1', data_lineage=original_lineage)
    add_step_to_graph(graph, 'step1', step_node)
    codeflash_return_value = codeflash_wrap(get_lineage_derived_from_control_flow, 'tests.inference.unit_tests.test_get_lineage_derived_from_control_flow__unit_test_1', None, 'test_result_contains_original_lineage_data', 'get_lineage_derived_from_control_flow', '5', codeflash_loop_index, control_flow_steps_selectors=['step1'], execution_graph=graph)
    codeflash_output = codeflash_return_value
    result = codeflash_output