from __future__ import annotations

import gc
import os
import time
from typing import Dict, List, Union

import pytest
from inference.core.workflows.core_steps.composition.workflow_caller.v1 import \
    _describe_outputs_from_spec
from inference.core.workflows.errors import WorkflowDefinitionError


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

class TestDescribeOutputsFromSpec:
    """Test suite for _describe_outputs_from_spec function."""

    def test_basic_workflow_with_single_output(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test basic workflow definition with a single output.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}], 'outputs': [{'name': 'predictions', 'selector': 'step1.predictions'}], 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_basic_workflow_with_single_output', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output

    def test_basic_workflow_with_wildcard_output(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test workflow definition with wildcard selector for all outputs.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}], 'outputs': [{'name': 'all_outputs', 'selector': 'step1.*'}], 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_basic_workflow_with_wildcard_output', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output

    def test_workflow_with_multiple_outputs(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test workflow with multiple output definitions.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}, {'name': 'step2', 'type': 'Roboflow'}], 'outputs': [{'name': 'output1', 'selector': 'step1.predictions'}, {'name': 'output2', 'selector': 'step2.predictions'}], 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_workflow_with_multiple_outputs', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output

    def test_output_returns_list_or_dict(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that output values are either List[str] or Dict[str, List[str]].'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}], 'outputs': [{'name': 'predictions', 'selector': 'step1.predictions'}], 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_output_returns_list_or_dict', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output
        for key, value in result.items():
            if isinstance(value, dict):
                for nested_key, nested_value in value.items():
                    pass

    def test_missing_steps_raises_error(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        "Test that missing 'steps' key raises WorkflowDefinitionError."
        spec = {'outputs': [{'name': 'predictions', 'selector': 'step1.predictions'}], 'dynamic_blocks_definitions': []}
        with pytest.raises(WorkflowDefinitionError):
            codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_missing_steps_raises_error', '_describe_outputs_from_spec', '2_0', codeflash_loop_index, spec=spec)

    def test_missing_outputs_raises_error(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        "Test that missing 'outputs' key raises WorkflowDefinitionError."
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}], 'dynamic_blocks_definitions': []}
        with pytest.raises(WorkflowDefinitionError):
            codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_missing_outputs_raises_error', '_describe_outputs_from_spec', '2_0', codeflash_loop_index, spec=spec)

    def test_empty_steps_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test workflow with empty steps list.'
        spec = {'steps': [], 'outputs': [{'name': 'predictions', 'selector': 'step1.predictions'}], 'dynamic_blocks_definitions': []}
        with pytest.raises(WorkflowDefinitionError):
            codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_empty_steps_list', '_describe_outputs_from_spec', '2_0', codeflash_loop_index, spec=spec)

    def test_empty_outputs_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test workflow with empty outputs list.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}], 'outputs': [], 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_empty_outputs_list', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output

    def test_step_without_name_raises_error(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        "Test that a step without 'name' raises WorkflowDefinitionError."
        spec = {'steps': [{'type': 'Roboflow'}], 'outputs': [{'name': 'predictions', 'selector': 'step1.predictions'}], 'dynamic_blocks_definitions': []}
        with pytest.raises(WorkflowDefinitionError):
            codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_step_without_name_raises_error', '_describe_outputs_from_spec', '2_0', codeflash_loop_index, spec=spec)

    def test_step_without_type_raises_error(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        "Test that a step without 'type' raises WorkflowDefinitionError."
        spec = {'steps': [{'name': 'step1'}], 'outputs': [{'name': 'predictions', 'selector': 'step1.predictions'}], 'dynamic_blocks_definitions': []}
        with pytest.raises(WorkflowDefinitionError):
            codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_step_without_type_raises_error', '_describe_outputs_from_spec', '2_0', codeflash_loop_index, spec=spec)

    def test_selector_references_nonexistent_step(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that referencing a nonexistent step in selector raises error.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}], 'outputs': [{'name': 'predictions', 'selector': 'step_that_does_not_exist.predictions'}], 'dynamic_blocks_definitions': []}
        with pytest.raises(WorkflowDefinitionError):
            codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_selector_references_nonexistent_step', '_describe_outputs_from_spec', '2_0', codeflash_loop_index, spec=spec)

    def test_selector_references_nonexistent_property(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that referencing a nonexistent property raises error.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}], 'outputs': [{'name': 'invalid_prop', 'selector': 'step1.property_that_does_not_exist'}], 'dynamic_blocks_definitions': []}
        with pytest.raises(WorkflowDefinitionError):
            codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_selector_references_nonexistent_property', '_describe_outputs_from_spec', '2_0', codeflash_loop_index, spec=spec)

    def test_spec_with_no_dynamic_blocks(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test spec with empty dynamic_blocks_definitions.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}], 'outputs': [{'name': 'predictions', 'selector': 'step1.predictions'}], 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_spec_with_no_dynamic_blocks', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output

    def test_spec_without_dynamic_blocks_key(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test spec without dynamic_blocks_definitions key.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}], 'outputs': [{'name': 'predictions', 'selector': 'step1.predictions'}]}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_spec_without_dynamic_blocks_key', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output

    def test_multiple_outputs_from_same_step(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test multiple outputs referencing the same step.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}], 'outputs': [{'name': 'predictions1', 'selector': 'step1.predictions'}, {'name': 'predictions2', 'selector': 'step1.predictions'}], 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_multiple_outputs_from_same_step', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output

    def test_workflow_with_many_steps(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test workflow with a large number of steps (100 steps).'
        steps = [{'name': f'step{i}', 'type': 'Roboflow'} for i in range(100)]
        outputs = [{'name': 'output_first', 'selector': 'step0.predictions'}, {'name': 'output_last', 'selector': 'step99.predictions'}]
        spec = {'steps': steps, 'outputs': outputs, 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_workflow_with_many_steps', '_describe_outputs_from_spec', '4', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output

    def test_workflow_with_many_outputs(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test workflow with a large number of outputs (100 outputs).'
        steps = [{'name': 'step1', 'type': 'Roboflow'}, {'name': 'step2', 'type': 'Roboflow'}]
        outputs = []
        for i in range(100):
            step_name = 'step1' if i % 2 == 0 else 'step2'
            outputs.append({'name': f'output{i}', 'selector': f'{step_name}.predictions'})
        spec = {'steps': steps, 'outputs': outputs, 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_workflow_with_many_outputs', '_describe_outputs_from_spec', '5', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output
        for i in range(100):
            pass

    def test_workflow_with_many_steps_and_outputs(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test workflow with many steps and many outputs (50 steps, 50 outputs).'
        steps = [{'name': f'step{i}', 'type': 'Roboflow'} for i in range(50)]
        outputs = [{'name': f'output{i}', 'selector': f'step{i}.predictions'} for i in range(50)]
        spec = {'steps': steps, 'outputs': outputs, 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_workflow_with_many_steps_and_outputs', '_describe_outputs_from_spec', '4', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output
        for i in range(50):
            pass

    def test_return_type_is_correct(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that the return type matches the expected signature.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}, {'name': 'step2', 'type': 'Roboflow'}], 'outputs': [{'name': 'output1', 'selector': 'step1.predictions'}, {'name': 'output2', 'selector': 'step2.*'}], 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_return_type_is_correct', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output
        for key, value in result.items():
            if isinstance(value, list):
                for elem in value:
                    pass
            elif isinstance(value, dict):
                for nested_key, nested_value in value.items():
                    for elem in nested_value:
                        pass

    def test_spec_with_special_characters_in_names(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test workflow with special characters in step and output names.'
        spec = {'steps': [{'name': 'step_1_a', 'type': 'Roboflow'}, {'name': 'step-2-b', 'type': 'Roboflow'}], 'outputs': [{'name': 'output_1', 'selector': 'step_1_a.predictions'}, {'name': 'output-2', 'selector': 'step-2-b.predictions'}], 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_spec_with_special_characters_in_names', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output

    def test_duplicate_output_names(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that duplicate output names are overwritten.'
        spec = {'steps': [{'name': 'step1', 'type': 'Roboflow'}, {'name': 'step2', 'type': 'Roboflow'}], 'outputs': [{'name': 'same_name', 'selector': 'step1.predictions'}, {'name': 'same_name', 'selector': 'step2.predictions'}], 'dynamic_blocks_definitions': []}
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_1', 'TestDescribeOutputsFromSpec', 'test_duplicate_output_names', '_describe_outputs_from_spec', '2', codeflash_loop_index, spec=spec)
        codeflash_output = codeflash_return_value
        result = codeflash_output