from __future__ import annotations

import gc
import os
import time

import networkx as nx
import pytest
from inference.core.workflows.errors import ControlFlowDefinitionError
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 \
    verify_compatibility_of_input_data_lineage_with_control_flow_lineage
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 _make_step_node_with_lineage(lineage):
    node = StepNode.__new__(StepNode)
    node.data_lineage = list(lineage)
    return node

def test_returns_immediately_when_inputs_lineage_is_empty():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    graph = nx.DiGraph()
    cf_selector = 'control_step_1'
    graph.add_node(cf_selector)
    graph.nodes[cf_selector][NODE_COMPILATION_OUTPUT_PROPERTY] = _make_step_node_with_lineage(['a', 'b'])
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_0', None, 'test_returns_immediately_when_inputs_lineage_is_empty', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4', codeflash_loop_index, step_name='my_step', inputs_lineage=[], control_flow_steps_selectors=[cf_selector], execution_graph=graph)

def test_no_error_when_control_flow_lineage_is_prefix_of_input():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    graph = nx.DiGraph()
    cf1 = 'cf1'
    cf2 = 'cf2'
    graph.add_node(cf1)
    graph.add_node(cf2)
    control_lineage = ['x', 'y']
    graph.nodes[cf1][NODE_COMPILATION_OUTPUT_PROPERTY] = _make_step_node_with_lineage(control_lineage)
    graph.nodes[cf2][NODE_COMPILATION_OUTPUT_PROPERTY] = _make_step_node_with_lineage(control_lineage)
    inputs_lineage = [['x', 'y', 'z']]
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_0', None, 'test_no_error_when_control_flow_lineage_is_prefix_of_input', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '9', codeflash_loop_index, step_name='consumer_step', inputs_lineage=inputs_lineage, control_flow_steps_selectors=[cf1, cf2], execution_graph=graph)

def test_raises_when_control_flow_lineage_not_covered_by_inputs():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    graph = nx.DiGraph()
    cf = 'cf_problem'
    graph.add_node(cf)
    graph.nodes[cf][NODE_COMPILATION_OUTPUT_PROPERTY] = _make_step_node_with_lineage(['a', 'b'])
    inputs_lineage = [['m', 'n']]
    with pytest.raises(ControlFlowDefinitionError) as excinfo:
        codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_0', None, 'test_raises_when_control_flow_lineage_not_covered_by_inputs', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '5_0', codeflash_loop_index, step_name='target_step', inputs_lineage=inputs_lineage, control_flow_steps_selectors=[cf], execution_graph=graph)
    err = excinfo.value

def test_empty_control_flow_lineage_is_ignored_and_does_not_raise():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    graph = nx.DiGraph()
    cf_empty = 'cf_empty'
    graph.add_node(cf_empty)
    graph.nodes[cf_empty][NODE_COMPILATION_OUTPUT_PROPERTY] = _make_step_node_with_lineage([])
    inputs_lineage = [['alpha', 'beta']]
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_0', None, 'test_empty_control_flow_lineage_is_ignored_and_does_not_raise', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '5', codeflash_loop_index, step_name='s', inputs_lineage=inputs_lineage, control_flow_steps_selectors=[cf_empty], execution_graph=graph)

def test_prefix_matching_with_repeated_and_special_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    graph = nx.DiGraph()
    cf = 'cf_special'
    graph.add_node(cf)
    graph.nodes[cf][NODE_COMPILATION_OUTPUT_PROPERTY] = _make_step_node_with_lineage(['#special'])
    inputs_lineage = [['#special', '#special'], ['#special', 'other']]
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_0', None, 'test_prefix_matching_with_repeated_and_special_characters', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '5', codeflash_loop_index, step_name='special_step', inputs_lineage=inputs_lineage, control_flow_steps_selectors=[cf], execution_graph=graph)

def test_large_scale_many_lineages_and_control_flow_steps():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    graph = nx.DiGraph()
    num_pairs = 500
    control_selectors = []
    inputs_lineage = []
    for i in range(num_pairs):
        lineage = [f'value_{i}', 'common_tail']
        inputs_lineage.append(lineage)
        cf_name = f'cf_{i}'
        control_selectors.append(cf_name)
        graph.add_node(cf_name)
        graph.nodes[cf_name][NODE_COMPILATION_OUTPUT_PROPERTY] = _make_step_node_with_lineage([f'value_{i}'])
    for i in range(10):
        cf_dup = f'cf_dup_{i}'
        graph.add_node(cf_dup)
        graph.nodes[cf_dup][NODE_COMPILATION_OUTPUT_PROPERTY] = _make_step_node_with_lineage([f'value_{i}'])
        control_selectors.append(cf_dup)
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_0', None, 'test_large_scale_many_lineages_and_control_flow_steps', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '6', codeflash_loop_index, step_name='bulk_step', inputs_lineage=inputs_lineage, control_flow_steps_selectors=control_selectors, execution_graph=graph)