from __future__ import annotations

import gc
import inspect
import os
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import dill as pickle
import pytest
from inference.core.cache.base import BaseCache
from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1 import \
    RoboflowDatasetUploadBlockV1
from inference.core.workflows.execution_engine.entities.base import Batch


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_run_raises_when_no_api_key():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV1(cache=cache, api_key=None, background_tasks=None, thread_pool_executor=None)
    images = Batch(content=[object()], indices=[(0,)])
    with pytest.raises(ValueError) as excinfo:
        _call__bound__arguments = inspect.signature(block.run).bind(images=images, predictions=None, target_project='proj', usage_quota_name='quota', minutely_usage_limit=1, persist_predictions=False, hourly_usage_limit=10, daily_usage_limit=100, max_image_size=(640, 480), compression_level=75, registration_tags=[], disable_sink=False, fire_and_forget=False, labeling_batch_prefix='prefix', labeling_batches_recreation_frequency='daily', image_name=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_raises_when_no_api_key', 'run', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_disable_sink_returns_message_for_each_image():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV1(cache=cache, api_key='dummy-key', background_tasks=None, thread_pool_executor=None)
    images = Batch(content=[object(), object(), object()], indices=[(0,), (1,), (2,)])
    _call__bound__arguments = inspect.signature(block.run).bind(images=images, predictions=None, target_project='proj', usage_quota_name='quota', minutely_usage_limit=1, persist_predictions=False, hourly_usage_limit=10, daily_usage_limit=100, max_image_size=(800, 600), compression_level=85, registration_tags=['a', 'b'], disable_sink=True, fire_and_forget=False, labeling_batch_prefix='prefix', labeling_batches_recreation_frequency='never', image_name=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_disable_sink_returns_message_for_each_image', 'run', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for entry in result:
        pass

def test_fire_and_forget_with_threadpool_returns_background_message():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    with ThreadPoolExecutor(max_workers=2) as executor:
        block = RoboflowDatasetUploadBlockV1(cache=cache, api_key='abc-123', background_tasks=None, thread_pool_executor=executor)
        images = Batch(content=[object(), object()], indices=[(0,), (1,)])
        _call__bound__arguments = inspect.signature(block.run).bind(images=images, predictions=None, target_project='proj', usage_quota_name='quota', minutely_usage_limit=5, persist_predictions=True, hourly_usage_limit=50, daily_usage_limit=500, max_image_size=(1024, 768), compression_level=90, registration_tags=[], disable_sink=False, fire_and_forget=True, labeling_batch_prefix='lbl', labeling_batches_recreation_frequency='daily', image_name=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_fire_and_forget_with_threadpool_returns_background_message', 'run', '1_2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    for entry in result:
        pass

def test_predictions_shorter_than_images_truncates_results():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    with ThreadPoolExecutor(max_workers=2) as executor:
        block = RoboflowDatasetUploadBlockV1(cache=cache, api_key='key', background_tasks=None, thread_pool_executor=executor)
        images = Batch(content=[object(), object(), object()], indices=[(0,), (1,), (2,)])
        predictions = [None, None]
        _call__bound__arguments = inspect.signature(block.run).bind(images=images, predictions=predictions, target_project='proj', usage_quota_name='quota', minutely_usage_limit=1, persist_predictions=False, hourly_usage_limit=10, daily_usage_limit=100, max_image_size=(640, 480), compression_level=75, registration_tags=[], disable_sink=False, fire_and_forget=True, labeling_batch_prefix='prefix', labeling_batches_recreation_frequency='weekly', image_name=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_predictions_shorter_than_images_truncates_results', 'run', '1_3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    for entry in result:
        pass

def test_large_scale_fire_and_forget_with_many_images():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    n = 1000
    images = Batch(content=[object() for _ in range(n)], indices=[(i,) for i in range(n)])
    with ThreadPoolExecutor(max_workers=20) as executor:
        block = RoboflowDatasetUploadBlockV1(cache=cache, api_key='big-test-key', background_tasks=None, thread_pool_executor=executor)
        _call__bound__arguments = inspect.signature(block.run).bind(images=images, predictions=None, target_project='proj', usage_quota_name='quota', minutely_usage_limit=1, persist_predictions=False, hourly_usage_limit=100, daily_usage_limit=1000, max_image_size=(1280, 720), compression_level=80, registration_tags=['scale-test'], disable_sink=False, fire_and_forget=True, labeling_batch_prefix='batch', labeling_batches_recreation_frequency='monthly', image_name=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_large_scale_fire_and_forget_with_many_images', 'run', '3_1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    for entry in result:
        pass