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.entities.operations import \
    OperationDefinition
from inference.core.workflows.core_steps.common.query_language.errors import \
    OperationTypeNotRecognisedError
from inference.core.workflows.core_steps.sinks.email_notification.v2 import \
    apply_operations_to_message_parameters


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_no_operations_returns_same_object_for_mutable_values():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mutable = [1, 2, 3]
    message_parameters = {'my_list': mutable}
    message_parameters_operations = {}
    _call__bound__arguments = inspect.signature(apply_operations_to_message_parameters).bind(message_parameters=message_parameters, message_parameters_operations=message_parameters_operations)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(apply_operations_to_message_parameters, 'tests.inference.unit_tests.test_apply_operations_to_message_parameters__unit_test_0', None, 'test_no_operations_returns_same_object_for_mutable_values', 'apply_operations_to_message_parameters', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_string_to_uppercase_operation_transforms_value():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original = 'Hello World'
    message_parameters = {'greeting': original}
    message_parameters_operations = {'greeting': [OperationDefinition(type='StringToUpperCase')]}
    _call__bound__arguments = inspect.signature(apply_operations_to_message_parameters).bind(message_parameters=message_parameters, message_parameters_operations=message_parameters_operations)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(apply_operations_to_message_parameters, 'tests.inference.unit_tests.test_apply_operations_to_message_parameters__unit_test_0', None, 'test_single_string_to_uppercase_operation_transforms_value', 'apply_operations_to_message_parameters', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_chain_operations_are_applied_in_order():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original = 'HeLLoPyTeSt'
    message_parameters = {'s': original}
    message_parameters_operations = {'s': [OperationDefinition(type='StringToUpperCase'), OperationDefinition(type='StringToLowerCase')]}
    _call__bound__arguments = inspect.signature(apply_operations_to_message_parameters).bind(message_parameters=message_parameters, message_parameters_operations=message_parameters_operations)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(apply_operations_to_message_parameters, 'tests.inference.unit_tests.test_apply_operations_to_message_parameters__unit_test_0', None, 'test_chain_operations_are_applied_in_order', 'apply_operations_to_message_parameters', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_parameters_without_associated_operations_stay_unchanged():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    message_parameters = {'keep': 'value'}
    message_parameters_operations = {'non_existent_key': [OperationDefinition(type='StringToUpperCase')], 'keep': []}
    _call__bound__arguments = inspect.signature(apply_operations_to_message_parameters).bind(message_parameters=message_parameters, message_parameters_operations=message_parameters_operations)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(apply_operations_to_message_parameters, 'tests.inference.unit_tests.test_apply_operations_to_message_parameters__unit_test_0', None, 'test_parameters_without_associated_operations_stay_unchanged', 'apply_operations_to_message_parameters', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_unregistered_operation_type_raises_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    message_parameters = {'a': 'b'}
    message_parameters_operations = {'a': [OperationDefinition(type='THIS_OPERATION_DOES_NOT_EXIST')]}
    with pytest.raises(OperationTypeNotRecognisedError):
        _call__bound__arguments = inspect.signature(apply_operations_to_message_parameters).bind(message_parameters=message_parameters, message_parameters_operations=message_parameters_operations)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(apply_operations_to_message_parameters, 'tests.inference.unit_tests.test_apply_operations_to_message_parameters__unit_test_0', None, 'test_unregistered_operation_type_raises_error', 'apply_operations_to_message_parameters', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_empty_message_parameters_returns_empty_dict():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    message_parameters = {}
    message_parameters_operations = {'anything': [OperationDefinition(type='StringToUpperCase')]}
    _call__bound__arguments = inspect.signature(apply_operations_to_message_parameters).bind(message_parameters=message_parameters, message_parameters_operations=message_parameters_operations)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(apply_operations_to_message_parameters, 'tests.inference.unit_tests.test_apply_operations_to_message_parameters__unit_test_0', None, 'test_empty_message_parameters_returns_empty_dict', 'apply_operations_to_message_parameters', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_none_and_falsy_operations_handling():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    message_parameters = {'x': 42, 'y': 'unchanged'}
    message_parameters_operations = {'x': None, 'y': []}
    _call__bound__arguments = inspect.signature(apply_operations_to_message_parameters).bind(message_parameters=message_parameters, message_parameters_operations=message_parameters_operations)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(apply_operations_to_message_parameters, 'tests.inference.unit_tests.test_apply_operations_to_message_parameters__unit_test_0', None, 'test_none_and_falsy_operations_handling', 'apply_operations_to_message_parameters', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_operations_do_not_affect_other_parameters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    message_parameters = {'p1': 'abc', 'p2': 'def'}
    message_parameters_operations = {'p1': [OperationDefinition(type='StringToUpperCase')]}
    _call__bound__arguments = inspect.signature(apply_operations_to_message_parameters).bind(message_parameters=message_parameters, message_parameters_operations=message_parameters_operations)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(apply_operations_to_message_parameters, 'tests.inference.unit_tests.test_apply_operations_to_message_parameters__unit_test_0', None, 'test_operations_do_not_affect_other_parameters', 'apply_operations_to_message_parameters', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_number_of_parameters_with_single_operation():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    message_parameters = {f'k{i}': f'value{i}' for i in range(n)}
    op = [OperationDefinition(type='StringToUpperCase')]
    message_parameters_operations = {f'k{i}': op for i in range(n)}
    _call__bound__arguments = inspect.signature(apply_operations_to_message_parameters).bind(message_parameters=message_parameters, message_parameters_operations=message_parameters_operations)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(apply_operations_to_message_parameters, 'tests.inference.unit_tests.test_apply_operations_to_message_parameters__unit_test_0', None, 'test_large_number_of_parameters_with_single_operation', 'apply_operations_to_message_parameters', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_long_chain_of_operations_on_single_parameter_performance_and_correctness():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    chain_length = 1000
    operations_chain = [OperationDefinition(type='StringToUpperCase')] * chain_length
    message_parameters = {'s': 'mixedCaseString'}
    message_parameters_operations = {'s': operations_chain}
    _call__bound__arguments = inspect.signature(apply_operations_to_message_parameters).bind(message_parameters=message_parameters, message_parameters_operations=message_parameters_operations)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(apply_operations_to_message_parameters, 'tests.inference.unit_tests.test_apply_operations_to_message_parameters__unit_test_0', None, 'test_long_chain_of_operations_on_single_parameter_performance_and_correctness', 'apply_operations_to_message_parameters', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output