from __future__ import annotations

import gc
import inspect
import os
import time
import types
from pathlib import Path
from typing import Any, Dict

import dill as pickle
import inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1 as v1
import pytest
from inference.core.roboflow_api import (annotate_image_at_roboflow,
                                         register_image_at_roboflow)
from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1 import \
    register_datapoint
from inference.core.workflows.execution_engine.constants import \
    INFERENCE_ID_KEY


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_stubs_capture(registration_response: Dict[str, Any], annotate_response: Dict[str, Any]=None):
    """
    Helper that returns a pair of stub functions (fake_register, fake_annotate)
    that capture the arguments they were called with into a dict
    and return predictable responses.
    """
    record: Dict[str, Any] = {}

    def fake_register_image_at_roboflow(api_key: str, dataset_id: str, local_image_id: str, image_bytes: bytes, batch_name: str, tags=None, inference_id=None, metadata=None):
        record['register_called'] = record.get('register_called', 0) + 1
        record['last_register_args'] = {'api_key': api_key, 'dataset_id': dataset_id, 'local_image_id': local_image_id, 'image_bytes': image_bytes, 'batch_name': batch_name, 'tags': list(tags) if tags is not None else [], 'inference_id': inference_id, 'metadata': metadata}
        return registration_response

    def fake_annotate_image_at_roboflow(api_key: str, dataset_id: str, local_image_id: str, roboflow_image_id: str, annotation_content: str, annotation_file_type: str, is_prediction: bool=True):
        record['annotate_called'] = record.get('annotate_called', 0) + 1
        record['last_annotate_args'] = {'api_key': api_key, 'dataset_id': dataset_id, 'local_image_id': local_image_id, 'roboflow_image_id': roboflow_image_id, 'annotation_content': annotation_content, 'annotation_file_type': annotation_file_type, 'is_prediction': is_prediction}
        return annotate_response or {}
    return (fake_register_image_at_roboflow, fake_annotate_image_at_roboflow, record)

