from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path

import dill as pickle
import pytest
from inference.core.workflows.core_steps.composition.workflow_caller.v1 import (
    WORKFLOW_CALLER_BLOCK_TYPE, _extract_workflow_caller_ids_from_spec)


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}######!')
    iteration = os.environ['CODEFLASH_TEST_ITERATION']
    with Path('{codeflash_run_tmp_dir_client_side}', f'test_return_values_{iteration}.bin').open('ab') as f:
        pickled_values = pickle.dumps((args, kwargs, exception)) if exception else pickle.dumps((args, kwargs, return_value))
        _test_name = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_line_id}".encode('ascii')
        f.write(len(_test_name).to_bytes(4, byteorder='big'))
        f.write(_test_name)
        f.write(codeflash_duration.to_bytes(8, byteorder='big'))
        f.write(len(pickled_values).to_bytes(4, byteorder='big'))
        f.write(pickled_values)
        f.write(codeflash_loop_index.to_bytes(8, byteorder='big'))
        f.write(len(invocation_id).to_bytes(4, byteorder='big'))
        f.write(invocation_id.encode('ascii'))
    if exception:
        raise exception
    return return_value

def test_single_valid_step_returns_tuple_of_strings():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    spec = {'steps': [{'type': WORKFLOW_CALLER_BLOCK_TYPE, 'workflow_id': 'workflow_123', 'workflow_version_id': 'version_1'}]}
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind(spec)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_single_valid_step_returns_tuple_of_strings', '_extract_workflow_caller_ids_from_spec', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_wrong_type_steps_are_ignored():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    spec = {'steps': [{'type': 'some/other@v1', 'workflow_id': 'should_not', 'workflow_version_id': 'appear'}]}
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind(spec)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_wrong_type_steps_are_ignored', '_extract_workflow_caller_ids_from_spec', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_selector_workflow_id_is_ignored():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    spec = {'steps': [{'type': WORKFLOW_CALLER_BLOCK_TYPE, 'workflow_id': '$selected_workflow', 'workflow_version_id': 'vX'}]}
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind(spec)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_selector_workflow_id_is_ignored', '_extract_workflow_caller_ids_from_spec', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_version_selector_becomes_none():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    spec = {'steps': [{'type': WORKFLOW_CALLER_BLOCK_TYPE, 'workflow_id': 'wf_ok', 'workflow_version_id': '$dynamic_version'}]}
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind(spec)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_version_selector_becomes_none', '_extract_workflow_caller_ids_from_spec', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_missing_steps_key_and_empty_steps_list_return_empty():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind({})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_missing_steps_key_and_empty_steps_list_return_empty', '_extract_workflow_caller_ids_from_spec', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result_missing = codeflash_output
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind({'steps': []})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_missing_steps_key_and_empty_steps_list_return_empty', '_extract_workflow_caller_ids_from_spec', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result_empty = codeflash_output

def test_step_with_no_type_is_ignored():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    spec = {'steps': [{'workflow_id': 'wf', 'workflow_version_id': 'v'}]}
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind(spec)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_step_with_no_type_is_ignored', '_extract_workflow_caller_ids_from_spec', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_non_string_workflow_id_is_ignored_but_non_string_version_is_preserved():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    spec = {'steps': [{'type': WORKFLOW_CALLER_BLOCK_TYPE, 'workflow_id': 12345, 'workflow_version_id': 'should_not_appear'}, {'type': WORKFLOW_CALLER_BLOCK_TYPE, 'workflow_id': 'wf_valid', 'workflow_version_id': 999}]}
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind(spec)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_non_string_workflow_id_is_ignored_but_non_string_version_is_preserved', '_extract_workflow_caller_ids_from_spec', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_string_workflow_id_is_considered_literal_and_kept():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    spec = {'steps': [{'type': WORKFLOW_CALLER_BLOCK_TYPE, 'workflow_id': ''}]}
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind(spec)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_empty_string_workflow_id_is_considered_literal_and_kept', '_extract_workflow_caller_ids_from_spec', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_non_string_version_types_are_returned_unchanged():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    complex_version = ['a', 'b', {'c': 3}]
    spec = {'steps': [{'type': WORKFLOW_CALLER_BLOCK_TYPE, 'workflow_id': 'wf_with_complex_version', 'workflow_version_id': complex_version}]}
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind(spec)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_non_string_version_types_are_returned_unchanged', '_extract_workflow_caller_ids_from_spec', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_many_steps_performance_and_correctness():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    steps = []
    expected = []
    total = 1000
    for i in range(total):
        mod = i % 4
        if mod == 0:
            wf_id = f'wf_{i}'
            if i % 8 == 0:
                wf_ver = f'v_{i}'
            else:
                wf_ver = f'${{selector_v_{i}}}'
            steps.append({'type': WORKFLOW_CALLER_BLOCK_TYPE, 'workflow_id': wf_id, 'workflow_version_id': wf_ver})
            expected.append((wf_id, None if isinstance(wf_ver, str) and wf_ver.startswith('$') else wf_ver))
        elif mod == 1:
            steps.append({'type': WORKFLOW_CALLER_BLOCK_TYPE, 'workflow_id': f'${{sel_wf_{i}}}', 'workflow_version_id': f'v_{i}'})
        elif mod == 2:
            steps.append({'type': 'not_the_workflow_caller_block', 'workflow_id': f'wf_{i}', 'workflow_version_id': f'v_{i}'})
        else:
            steps.append({'type': WORKFLOW_CALLER_BLOCK_TYPE, 'workflow_id': i, 'workflow_version_id': f'v_{i}'})
    spec = {'steps': steps}
    _call__bound__arguments = inspect.signature(_extract_workflow_caller_ids_from_spec).bind(spec)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_extract_workflow_caller_ids_from_spec, 'tests.inference.unit_tests.test__extract_workflow_caller_ids_from_spec__unit_test_0', None, 'test_large_scale_many_steps_performance_and_correctness', '_extract_workflow_caller_ids_from_spec', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output