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.execution_engine.entities.base import \
    WorkflowImageData
from inference.enterprise.workflows.enterprise_blocks.sinks.event_writer import \
    v1 as event_writer_v1
from inference.enterprise.workflows.enterprise_blocks.sinks.event_writer.v1 import \
    _build_image_entry


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
_build_image_entry = event_writer_v1._build_image_entry

def test_basic_output_only_includes_base64Image():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    out = WorkflowImageData(parent_metadata=None, base64_image='OUT_BASE64')
    _call__bound__arguments = inspect.signature(_build_image_entry).bind(output_image=out)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_build_image_entry, 'tests.inference.unit_tests.test__build_image_entry__unit_test_0', None, 'test_basic_output_only_includes_base64Image', '_build_image_entry', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    entry = codeflash_output

def test_with_input_and_label_includes_all_fields():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    out = WorkflowImageData(parent_metadata=None, base64_image='OUT_B64')
    inp = WorkflowImageData(parent_metadata=None, base64_image='IN_B64')
    _call__bound__arguments = inspect.signature(_build_image_entry).bind(output_image=out, input_image=inp, image_label='cat')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_build_image_entry, 'tests.inference.unit_tests.test__build_image_entry__unit_test_0', None, 'test_with_input_and_label_includes_all_fields', '_build_image_entry', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    entry = codeflash_output
    for key in ('objectDetections', 'classifications', 'instanceSegmentations', 'keypoints'):
        pass

def test_classifications_from_dict_are_converted_and_floatified():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    out = WorkflowImageData(parent_metadata=None, base64_image='OUT')
    classifications = {'class_name': ['dog', 'cat', 'rabbit'], 'confidence': [0.9, '0.75']}
    _call__bound__arguments = inspect.signature(_build_image_entry).bind(output_image=out, classifications=classifications)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_build_image_entry, 'tests.inference.unit_tests.test__build_image_entry__unit_test_0', None, 'test_classifications_from_dict_are_converted_and_floatified', '_build_image_entry', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    entry = codeflash_output
    cls = entry['classifications']

def test_classifications_empty_dict_omitted_from_entry():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    out = WorkflowImageData(parent_metadata=None, base64_image='B64')
    _call__bound__arguments = inspect.signature(_build_image_entry).bind(output_image=out, classifications={})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_build_image_entry, 'tests.inference.unit_tests.test__build_image_entry__unit_test_0', None, 'test_classifications_empty_dict_omitted_from_entry', '_build_image_entry', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    entry = codeflash_output

def test_invalid_WorkflowImageData_construction_raises_ValueError():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    with pytest.raises(ValueError):
        WorkflowImageData(parent_metadata=None)

def test_classification_confidences_with_shorter_list_default_to_zero():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    out = WorkflowImageData(parent_metadata=None, base64_image='O')
    classifications = {'class_name': ['a', 'b', 'c', 'd'], 'confidence': [1]}
    _call__bound__arguments = inspect.signature(_build_image_entry).bind(output_image=out, classifications=classifications)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_build_image_entry, 'tests.inference.unit_tests.test__build_image_entry__unit_test_0', None, 'test_classification_confidences_with_shorter_list_default_to_zero', '_build_image_entry', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    entry = codeflash_output
    cls = entry['classifications']

def test_large_scale_classifications_performance_and_correctness():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    out = WorkflowImageData(parent_metadata=None, base64_image='LARGE_OUT')
    N = 1000
    class_names = [f'class_{i}' for i in range(N)]
    confidences = [i / N for i in range(N)]
    classifications = {'class_name': class_names, 'confidence': confidences}
    _call__bound__arguments = inspect.signature(_build_image_entry).bind(output_image=out, classifications=classifications)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_build_image_entry, 'tests.inference.unit_tests.test__build_image_entry__unit_test_0', None, 'test_large_scale_classifications_performance_and_correctness', '_build_image_entry', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    entry = codeflash_output
    cls = entry['classifications']
    for idx in (0, 10, 500, 999):
        pass

def test_entry_omits_detection_keys_when_converters_return_empty_lists():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    out = WorkflowImageData(parent_metadata=None, base64_image='OUT2')
    _call__bound__arguments = inspect.signature(_build_image_entry).bind(output_image=out, object_detections=None, classifications=None, instance_segmentations=None, keypoint_detections=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(_build_image_entry, 'tests.inference.unit_tests.test__build_image_entry__unit_test_0', None, 'test_entry_omits_detection_keys_when_converters_return_empty_lists', '_build_image_entry', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    entry = codeflash_output