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.cache.base import BaseCache
from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v2 import \
    RoboflowDatasetUploadBlockV2
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 _make_batch(n):
    """
    Helper to create a Batch instance with `n` simple elements and corresponding indices.

    We use integers as content items because the RoboflowDatasetUploadBlockV2.run method
    only uses the batch elements as opaque items when the sampling branch is skipped.
    The Batch class is constructed via Batch.init (a real constructor from the codebase).
    """
    content = list(range(n))
    indices = [(i,) for i in range(n)]
    return Batch.init(content=content, indices=indices)

def _common_run(block, images, data_percentage=0.0, predictions=None, image_name=None, disable_sink=False):
    """
    Convenience wrapper to call block.run with consistent dummy parameters.
    It keeps all other arguments constant and focuses tests on behavior of interest.
    """
    return block.run(images=images, predictions=predictions, target_project='proj', usage_quota_name='quota', data_percentage=data_percentage, minutely_usage_limit=1, persist_predictions=False, hourly_usage_limit=10, daily_usage_limit=100, max_image_size=(1024, 1024), compression_level=50, registration_tags=[], disable_sink=disable_sink, fire_and_forget=False, labeling_batch_prefix='', labeling_batches_recreation_frequency='never', image_name=image_name)

def test_run_raises_when_no_api_key():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key=None, background_tasks=None, thread_pool_executor=None)
    images = _make_batch(1)
    with pytest.raises(ValueError) as excinfo:
        _common_run(block, images)

def test_disable_sink_returns_expected_message_and_length():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='valid-key', background_tasks=None, thread_pool_executor=None)
    images = _make_batch(3)
    result = _common_run(block, images, disable_sink=True)
    for entry in result:
        pass

def test_sampling_skips_registration_when_data_percentage_zero():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='valid-key', background_tasks=None, thread_pool_executor=None)
    images = _make_batch(4)
    result = _common_run(block, images, data_percentage=0.0)
    for entry in result:
        pass

def test_predictions_shorter_than_images_truncates_results_due_to_zip_behavior():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='valid-key', background_tasks=None, thread_pool_executor=None)
    images = _make_batch(5)
    predictions = Batch.init(content=[{'a': 1}, {'b': 2}], indices=[(0,), (1,)])
    result = _common_run(block, images, data_percentage=0.0, predictions=predictions)

def test_empty_images_batch_returns_empty_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='valid-key', background_tasks=None, thread_pool_executor=None)
    images = Batch.init(content=[], indices=[])
    result = _common_run(block, images, data_percentage=0.0)

def test_large_scale_sampling_skip_handles_1000_elements():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='valid-key', background_tasks=None, thread_pool_executor=None)
    n = 1000
    images = _make_batch(n)
    result = _common_run(block, images, data_percentage=0.0)

def test_large_scale_disable_sink_with_1000_elements():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='valid-key', background_tasks=None, thread_pool_executor=None)
    n = 1000
    images = _make_batch(n)
    result = _common_run(block, images, disable_sink=True)
    for entry in result:
        pass