from __future__ import annotations

import gc
import os
import time
import types
from concurrent.futures import ThreadPoolExecutor

import pytest
from fastapi import BackgroundTasks
from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload import \
    v1 as rf
from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1 import \
    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}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

def test_background_tasks_fire_and_forget_does_not_execute_registration(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Basic check:\n    - When fire_and_forget is True and a FastAPI BackgroundTasks instance is provided,\n      register_datapoint_at_roboflow should schedule the registration and return the background message.\n    - The execute_registration function must NOT be executed synchronously.\n    '
    executed = {'flag': False}

    def fake_execute_registration(*args, **kwargs):
        executed['flag'] = True
        return (True, 'should not be returned')
    monkeypatch.setattr(rf, 'execute_registration', fake_execute_registration)
    bg_tasks = BackgroundTasks()
    codeflash_return_value = codeflash_wrap(rf.register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_register_datapoint_at_roboflow__unit_test_0', None, 'test_background_tasks_fire_and_forget_does_not_execute_registration', 'register_datapoint_at_roboflow', '5', codeflash_loop_index, image=None, prediction=None, target_project='proj', usage_quota_name='q', 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=True, labeling_batch_prefix='p', new_labeling_batch_frequency='never', cache=None, background_tasks=bg_tasks, thread_pool_executor=None, api_key='key', image_name=None, metadata=None)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    scheduled = bg_tasks.tasks[0]

def test_synchronous_execution_calls_execute_registration_and_returns_its_result(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Basic synchronous flow:\n    - When fire_and_forget is False and no executors/background tasks provided,\n      the function should call execute_registration synchronously and return its result.\n    '
    captured = {'called': False, 'kwargs': None}

    def fake_execute_registration(image, prediction, target_project, persist_predictions, usage_quota_name, minutely_usage_limit, hourly_usage_limit, daily_usage_limit, max_image_size, compression_level, registration_tags, labeling_batch_prefix, new_labeling_batch_frequency, cache, api_key, image_name=None, metadata=None):
        captured['called'] = True
        captured['kwargs'] = {'target_project': target_project, 'persist_predictions': persist_predictions, 'usage_quota_name': usage_quota_name, 'api_key': api_key}
        return (True, 'sync-complete')
    monkeypatch.setattr(rf, 'execute_registration', fake_execute_registration)
    codeflash_return_value = codeflash_wrap(rf.register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_register_datapoint_at_roboflow__unit_test_0', None, 'test_synchronous_execution_calls_execute_registration_and_returns_its_result', 'register_datapoint_at_roboflow', '4', codeflash_loop_index, image=None, prediction=None, target_project='my_project', usage_quota_name='quota_xyz', persist_predictions=True, minutely_usage_limit=2, hourly_usage_limit=3, daily_usage_limit=4, max_image_size=(50, 50), compression_level=80, registration_tags=['t1'], fire_and_forget=False, labeling_batch_prefix='batch', new_labeling_batch_frequency='never', cache=None, background_tasks=None, thread_pool_executor=None, api_key='super_key', image_name='img_123', metadata={'k': 'v'})
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_fire_and_forget_prefers_background_tasks_over_thread_pool(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    If both background_tasks and a thread pool executor are provided, the function should prefer\n    adding a background task (background_tasks path) and NOT submit to the executor.\n    '
    executor_calls = {'count': 0}
    executor = ThreadPoolExecutor(max_workers=1)

    def recorder_submit(arg):
        executor_calls['count'] += 1
        return None
    executor.submit = recorder_submit

    def raising_execute_registration(*a, **kw):
        raise RuntimeError('Should not be executed synchronously or by the executor in this test')
    monkeypatch.setattr(rf, 'execute_registration', raising_execute_registration)
    bg_tasks = BackgroundTasks()
    codeflash_return_value = codeflash_wrap(rf.register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_register_datapoint_at_roboflow__unit_test_0', None, 'test_fire_and_forget_prefers_background_tasks_over_thread_pool', 'register_datapoint_at_roboflow', '8', codeflash_loop_index, image=None, prediction=None, target_project='tp', usage_quota_name='q', persist_predictions=False, minutely_usage_limit=0, hourly_usage_limit=0, daily_usage_limit=0, max_image_size=(10, 10), compression_level=50, registration_tags=[], fire_and_forget=True, labeling_batch_prefix='pfx', new_labeling_batch_frequency='never', cache=None, background_tasks=bg_tasks, thread_pool_executor=executor, api_key='api', image_name=None, metadata=None)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_thread_pool_executor_used_when_no_background_tasks(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    When no background_tasks is provided but a thread pool executor is,\n    the function should call executor.submit with the registration task and return the background message.\n    We monkeypatch executor.submit to record the passed callable without executing it.\n    '
    submitted = {'callable': None}
    executor = ThreadPoolExecutor(max_workers=1)

    def fake_submit(fn):
        submitted['callable'] = fn
        return None
    executor.submit = fake_submit

    def raising_execute_registration(*a, **kw):
        raise RuntimeError('execute_registration should not run when fire_and_forget uses executor in test')
    monkeypatch.setattr(rf, 'execute_registration', raising_execute_registration)
    codeflash_return_value = codeflash_wrap(rf.register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_register_datapoint_at_roboflow__unit_test_0', None, 'test_thread_pool_executor_used_when_no_background_tasks', 'register_datapoint_at_roboflow', '7', codeflash_loop_index, image=None, prediction=None, target_project='tp2', usage_quota_name='q2', persist_predictions=False, minutely_usage_limit=0, hourly_usage_limit=0, daily_usage_limit=0, max_image_size=(5, 5), compression_level=50, registration_tags=[], fire_and_forget=True, labeling_batch_prefix='b', new_labeling_batch_frequency='never', cache=None, background_tasks=None, thread_pool_executor=executor, api_key='k', image_name=None, metadata=None)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_synchronous_invocations(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Call register_datapoint_at_roboflow synchronously many times (1000) with a safe stub for execute_registration.\n    This ensures the function behaves deterministically under many repeated synchronous calls.\n    '
    calls = {'count': 0}

    def fake_execute_registration(*args, **kwargs):
        calls['count'] += 1
        return (False, f"ok-{calls['count']}")
    monkeypatch.setattr(rf, 'execute_registration', fake_execute_registration)
    N = 1000
    for i in range(N):
        codeflash_return_value = codeflash_wrap(rf.register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_register_datapoint_at_roboflow__unit_test_0', None, 'test_large_scale_synchronous_invocations', 'register_datapoint_at_roboflow', '5_0', codeflash_loop_index, image=None, prediction=None, target_project=f'p{i}', usage_quota_name='quota', persist_predictions=False, minutely_usage_limit=1, hourly_usage_limit=1, daily_usage_limit=1, max_image_size=(200, 200), compression_level=85, 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=f'key{i}', image_name=None, metadata=None)
        codeflash_output = codeflash_return_value
        result = codeflash_output

def test_large_scale_background_tasks_scheduling():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Schedule a large number of background tasks (1000) and ensure they are all collected.\n    We avoid executing registration by not providing a real execute_registration invocation.\n    '
    bg = BackgroundTasks()
    N = 1000
    for i in range(N):
        codeflash_return_value = codeflash_wrap(rf.register_datapoint_at_roboflow, 'tests.inference.unit_tests.test_register_datapoint_at_roboflow__unit_test_0', None, 'test_large_scale_background_tasks_scheduling', 'register_datapoint_at_roboflow', '3_0', codeflash_loop_index, image=None, prediction=None, target_project=f'proj-{i}', usage_quota_name='q', persist_predictions=False, minutely_usage_limit=1, hourly_usage_limit=1, daily_usage_limit=1, max_image_size=(64, 64), compression_level=70, registration_tags=[], fire_and_forget=True, labeling_batch_prefix='bp', new_labeling_batch_frequency='never', cache=None, background_tasks=bg, thread_pool_executor=None, api_key=f'apikey-{i}', image_name=None, metadata=None)
        codeflash_output = codeflash_return_value
        result = codeflash_output