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.common.query_language.errors import \
    InvalidInputTypeError
from inference.core.workflows.core_steps.common.query_language.operations.sequences.base import \
    sequence_apply


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_sequence_apply_with_list_and_kwargs():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    values = [1, 2, 3]

    def add_fn(x, add=0):
        return x + add
    _call__bound__arguments = inspect.signature(sequence_apply).bind(values, add_fn, execution_context='basic_ctx', add=10)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_with_list_and_kwargs', 'sequence_apply', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_sequence_apply_with_tuple_and_string_suffix():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    values = ('apple', 'banana')

    def suffix_fn(s, suffix=''):
        return s + suffix
    _call__bound__arguments = inspect.signature(sequence_apply).bind(values, suffix_fn, execution_context='tuple_ctx', suffix='_X')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_with_tuple_and_string_suffix', 'sequence_apply', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_sequence_apply_with_empty_sequence_returns_empty_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    values = []

    def identity(x):
        return x
    _call__bound__arguments = inspect.signature(sequence_apply).bind(values, identity, execution_context='empty_ctx')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_with_empty_sequence_returns_empty_list', 'sequence_apply', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_sequence_apply_iterates_string_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    value = 'abC'

    def to_lower(ch):
        return ch.lower()
    _call__bound__arguments = inspect.signature(sequence_apply).bind(value, to_lower, execution_context='string_ctx')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_iterates_string_characters', 'sequence_apply', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_sequence_apply_with_generator_input():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    gen = (i * 2 for i in range(3))

    def inc(x, inc_by=1):
        return x + inc_by
    _call__bound__arguments = inspect.signature(sequence_apply).bind(gen, inc, execution_context='gen_ctx', inc_by=1)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_with_generator_input', 'sequence_apply', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_sequence_apply_with_dict_iterates_keys_in_insertion_order():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    d = {'first': 1, 'second': 2}

    def tag_key(k, tag=''):
        return f'{k}{tag}'
    _call__bound__arguments = inspect.signature(sequence_apply).bind(d, tag_key, execution_context='dict_ctx', tag='_K')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_with_dict_iterates_keys_in_insertion_order', 'sequence_apply', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_sequence_apply_with_none_in_sequence_and_special_handling():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    values = [None, 2]

    def handle_none(v):
        return 'NONE' if v is None else v * 2
    _call__bound__arguments = inspect.signature(sequence_apply).bind(values, handle_none, execution_context='none_ctx')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_with_none_in_sequence_and_special_handling', 'sequence_apply', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_sequence_apply_with_non_iterable_raises_invalid_input_type_error_and_contains_context():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    non_iterable = 42
    with pytest.raises(InvalidInputTypeError) as excinfo:
        _call__bound__arguments = inspect.signature(sequence_apply).bind(non_iterable, lambda x: x, execution_context='non_iter_ctx')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_with_non_iterable_raises_invalid_input_type_error_and_contains_context', 'sequence_apply', '1_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    exc = excinfo.value

def test_sequence_apply_wraps_typeerror_and_valueerror_from_element_processing_into_invalid_input_type_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    values = [1, 2, 3]

    def sometimes_type_error(x):
        if x == 2:
            raise TypeError('bad type')
        return x
    with pytest.raises(InvalidInputTypeError) as excinfo_type:
        _call__bound__arguments = inspect.signature(sequence_apply).bind(values, sometimes_type_error, execution_context='wrap_ctx')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_wraps_typeerror_and_valueerror_from_element_processing_into_invalid_input_type_error', 'sequence_apply', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    exc_type = excinfo_type.value

    def sometimes_value_error(x):
        if x == 3:
            raise ValueError('bad value')
        return x
    with pytest.raises(InvalidInputTypeError) as excinfo_value:
        _call__bound__arguments = inspect.signature(sequence_apply).bind(values, sometimes_value_error, execution_context='wrap_ctx_val')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_wraps_typeerror_and_valueerror_from_element_processing_into_invalid_input_type_error', 'sequence_apply', '5_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    exc_value = excinfo_value.value

def test_sequence_apply_does_not_swallow_other_exceptions_runtimeerror_propagates():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    values = [0, 1]

    def sometimes_runtime_error(x):
        if x == 1:
            raise RuntimeError('boom')
        return x
    with pytest.raises(RuntimeError):
        _call__bound__arguments = inspect.signature(sequence_apply).bind(values, sometimes_runtime_error, execution_context='runtime_ctx')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_does_not_swallow_other_exceptions_runtimeerror_propagates', 'sequence_apply', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_sequence_apply_large_scale_simple_mapping():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    values = list(range(n))

    def add_offset(x, offset=0):
        return x + offset
    offset = 1000
    _call__bound__arguments = inspect.signature(sequence_apply).bind(values, add_offset, execution_context='large_ctx', offset=offset)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_large_scale_simple_mapping', 'sequence_apply', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_sequence_apply_large_scale_with_generator_1000_elements():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    gen = (i * 3 for i in range(n))

    def times_two(x):
        return x * 2
    _call__bound__arguments = inspect.signature(sequence_apply).bind(gen, times_two, execution_context='large_gen_ctx')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sequence_apply, 'tests.inference.unit_tests.test_sequence_apply__unit_test_0', None, 'test_sequence_apply_large_scale_with_generator_1000_elements', 'sequence_apply', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output