from __future__ import annotations

import base64
import datetime
import gc
import os
import sys
import time
import types
from enum import Enum

import pytest
from inference.core.interfaces.webrtc_worker.webrtc import VideoFrameProcessor

def codeflash_wrap(
    codeflash_wrapped: Callable[..., Any],
    codeflash_test_module_name: str,
    codeflash_test_class_name: Optional[str],
    codeflash_test_name: str,
    codeflash_function_name: str,
    codeflash_line_id: str,
    codeflash_loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> str:
    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}"
    exception = None
    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}######$!")
    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}:{codeflash_duration}######!")
    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'}
    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, fields, output, DataOutputMode.NONE)
    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'}
    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, fields, output, DataOutputMode.NONE)
    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}
    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, fields, output, DataOutputMode.NONE)
    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}
    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, ['img'], output_err, DataOutputMode.NONE)
    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}
    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, fields, output, DataOutputMode.ALL)
    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}
    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, fields, output, DataOutputMode.NONE)
    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
    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, fields, workflow_output, DataOutputMode.NONE)
    serialized, errors = codeflash_return_value
    img_entry = serialized['f0']
    nested_entry = serialized['f10']