def test_duplicated_image_returns_duplicated_status():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_register = v1.register_image_at_roboflow
    original_annotate = v1.annotate_image_at_roboflow
    try:
        fake_reg, fake_ann, record = _make_stubs_capture({'duplicate': True, 'id': 'ignored'})
        v1.register_image_at_roboflow = fake_reg
        v1.annotate_image_at_roboflow = fake_ann
        _call__bound__arguments = inspect.signature(v1.register_datapoint).bind(target_project='proj', encoded_image=b'\x00\x01', local_image_id='local1', prediction=None, api_key='key', batch_name='batchA', tags=['t1', 't2'], metadata=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(v1.register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_0', None, 'test_duplicated_image_returns_duplicated_status', 'register_datapoint', '2_3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        v1.register_image_at_roboflow = original_register
        v1.annotate_image_at_roboflow = original_annotate

def test_successful_registration_without_prediction_skips_annotation():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_register = v1.register_image_at_roboflow
    original_annotate = v1.annotate_image_at_roboflow
    try:
        fake_reg, fake_ann, record = _make_stubs_capture({'duplicate': False, 'id': 'rf-id-123'})
        v1.register_image_at_roboflow = fake_reg
        v1.annotate_image_at_roboflow = fake_ann
        _call__bound__arguments = inspect.signature(v1.register_datapoint).bind(target_project='projX', encoded_image=b'\xff', local_image_id='img42', prediction=None, api_key='apikey', batch_name='batch1', tags=[], metadata={'meta': 'value'})
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(v1.register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_0', None, 'test_successful_registration_without_prediction_skips_annotation', 'register_datapoint', '2_3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        v1.register_image_at_roboflow = original_register
        v1.annotate_image_at_roboflow = original_annotate

def test_annotation_with_predicted_classes_uses_comma_separated_txt_format():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_register = v1.register_image_at_roboflow
    original_annotate = v1.annotate_image_at_roboflow
    try:
        fake_reg, fake_ann, record = _make_stubs_capture({'duplicate': False, 'id': 'rf-annot-1'})
        v1.register_image_at_roboflow = fake_reg
        v1.annotate_image_at_roboflow = fake_ann
        prediction = {'predicted_classes': ['cat', 'dog']}
        _call__bound__arguments = inspect.signature(v1.register_datapoint).bind(target_project='animals', encoded_image=b'imgbytes', local_image_id='img-cd', prediction=prediction, api_key='k', batch_name='b', tags=['pets'], metadata=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(v1.register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_0', None, 'test_annotation_with_predicted_classes_uses_comma_separated_txt_format', 'register_datapoint', '2_4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
        ann_args = record['last_annotate_args']
    finally:
        v1.register_image_at_roboflow = original_register
        v1.annotate_image_at_roboflow = original_annotate

def test_annotation_with_top_key_uses_top_value_txt_format():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_register = v1.register_image_at_roboflow
    original_annotate = v1.annotate_image_at_roboflow
    try:
        fake_reg, fake_ann, record = _make_stubs_capture({'duplicate': False, 'id': 'rf-top-1'})
        v1.register_image_at_roboflow = fake_reg
        v1.annotate_image_at_roboflow = fake_ann
        prediction = {'top': 'best-class-only'}
        _call__bound__arguments = inspect.signature(v1.register_datapoint).bind(target_project='projTop', encoded_image=b'', local_image_id='top1', prediction=prediction, api_key='keytop', batch_name='batchTop', tags=[], metadata=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(v1.register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_0', None, 'test_annotation_with_top_key_uses_top_value_txt_format', 'register_datapoint', '2_4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
        ann_args = record['last_annotate_args']
    finally:
        v1.register_image_at_roboflow = original_register
        v1.annotate_image_at_roboflow = original_annotate

def test_prediction_dict_missing_required_keys_forbids_annotation():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_register = v1.register_image_at_roboflow
    original_annotate = v1.annotate_image_at_roboflow
    try:
        fake_reg, fake_ann, record = _make_stubs_capture({'duplicate': False, 'id': 'rf-no-annot'})
        v1.register_image_at_roboflow = fake_reg
        v1.annotate_image_at_roboflow = fake_ann
        prediction = {'some_other_key': 123}
        _call__bound__arguments = inspect.signature(v1.register_datapoint).bind(target_project='projNoAnn', encoded_image=b'123', local_image_id='noann', prediction=prediction, api_key='k', batch_name='b', tags=['x'], metadata=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(v1.register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_0', None, 'test_prediction_dict_missing_required_keys_forbids_annotation', 'register_datapoint', '2_4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        v1.register_image_at_roboflow = original_register
        v1.annotate_image_at_roboflow = original_annotate

def test_inference_id_from_prediction_dict_is_passed_to_register():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_register = v1.register_image_at_roboflow
    original_annotate = v1.annotate_image_at_roboflow
    try:
        fake_reg, fake_ann, record = _make_stubs_capture({'duplicate': False, 'id': 'rf-id-forward'})
        v1.register_image_at_roboflow = fake_reg
        v1.annotate_image_at_roboflow = fake_ann
        inference_key = v1.INFERENCE_ID_KEY
        prediction = {inference_key: 'inference-XYZ', 'predicted_classes': ['a']}
        _call__bound__arguments = inspect.signature(v1.register_datapoint).bind(target_project='projInfer', encoded_image=b'i', local_image_id='in1', prediction=prediction, api_key='k', batch_name='b', tags=[], metadata=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(v1.register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_0', None, 'test_inference_id_from_prediction_dict_is_passed_to_register', 'register_datapoint', '2_5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        _ = codeflash_output
    finally:
        v1.register_image_at_roboflow = original_register
        v1.annotate_image_at_roboflow = original_annotate

def test_large_tags_list_is_transmitted_to_register_and_handled_correctly():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_register = v1.register_image_at_roboflow
    original_annotate = v1.annotate_image_at_roboflow
    try:
        large_tags = [f'tag{i}' for i in range(1000)]
        fake_reg, fake_ann, record = _make_stubs_capture({'duplicate': False, 'id': 'rf-large-tags'})
        v1.register_image_at_roboflow = fake_reg
        v1.annotate_image_at_roboflow = fake_ann
        _call__bound__arguments = inspect.signature(v1.register_datapoint).bind(target_project='projLargeTags', encoded_image=b'bytes', local_image_id='large1', prediction={'predicted_classes': ['only']}, api_key='key', batch_name='batchLT', tags=large_tags, metadata=None)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(v1.register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_0', None, 'test_large_tags_list_is_transmitted_to_register_and_handled_correctly', 'register_datapoint', '2_4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        v1.register_image_at_roboflow = original_register
        v1.annotate_image_at_roboflow = original_annotate

def test_large_scale_repeated_registrations_are_deterministic_and_performant():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_register = v1.register_image_at_roboflow
    original_annotate = v1.annotate_image_at_roboflow
    try:
        fake_reg, fake_ann, record = _make_stubs_capture({'duplicate': False, 'id': 'rf-batch'})
        v1.register_image_at_roboflow = fake_reg
        v1.annotate_image_at_roboflow = fake_ann
        iterations = 1000
        for i in range(iterations):
            if i % 2 == 0:
                pred = {'predicted_classes': [f'cls{i}']}
            else:
                pred = {'some_other': i}
            _call__bound__arguments = inspect.signature(v1.register_datapoint).bind(target_project='projBulk', encoded_image=bytes([i % 256]), local_image_id=f'img{i}', prediction=pred, api_key='bulkkey', batch_name='bulk', tags=['bulk'], metadata=None)
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(v1.register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_0', None, 'test_large_scale_repeated_registrations_are_deterministic_and_performant', 'register_datapoint', '2_4_1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
            codeflash_output = codeflash_return_value
            res = codeflash_output
    finally:
        v1.register_image_at_roboflow = original_register
        v1.annotate_image_at_roboflow = original_annotate