from __future__ import annotations

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

import dill as pickle
import pytest
from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v2 import \
    maybe_register_datapoint_at_roboflow


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_skip_when_data_percentage_zero():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Basic test:\n    When data_percentage is 0, maybe_register_datapoint_at_roboflow must always\n    skip registration regardless of random state and return the expected message.\n    '
    _call__bound__arguments = inspect.signature(maybe_register_datapoint_at_roboflow).bind(image=None, prediction=None, target_project='project_x', usage_quota_name='quota', data_percentage=0.0, persist_predictions=False, minutely_usage_limit=1, hourly_usage_limit=1, daily_usage_limit=1, max_image_size=(100, 100), compression_level=75, registration_tags=[], fire_and_forget=False, labeling_batch_prefix='prefix', new_labeling_batch_frequency='never', cache=None, background_tasks=None, thread_pool_executor=None, api_key='api_key_value', image_name=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(maybe_register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_maybe_register_datapoint_at_roboflow__unit_test_0', None, 'test_skip_when_data_percentage_zero', 'maybe_register_datapoint_at_roboflow', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_fire_and_forget_with_thread_pool_submits_task_and_returns_immediate_message():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "\n    Basic test:\n    When data_percentage is 100 (always register) and fire_and_forget=True with a\n    provided ThreadPoolExecutor, the function should submit the registration task\n    to the executor and return the immediate 'background' message. We do not\n    wait for the background task to complete.\n    "
    random.seed(0)
    executor = ThreadPoolExecutor(max_workers=1)
    try:
        _call__bound__arguments = inspect.signature(maybe_register_datapoint_at_roboflow).bind(image=None, prediction=None, target_project='project_x', usage_quota_name='quota', data_percentage=100.0, persist_predictions=False, minutely_usage_limit=1, hourly_usage_limit=1, daily_usage_limit=1, max_image_size=(100, 100), compression_level=75, registration_tags=['tag1'], fire_and_forget=True, labeling_batch_prefix='prefix', new_labeling_batch_frequency='never', cache=None, background_tasks=None, thread_pool_executor=executor, api_key='api_key_value', image_name='img_123')
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(maybe_register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_maybe_register_datapoint_at_roboflow__unit_test_0', None, 'test_fire_and_forget_with_thread_pool_submits_task_and_returns_immediate_message', 'maybe_register_datapoint_at_roboflow', '3_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        executor.shutdown(wait=False)

def test_data_percentage_above_100_treated_as_full_probability_and_schedules_task():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Edge test:\n    A data_percentage > 100 should be treated as probability > 1 after normalization,\n    which makes random.random() < normalized_probability always True. The function\n    must therefore attempt registration. We again use fire_and_forget with a thread\n    pool so the call returns immediately and deterministically.\n    '
    executor = ThreadPoolExecutor(max_workers=1)
    try:
        _call__bound__arguments = inspect.signature(maybe_register_datapoint_at_roboflow).bind(image=None, prediction=None, target_project='project_special_chars!@#', usage_quota_name='quota-with-special:chars', data_percentage=150.0, persist_predictions=True, minutely_usage_limit=0, hourly_usage_limit=0, daily_usage_limit=0, max_image_size=(1, 1), compression_level=100, registration_tags=['weird tag', 'ümlaut'], fire_and_forget=True, labeling_batch_prefix='batch_prefix', new_labeling_batch_frequency='daily', cache=None, background_tasks=None, thread_pool_executor=executor, api_key='api_key_value', image_name=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(maybe_register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_maybe_register_datapoint_at_roboflow__unit_test_0', None, 'test_data_percentage_above_100_treated_as_full_probability_and_schedules_task', 'maybe_register_datapoint_at_roboflow', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        executor.shutdown(wait=False)

def test_negative_data_percentage_always_skips_registration():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Edge test:\n    Negative data_percentage should normalize to a negative probability, and thus\n    random.random() < normalized_probability is always False -> registration skipped.\n    '
    random.seed(12345)
    _call__bound__arguments = inspect.signature(maybe_register_datapoint_at_roboflow).bind(image=None, prediction=None, target_project='p', usage_quota_name='q', data_percentage=-10.0, persist_predictions=False, minutely_usage_limit=1, hourly_usage_limit=1, daily_usage_limit=1, max_image_size=(50, 50), compression_level=50, registration_tags=[], fire_and_forget=False, labeling_batch_prefix='', new_labeling_batch_frequency='never', cache=None, background_tasks=None, thread_pool_executor=None, api_key='k', image_name=None)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(maybe_register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_maybe_register_datapoint_at_roboflow__unit_test_0', None, 'test_negative_data_percentage_always_skips_registration', 'maybe_register_datapoint_at_roboflow', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_many_invocations_with_zero_percentage_are_all_skipped():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "\n    Large-scale test:\n    Call the function 1000 times with data_percentage = 0 and assert it always\n    returns the 'skipped due to sampling settings' result. This verifies that\n    the sampling branch is cheap and deterministic over many iterations.\n    "
    for _ in range(1000):
        _call__bound__arguments = inspect.signature(maybe_register_datapoint_at_roboflow).bind(image=None, prediction=None, target_project='bulk_test', usage_quota_name='bulk_quota', data_percentage=0.0, persist_predictions=False, minutely_usage_limit=1, hourly_usage_limit=1, daily_usage_limit=1, max_image_size=(10, 10), compression_level=50, registration_tags=[], fire_and_forget=False, labeling_batch_prefix='', new_labeling_batch_frequency='never', cache=None, background_tasks=None, thread_pool_executor=None, api_key='k', image_name=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(maybe_register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_maybe_register_datapoint_at_roboflow__unit_test_0', None, 'test_many_invocations_with_zero_percentage_are_all_skipped', 'maybe_register_datapoint_at_roboflow', '1_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output