from __future__ import annotations

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

import dill as pickle
import pytest
from inference.core.interfaces.webrtc_worker.webrtc import default_encoder


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_basic_bytes_are_base64_encoded_to_ascii_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    data = b'hello'
    expected = base64.b64encode(data).decode('ascii')
    _call__bound__arguments = inspect.signature(default_encoder).bind(data)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(default_encoder, 'tests.inference.unit_tests.test_default_encoder__unit_test_0', None, 'test_basic_bytes_are_base64_encoded_to_ascii_string', 'default_encoder', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_bytes_encodes_to_empty_string():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    data = b''
    expected = base64.b64encode(data).decode('ascii')
    _call__bound__arguments = inspect.signature(default_encoder).bind(data)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(default_encoder, 'tests.inference.unit_tests.test_default_encoder__unit_test_0', None, 'test_empty_bytes_encodes_to_empty_string', 'default_encoder', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_non_bytes_types_are_returned_unchanged_value():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    s = 'aGVsbG8='
    _call__bound__arguments = inspect.signature(default_encoder).bind(s)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(default_encoder, 'tests.inference.unit_tests.test_default_encoder__unit_test_0', None, 'test_non_bytes_types_are_returned_unchanged_value', 'default_encoder', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result_s = codeflash_output
    n = 12345
    _call__bound__arguments = inspect.signature(default_encoder).bind(n)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(default_encoder, 'tests.inference.unit_tests.test_default_encoder__unit_test_0', None, 'test_non_bytes_types_are_returned_unchanged_value', 'default_encoder', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result_n = codeflash_output
    none_val = None
    _call__bound__arguments = inspect.signature(default_encoder).bind(none_val)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(default_encoder, 'tests.inference.unit_tests.test_default_encoder__unit_test_0', None, 'test_non_bytes_types_are_returned_unchanged_value', 'default_encoder', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result_none = codeflash_output

def test_mutable_bytes_like_objects_are_not_encoded_and_identity_preserved():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    ba = bytearray(b'abc')
    _call__bound__arguments = inspect.signature(default_encoder).bind(ba)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(default_encoder, 'tests.inference.unit_tests.test_default_encoder__unit_test_0', None, 'test_mutable_bytes_like_objects_are_not_encoded_and_identity_preserved', 'default_encoder', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    res_ba = codeflash_output
    mv = memoryview(b'xyz')
    _call__bound__arguments = inspect.signature(default_encoder).bind(mv)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(default_encoder, 'tests.inference.unit_tests.test_default_encoder__unit_test_0', None, 'test_mutable_bytes_like_objects_are_not_encoded_and_identity_preserved', 'default_encoder', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    res_mv = codeflash_output

def test_non_ascii_bytes_are_encoded_correctly():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    data = bytes([255, 0, 16, 128])
    expected = base64.b64encode(data).decode('ascii')
    _call__bound__arguments = inspect.signature(default_encoder).bind(data)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(default_encoder, 'tests.inference.unit_tests.test_default_encoder__unit_test_0', None, 'test_non_ascii_bytes_are_encoded_correctly', 'default_encoder', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_many_iterations_stability_and_performance_loop_1000():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    for i in range(1000):
        payload = bytes([i * 37 % 256])
        expected = base64.b64encode(payload).decode('ascii')
        _call__bound__arguments = inspect.signature(default_encoder).bind(payload)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(default_encoder, 'tests.inference.unit_tests.test_default_encoder__unit_test_0', None, 'test_many_iterations_stability_and_performance_loop_1000', 'default_encoder', '0_2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_list_of_1000_byte_values_all_encoded_correctly():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = [bytes([i % 256]) for i in range(1000)]
    outputs = [default_encoder(x) for x in inputs]
    expected = [base64.b64encode(x).decode('ascii') for x in inputs]
    for out, exp in zip(outputs, expected):
        pass

def test_bytearray_mutation_after_call_reflects_same_object_identity():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    ba = bytearray(b'initial')
    _call__bound__arguments = inspect.signature(default_encoder).bind(ba)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(default_encoder, 'tests.inference.unit_tests.test_default_encoder__unit_test_0', None, 'test_bytearray_mutation_after_call_reflects_same_object_identity', 'default_encoder', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    returned = codeflash_output
    ba.extend(b'_more')