from __future__ import annotations

import gc
import inspect
import os
import time

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 (
    ExecutionGraphNode, StepNode)
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    _collect_unique_control_flow_lineages_with_step_mapping
from inference.core.workflows.execution_engine.v1.compiler.utils import \
    identify_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(selector: str, data_lineage):
    """
    Create a StepNode instance (real class) and populate the attributes that the
    code under test expects (node_category, name, selector, data_lineage).

    We avoid calling the class constructor directly because the exact signature of
    StepNode is not guaranteed in the provided context. Using __new__ gives us a
    real instance (isinstance will pass) while letting us set the relevant fields.
    """
    node = StepNode.__new__(StepNode)
    setattr(node, 'node_category', None)
    setattr(node, 'name', selector + '_name')
    setattr(node, 'selector', selector)
    setattr(node, 'data_lineage', data_lineage)
    return node

def test_single_non_empty_lineage_basic():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    selector = 'step_a'
    lineage = ['input1', 'input2']
    node = _make_step_node(selector=selector, data_lineage=lineage)
    g.add_node(selector, **{NODE_COMPILATION_OUTPUT_PROPERTY: node})
    codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_0', None, 'test_single_non_empty_lineage_basic', '_collect_unique_control_flow_lineages_with_step_mapping', '5', codeflash_loop_index, control_flow_steps_selectors=[selector], execution_graph=g)
    unique_lineages, lineage_map = codeflash_return_value
    lid = identify_lineage(lineage)

def test_deduplication_identical_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    selector1 = 'step_1'
    selector2 = 'step_2'
    lineage = ['shared', 'value']
    node1 = _make_step_node(selector=selector1, data_lineage=lineage)
    node2 = _make_step_node(selector=selector2, data_lineage=lineage)
    g.add_node(selector1, **{NODE_COMPILATION_OUTPUT_PROPERTY: node1})
    g.add_node(selector2, **{NODE_COMPILATION_OUTPUT_PROPERTY: node2})
    codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_0', None, 'test_deduplication_identical_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '8', codeflash_loop_index, control_flow_steps_selectors=[selector1, selector2], execution_graph=g)
    unique_lineages, lineage_map = codeflash_return_value
    lid = identify_lineage(lineage)

def test_empty_lineage_is_not_in_unique_but_is_mapped():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    sel_empty = 'empty_step'
    sel_nonempty = 'nonempty_step'
    empty_node = _make_step_node(selector=sel_empty, data_lineage=[])
    nonempty_lineage = ['x']
    nonempty_node = _make_step_node(selector=sel_nonempty, data_lineage=nonempty_lineage)
    g.add_node(sel_empty, **{NODE_COMPILATION_OUTPUT_PROPERTY: empty_node})
    g.add_node(sel_nonempty, **{NODE_COMPILATION_OUTPUT_PROPERTY: nonempty_node})
    codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_0', None, 'test_empty_lineage_is_not_in_unique_but_is_mapped', '_collect_unique_control_flow_lineages_with_step_mapping', '8', codeflash_loop_index, control_flow_steps_selectors=[sel_empty, sel_nonempty], execution_graph=g)
    unique_lineages, lineage_map = codeflash_return_value
    empty_lid = identify_lineage([])
    nonempty_lid = identify_lineage(nonempty_lineage)

def test_missing_node_raises_assumption_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    missing_selector = 'i_do_not_exist'
    with pytest.raises(AssumptionError):
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_0', None, 'test_missing_node_raises_assumption_error', '_collect_unique_control_flow_lineages_with_step_mapping', '2_0', codeflash_loop_index, control_flow_steps_selectors=[missing_selector], execution_graph=g)

def test_missing_compilation_output_property_raises_assumption_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    selector = 'node_no_property'
    g.add_node(selector)
    with pytest.raises(AssumptionError):
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_0', None, 'test_missing_compilation_output_property_raises_assumption_error', '_collect_unique_control_flow_lineages_with_step_mapping', '3_0', codeflash_loop_index, control_flow_steps_selectors=[selector], execution_graph=g)

def test_wrong_type_for_compilation_output_raises_assumption_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    selector = 'wrong_type_node'
    exec_node = ExecutionGraphNode(node_category=None, name='nm', selector=selector, data_lineage=['a'])
    g.add_node(selector, **{NODE_COMPILATION_OUTPUT_PROPERTY: exec_node})
    with pytest.raises(AssumptionError):
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_0', None, 'test_wrong_type_for_compilation_output_raises_assumption_error', '_collect_unique_control_flow_lineages_with_step_mapping', '4_0', codeflash_loop_index, control_flow_steps_selectors=[selector], execution_graph=g)

def test_ordering_of_unique_lineages_matches_first_occurrence():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    s1, s2, s3, s4 = ('s1', 's2', 's3', 's4')
    L1 = ['a']
    L2 = ['b']
    L3 = ['c', 'd']
    g.add_node(s1, **{NODE_COMPILATION_OUTPUT_PROPERTY: _make_step_node(s1, L1)})
    g.add_node(s2, **{NODE_COMPILATION_OUTPUT_PROPERTY: _make_step_node(s2, L2)})
    g.add_node(s3, **{NODE_COMPILATION_OUTPUT_PROPERTY: _make_step_node(s3, L1)})
    g.add_node(s4, **{NODE_COMPILATION_OUTPUT_PROPERTY: _make_step_node(s4, L3)})
    codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_0', None, 'test_ordering_of_unique_lineages_matches_first_occurrence', '_collect_unique_control_flow_lineages_with_step_mapping', '9', codeflash_loop_index, control_flow_steps_selectors=[s1, s2, s3, s4], execution_graph=g)
    unique_lineages, lineage_map = codeflash_return_value

def test_large_scale_with_many_selectors_and_repeated_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    total = 1000
    unique_count = 25
    g = DiGraph()
    selectors = []
    unique_lineages_templates = [[f'val_{i}'] for i in range(unique_count)]
    for i in range(total):
        sel = f'node_{i}'
        lineage = unique_lineages_templates[i % unique_count]
        g.add_node(sel, **{NODE_COMPILATION_OUTPUT_PROPERTY: _make_step_node(sel, lineage)})
        selectors.append(sel)
    codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_0', None, 'test_large_scale_with_many_selectors_and_repeated_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '6', codeflash_loop_index, control_flow_steps_selectors=selectors, execution_graph=g)
    unique_lineages_result, lineage_map = codeflash_return_value
    expected_first_seen = [unique_lineages_templates[i] for i in range(unique_count)]
    expected_count_per_lineage = total // unique_count
    total_mapped_selectors = 0
    for lid, sel_list in lineage_map.items():
        total_mapped_selectors += len(sel_list)
    for idx in [0, 5, unique_count - 1]:
        lid = identify_lineage(unique_lineages_templates[idx])

def test_empty_control_flow_selectors_returns_empty_collections():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    g = DiGraph()
    codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_0', None, 'test_empty_control_flow_selectors_returns_empty_collections', '_collect_unique_control_flow_lineages_with_step_mapping', '1', codeflash_loop_index, control_flow_steps_selectors=[], execution_graph=g)
    unique_lineages, mapping = codeflash_return_value