from __future__ import annotations

import gc
import os
import time

import numpy as np
import pytest
from inference.core.workflows.core_steps.common.deserializers import \
    deserialize_image_kind
from inference.core.workflows.errors import RuntimeInputError
from inference.core.workflows.execution_engine.constants import (
    IMAGE_NAME_KEY, PARENT_ID_KEY, PARENT_ORIGIN_KEY, ROOT_PARENT_ID_KEY,
    ROOT_PARENT_ORIGIN_KEY)
from inference.core.workflows.execution_engine.entities.base import (
    ImageParentMetadata, ParentOrigin, VideoMetadata, WorkflowImageData)


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}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

def test_pass_through_when_workflow_image_data_instance():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    parent_meta = ImageParentMetadata(parent_id='parent-123')
    arr = np.zeros((2, 2, 3), dtype=np.uint8)
    w = WorkflowImageData(parent_metadata=parent_meta, numpy_image=arr)
    codeflash_return_value = codeflash_wrap(deserialize_image_kind, 'tests.inference.unit_tests.test_deserialize_image_kind__unit_test_0', None, 'test_pass_through_when_workflow_image_data_instance', 'deserialize_image_kind', '3', codeflash_loop_index, parameter='p', image=w)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_numpy_array_input_creates_workflow_image_data_with_parent_metadata():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    arr = np.ones((3, 3, 3), dtype=np.uint8)
    param_name = 'my_param'
    codeflash_return_value = codeflash_wrap(deserialize_image_kind, 'tests.inference.unit_tests.test_deserialize_image_kind__unit_test_0', None, 'test_numpy_array_input_creates_workflow_image_data_with_parent_metadata', 'deserialize_image_kind', '2', codeflash_loop_index, parameter=param_name, image=arr)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_dict_with_value_numpy_array_is_handled_and_image_name_preserved():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    arr = np.full((1, 1, 3), 255, dtype=np.uint8)
    image_name = 'pic.png'
    image_dict = {PARENT_ID_KEY: 'explicit-parent', IMAGE_NAME_KEY: image_name, 'value': arr}
    codeflash_return_value = codeflash_wrap(deserialize_image_kind, 'tests.inference.unit_tests.test_deserialize_image_kind__unit_test_0', None, 'test_dict_with_value_numpy_array_is_handled_and_image_name_preserved', 'deserialize_image_kind', '3', codeflash_loop_index, parameter='paramX', image=image_dict)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_non_supported_type_raises_runtime_input_error():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    with pytest.raises(RuntimeInputError) as excinfo:
        codeflash_return_value = codeflash_wrap(deserialize_image_kind, 'tests.inference.unit_tests.test_deserialize_image_kind__unit_test_0', None, 'test_non_supported_type_raises_runtime_input_error', 'deserialize_image_kind', '0_0', codeflash_loop_index, parameter='p', image=12345)
    err = excinfo.value

def test_malformed_parent_origin_dict_raises_runtime_input_error_with_inner_validation_error():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    bad_parent_origin = {'offset_x': 'not-an-int'}
    image_dict = {PARENT_ID_KEY: 'pid', PARENT_ORIGIN_KEY: bad_parent_origin, 'value': np.zeros((1, 1, 3), dtype=np.uint8)}
    with pytest.raises(RuntimeInputError) as excinfo:
        codeflash_return_value = codeflash_wrap(deserialize_image_kind, 'tests.inference.unit_tests.test_deserialize_image_kind__unit_test_0', None, 'test_malformed_parent_origin_dict_raises_runtime_input_error_with_inner_validation_error', 'deserialize_image_kind', '2_0', codeflash_loop_index, parameter='paramY', image=image_dict)
    exc = excinfo.value
    inner_type = getattr(exc, 'inner_error_type', None)

def test_dict_with_non_image_value_string_triggers_runtime_input_error_when_loading_fails():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    bad_value = 'this-is-not-a-base64-or-pickle-or-path'
    image_dict = {PARENT_ID_KEY: 'pid2', 'value': bad_value}
    with pytest.raises(RuntimeInputError) as excinfo:
        codeflash_return_value = codeflash_wrap(deserialize_image_kind, 'tests.inference.unit_tests.test_deserialize_image_kind__unit_test_0', None, 'test_dict_with_non_image_value_string_triggers_runtime_input_error_when_loading_fails', 'deserialize_image_kind', '2_0', codeflash_loop_index, parameter='paramZ', image=image_dict, prevent_local_images_loading=True)
    exc = excinfo.value

def test_large_number_of_numpy_inputs_quick_loop():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    for i in range(n):
        a = np.full((2, 2, 3), i % 256, dtype=np.uint8)
        codeflash_return_value = codeflash_wrap(deserialize_image_kind, 'tests.inference.unit_tests.test_deserialize_image_kind__unit_test_0', None, 'test_large_number_of_numpy_inputs_quick_loop', 'deserialize_image_kind', '1_1', codeflash_loop_index, parameter=f'param_{i}', image=a)
        codeflash_output = codeflash_return_value
        res = codeflash_output

def test_batch_of_dicts_with_numpy_values_processed_in_loop():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    count = 500
    dicts = []
    for i in range(count):
        dicts.append({PARENT_ID_KEY: f'parent_{i}', 'value': np.zeros((1, 1, 3), dtype=np.uint8)})
    for i, d in enumerate(dicts):
        codeflash_return_value = codeflash_wrap(deserialize_image_kind, 'tests.inference.unit_tests.test_deserialize_image_kind__unit_test_0', None, 'test_batch_of_dicts_with_numpy_values_processed_in_loop', 'deserialize_image_kind', '3_0', codeflash_loop_index, parameter='batch', image=d)
        codeflash_output = codeflash_return_value
        out = codeflash_output