from __future__ import annotations

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

import dill as pickle
import numpy as np
import pytest
from inference.core.workflows.execution_engine.entities.base import \
    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}######!')
    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_copy_preserves_fields_and_returns_new_instance():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    parent = object()
    origin = WorkflowImageData(parent, base64_image='orig_base64')
    _call__bound__arguments = inspect.signature(WorkflowImageData.copy_and_replace).bind(origin)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(WorkflowImageData.copy_and_replace, 'tests.inference.unit_tests.test_copy_and_replace__unit_test_0', None, 'test_basic_copy_preserves_fields_and_returns_new_instance', 'copy_and_replace', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    copied = codeflash_output

def test_copy_and_replace_replaces_image_fields_when_any_image_key_present():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    parent = object()
    origin = WorkflowImageData(parent, base64_image='keep_this_base64')
    new_array = np.array([1, 2, 3])
    _call__bound__arguments = inspect.signature(WorkflowImageData.copy_and_replace).bind(origin, numpy_image=new_array, base64_image=None, image_reference=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(WorkflowImageData.copy_and_replace, 'tests.inference.unit_tests.test_copy_and_replace__unit_test_0', None, 'test_copy_and_replace_replaces_image_fields_when_any_image_key_present', 'copy_and_replace', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    replaced = codeflash_output

def test_copy_and_replace_raises_if_all_image_representations_are_cleared():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    parent = object()
    origin = WorkflowImageData(parent, base64_image='present')
    with pytest.raises(ValueError):
        _call__bound__arguments = inspect.signature(WorkflowImageData.copy_and_replace).bind(origin, numpy_image=None, base64_image=None, image_reference=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(WorkflowImageData.copy_and_replace, 'tests.inference.unit_tests.test_copy_and_replace__unit_test_0', None, 'test_copy_and_replace_raises_if_all_image_representations_are_cleared', 'copy_and_replace', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_parent_and_workflow_root_swapping_when_they_are_identical():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_parent = object()
    origin = WorkflowImageData(original_parent, base64_image='b64')
    new_parent = object()
    _call__bound__arguments = inspect.signature(WorkflowImageData.copy_and_replace).bind(origin, parent_metadata=new_parent)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(WorkflowImageData.copy_and_replace, 'tests.inference.unit_tests.test_copy_and_replace__unit_test_0', None, 'test_parent_and_workflow_root_swapping_when_they_are_identical', 'copy_and_replace', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    replaced = codeflash_output

def test_parent_replacement_does_not_change_workflow_root_if_distinct():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    parent = object()
    distinct_root = object()
    origin = WorkflowImageData(parent, workflow_root_ancestor_metadata=distinct_root, base64_image='b64')
    new_parent = object()
    _call__bound__arguments = inspect.signature(WorkflowImageData.copy_and_replace).bind(origin, parent_metadata=new_parent)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(WorkflowImageData.copy_and_replace, 'tests.inference.unit_tests.test_copy_and_replace__unit_test_0', None, 'test_parent_replacement_does_not_change_workflow_root_if_distinct', 'copy_and_replace', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    replaced = codeflash_output

def test_workflow_root_replacement_updates_parent_when_they_are_identical():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_parent = object()
    origin = WorkflowImageData(original_parent, base64_image='b64')
    new_root = object()
    _call__bound__arguments = inspect.signature(WorkflowImageData.copy_and_replace).bind(origin, workflow_root_ancestor_metadata=new_root)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(WorkflowImageData.copy_and_replace, 'tests.inference.unit_tests.test_copy_and_replace__unit_test_0', None, 'test_workflow_root_replacement_updates_parent_when_they_are_identical', 'copy_and_replace', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    replaced = codeflash_output

def test_video_metadata_and_image_name_are_updated_if_provided():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    parent = object()
    origin = WorkflowImageData(parent, base64_image='b64')
    video_meta = {'fps': 30}
    _call__bound__arguments = inspect.signature(WorkflowImageData.copy_and_replace).bind(origin, video_metadata=video_meta, image_name='new_name')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(WorkflowImageData.copy_and_replace, 'tests.inference.unit_tests.test_copy_and_replace__unit_test_0', None, 'test_video_metadata_and_image_name_are_updated_if_provided', 'copy_and_replace', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    replaced = codeflash_output

def test_numpy_array_identity_and_mutation_reflected_in_copy_when_not_replaced():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    parent = object()
    arr = np.zeros((2, 2), dtype=int)
    origin = WorkflowImageData(parent, numpy_image=arr)
    _call__bound__arguments = inspect.signature(WorkflowImageData.copy_and_replace).bind(origin)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(WorkflowImageData.copy_and_replace, 'tests.inference.unit_tests.test_copy_and_replace__unit_test_0', None, 'test_numpy_array_identity_and_mutation_reflected_in_copy_when_not_replaced', 'copy_and_replace', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    copied = codeflash_output
    arr[0, 0] = 42

def test_constructor_raises_if_no_image_representation_provided():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    with pytest.raises(ValueError):
        WorkflowImageData(parent_metadata=object())
    with pytest.raises(ValueError):
        WorkflowImageData(parent_metadata=object(), numpy_image=None)

def test_large_scale_many_sequential_copies_update_image_name_correctly():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    parent = object()
    origin = WorkflowImageData(parent, base64_image='persistent_base64')
    current = origin
    for i in range(1000):
        _call__bound__arguments = inspect.signature(WorkflowImageData.copy_and_replace).bind(current, image_name=f'name_{i}')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(WorkflowImageData.copy_and_replace, 'tests.inference.unit_tests.test_copy_and_replace__unit_test_0', None, 'test_large_scale_many_sequential_copies_update_image_name_correctly', 'copy_and_replace', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        current = codeflash_output