from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path
from typing import List, Optional, Tuple

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, ImageParentMetadata, 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 make_workflow_image(id_suffix: str='0', base64_value: str='a') -> WorkflowImageData:
    """
    Helper to construct a minimal, valid WorkflowImageData instance.
    We provide a real ImageParentMetadata (from the real module) and a non-empty
    base64_image to satisfy constructor validation.
    """
    parent_metadata = ImageParentMetadata(parent_id=f'parent_{id_suffix}', origin_coordinates=None)
    return WorkflowImageData(parent_metadata=parent_metadata, base64_image=base64_value)

def make_batch_of_images(n: int) -> Batch[WorkflowImageData]:
    """
    Helper to create a Batch containing n WorkflowImageData objects with matching indices.
    Uses Batch.init which enforces index/content length equality.
    """
    images = [make_workflow_image(id_suffix=str(i)) for i in range(n)]
    indices = [(i,) for i in range(n)]
    return Batch.init(content=images, indices=indices)

def test_run_raises_when_api_key_missing():
    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_of_images(1)
    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', data_percentage=0.0, minutely_usage_limit=0, persist_predictions=False, hourly_usage_limit=0, daily_usage_limit=0, max_image_size=(100, 100), compression_level=5, registration_tags=[], disable_sink=False, fire_and_forget=False, labeling_batch_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_run_raises_when_api_key_missing', 'run', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_disable_sink_returns_disabled_messages_single_image():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='dummy-key', background_tasks=None, thread_pool_executor=None)
    images = make_batch_of_images(1)
    _call__bound__arguments = inspect.signature(block.run).bind(images=images, predictions=None, target_project='proj', usage_quota_name='quota', data_percentage=0.0, minutely_usage_limit=0, persist_predictions=False, hourly_usage_limit=0, daily_usage_limit=0, max_image_size=(100, 100), compression_level=5, registration_tags=['tag'], disable_sink=True, fire_and_forget=False, labeling_batch_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_disabled_messages_single_image', 'run', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_run_skips_registration_when_data_percentage_zero_multiple_images():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='dummy-key', background_tasks=None, thread_pool_executor=None)
    images = make_batch_of_images(3)
    _call__bound__arguments = inspect.signature(block.run).bind(images=images, predictions=None, target_project='proj', usage_quota_name='quota', data_percentage=0.0, minutely_usage_limit=0, persist_predictions=False, hourly_usage_limit=0, daily_usage_limit=0, max_image_size=(200, 200), compression_level=1, registration_tags=['a', 'b'], disable_sink=False, fire_and_forget=False, labeling_batch_prefix='lbl', 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_run_skips_registration_when_data_percentage_zero_multiple_images', '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_run_accepts_explicit_predictions_and_image_name_batch():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='key', background_tasks=None, thread_pool_executor=None)
    images = make_batch_of_images(2)
    predictions = [{'boxes': []}, {'boxes': []}]
    image_names = Batch.init(content=['img_a', 'img_b'], indices=[(0,), (1,)])
    _call__bound__arguments = inspect.signature(block.run).bind(images=images, predictions=predictions, target_project='proj', usage_quota_name='q', data_percentage=0.0, minutely_usage_limit=1, persist_predictions=True, hourly_usage_limit=10, daily_usage_limit=100, max_image_size=(50, 50), compression_level=3, registration_tags=['t'], disable_sink=False, fire_and_forget=False, labeling_batch_prefix='p', labeling_batches_recreation_frequency='never', image_name=image_names)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_run_accepts_explicit_predictions_and_image_name_batch', 'run', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_negative_data_percentage_treated_as_zero_sampling():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='key', background_tasks=None, thread_pool_executor=None)
    images = make_batch_of_images(2)
    _call__bound__arguments = inspect.signature(block.run).bind(images=images, predictions=None, target_project='proj', usage_quota_name='q', data_percentage=-10.0, minutely_usage_limit=0, persist_predictions=False, hourly_usage_limit=0, daily_usage_limit=0, max_image_size=(10, 10), compression_level=0, registration_tags=[], disable_sink=False, fire_and_forget=False, labeling_batch_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_negative_data_percentage_treated_as_zero_sampling', 'run', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_batch_of_1000_images_is_processed_quickly_and_deterministically():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    cache = BaseCache()
    block = RoboflowDatasetUploadBlockV2(cache=cache, api_key='large-test-key', background_tasks=None, thread_pool_executor=None)
    n = 1000
    images = make_batch_of_images(n)
    _call__bound__arguments = inspect.signature(block.run).bind(images=images, predictions=None, target_project='proj-large', usage_quota_name='quota-large', data_percentage=0.0, minutely_usage_limit=0, persist_predictions=False, hourly_usage_limit=0, daily_usage_limit=0, max_image_size=(64, 64), compression_level=1, registration_tags=['bulk'], disable_sink=False, fire_and_forget=False, labeling_batch_prefix='batch', 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_large_scale_batch_of_1000_images_is_processed_quickly_and_deterministically', 'run', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for idx in (0, n // 2, n - 1):
        pass