from __future__ import annotations

import base64
import datetime
import gc
import inspect
import os
import sys
import time
import types
from enum import Enum
from pathlib import Path

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


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

class DataOutputMode(Enum):
    NONE = 0
    ALL = 1
    SPECIFIC = 2

class VideoMetadata:

    def __init__(self, **kwargs):
        self._data = kwargs

    def dict(self):
        return dict(self._data)

class WorkflowImageData:

    def __init__(self, numpy_image, video_metadata=None):
        self.numpy_image = numpy_image
        self.video_metadata = video_metadata

def test_basic_serialization_of_primitives():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    fields = ['num', 'text']
    output = {'num': 42, 'text': 'hello'}
    _call__bound__arguments = inspect.signature(VideoFrameProcessor.serialize_outputs_sync).bind(fields, output, DataOutputMode.NONE)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(VideoFrameProcessor.serialize_outputs_sync, 'tests.inference.unit_tests.test_serialize_outputs_sync__unit_test_0', None, 'test_basic_serialization_of_primitives', 'serialize_outputs_sync', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    serialized, errors = codeflash_return_value

def test_missing_field_reports_error_and_skips_field():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    fields = ['present', 'missing']
    output = {'present': 'ok'}
    _call__bound__arguments = inspect.signature(VideoFrameProcessor.serialize_outputs_sync).bind(fields, output, DataOutputMode.NONE)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(VideoFrameProcessor.serialize_outputs_sync, 'tests.inference.unit_tests.test_serialize_outputs_sync__unit_test_0', None, 'test_missing_field_reports_error_and_skips_field', 'serialize_outputs_sync', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    serialized, errors = codeflash_return_value

def test_workflow_image_data_serialization_and_error_handling():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    image = WorkflowImageData(numpy_image='OK_BYTES', video_metadata=VideoMetadata(camera='cam1'))
    fields = ['img']
    output = {'img': image}
    _call__bound__arguments = inspect.signature(VideoFrameProcessor.serialize_outputs_sync).bind(fields, output, DataOutputMode.NONE)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(VideoFrameProcessor.serialize_outputs_sync, 'tests.inference.unit_tests.test_serialize_outputs_sync__unit_test_0', None, 'test_workflow_image_data_serialization_and_error_handling', 'serialize_outputs_sync', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    serialized, errors = codeflash_return_value
    img_dict = serialized['img']
    decoded = base64.b64decode(img_dict['value'].encode('ascii'))
    image_err = WorkflowImageData(numpy_image='RAISE', video_metadata=None)
    output_err = {'img': image_err}
    _call__bound__arguments = inspect.signature(VideoFrameProcessor.serialize_outputs_sync).bind(['img'], output_err, DataOutputMode.NONE)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(VideoFrameProcessor.serialize_outputs_sync, 'tests.inference.unit_tests.test_serialize_outputs_sync__unit_test_0', None, 'test_workflow_image_data_serialization_and_error_handling', 'serialize_outputs_sync', '8', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    serialized_err, errors_err = codeflash_return_value

def test_data_output_mode_all_skips_images():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    img1 = WorkflowImageData(numpy_image='OK1', video_metadata=None)
    fields = ['img1']
    output = {'img1': img1}
    _call__bound__arguments = inspect.signature(VideoFrameProcessor.serialize_outputs_sync).bind(fields, output, DataOutputMode.ALL)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(VideoFrameProcessor.serialize_outputs_sync, 'tests.inference.unit_tests.test_serialize_outputs_sync__unit_test_0', None, 'test_data_output_mode_all_skips_images', 'serialize_outputs_sync', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    serialized, errors = codeflash_return_value

def test_serialize_datetime_and_nested_structures():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    now = datetime.datetime(2022, 1, 2, 3, 4, 5)
    nested = {'time': now, 'list': [1, now, {'inner_time': now}]}
    fields = ['n']
    output = {'n': nested}
    _call__bound__arguments = inspect.signature(VideoFrameProcessor.serialize_outputs_sync).bind(fields, output, DataOutputMode.NONE)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(VideoFrameProcessor.serialize_outputs_sync, 'tests.inference.unit_tests.test_serialize_outputs_sync__unit_test_0', None, 'test_serialize_datetime_and_nested_structures', 'serialize_outputs_sync', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    serialized, errors = codeflash_return_value

def test_large_scale_serialization_many_fields():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    count = 200
    fields = [f'f{i}' for i in range(count)]
    now = datetime.datetime(2020, 5, 17, 12, 0, 0)
    workflow_output = {}
    for i in range(count):
        if i % 50 == 0:
            workflow_output[f'f{i}'] = WorkflowImageData(numpy_image=f'IMG{i}', video_metadata=VideoMetadata(idx=i))
        elif i % 10 == 0:
            workflow_output[f'f{i}'] = {'nested': [i, now, {'deep': now}]}
        elif i % 2 == 0:
            workflow_output[f'f{i}'] = i
        else:
            workflow_output[f'f{i}'] = now
    _call__bound__arguments = inspect.signature(VideoFrameProcessor.serialize_outputs_sync).bind(fields, workflow_output, DataOutputMode.NONE)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(VideoFrameProcessor.serialize_outputs_sync, 'tests.inference.unit_tests.test_serialize_outputs_sync__unit_test_0', None, 'test_large_scale_serialization_many_fields', 'serialize_outputs_sync', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    serialized, errors = codeflash_return_value
    img_entry = serialized['f0']
    nested_entry = serialized['f10']