from __future__ import annotations

import gc
import os
import time

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

def test_basic_delegation_returns_value_and_receives_same_object(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Basic test:\n    - Replace the real describe_workflow_outputs with a lightweight function.\n    - Ensure the wrapper returns exactly what that function returns.\n    - Ensure the wrapper passes the exact same spec object (identity) through.\n    '
    received = {}
    import inference.core.workflows.execution_engine.v1.introspection.outputs_discovery as outputs_mod

    def fake_describe_workflow_outputs(definition):
        received['obj'] = definition
        return {'my_output': ['string', 'number']}
    monkeypatch.setattr(outputs_mod, 'describe_workflow_outputs', fake_describe_workflow_outputs)
    spec = {'steps': [], 'outputs': []}
    codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_0', None, 'test_basic_delegation_returns_value_and_receives_same_object', '_describe_outputs_from_spec', '6', codeflash_loop_index, spec)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_propagates_workflow_definition_error_from_inner_function(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Edge case: If the inner describe_workflow_outputs raises a WorkflowDefinitionError,\n    the wrapper must propagate that same exception type to the caller.\n    '
    import inference.core.workflows.execution_engine.v1.introspection.outputs_discovery as outputs_mod

    def fake_raiser(definition):
        raise WorkflowDefinitionError(public_message='invalid workflow', context='unit_test_describe_outputs')
    monkeypatch.setattr(outputs_mod, 'describe_workflow_outputs', fake_raiser)
    with pytest.raises(WorkflowDefinitionError):
        codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_0', None, 'test_propagates_workflow_definition_error_from_inner_function', '_describe_outputs_from_spec', '4_0', codeflash_loop_index, {'steps': [], 'outputs': []})

def test_large_scale_spec_passed_and_validated_by_inner_function(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Large-scale test:\n    - Construct a large spec with up to 1000 steps and 1000 outputs to validate\n      that the wrapper correctly passes large structures through without modification.\n    - The fake inner function will validate sizes and return a computed summary.\n    '
    import inference.core.workflows.execution_engine.v1.introspection.outputs_discovery as outputs_mod
    N = 1000
    steps = [{'name': f'step_{i}', 'type': f'type_{i}'} for i in range(N)]
    outputs = [{'name': f'out_{i}', 'selector': f'step_{i}.value'} for i in range(N)]
    extra_payload = list(range(N))
    spec = {'steps': steps, 'outputs': outputs, 'extra_payload': extra_payload}

    def fake_describe_workflow_outputs(definition):
        return {'summary_count': [str(len(definition.get('extra_payload', [])))]}
    monkeypatch.setattr(outputs_mod, 'describe_workflow_outputs', fake_describe_workflow_outputs)
    codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_0', None, 'test_large_scale_spec_passed_and_validated_by_inner_function', '_describe_outputs_from_spec', '9', codeflash_loop_index, spec)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_wrapper_accepts_none_spec_and_passes_it_through(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Edge test:\n    - The wrapper does not itself validate the input shape; it simply delegates.\n    - Passing None should result in the inner function receiving None. We verify that.\n    '
    import inference.core.workflows.execution_engine.v1.introspection.outputs_discovery as outputs_mod

    def fake_describe_workflow_outputs(definition):
        return {}
    monkeypatch.setattr(outputs_mod, 'describe_workflow_outputs', fake_describe_workflow_outputs)
    codeflash_return_value = codeflash_wrap(_describe_outputs_from_spec, 'tests.inference.unit_tests.test__describe_outputs_from_spec__unit_test_0', None, 'test_wrapper_accepts_none_spec_and_passes_it_through', '_describe_outputs_from_spec', '4', codeflash_loop_index, None)
    codeflash_output = codeflash_return_value
    result = codeflash_output