from __future__ import annotations

import gc
import json
import os
import time
from unittest.mock import MagicMock, patch

import numpy as np
import pytest
import supervision as sv
from inference.core.workflows.core_steps.common.serializers import \
    serialise_sv_detections
from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1 import (
    DUPLICATED_STATUS, encode_prediction, is_prediction_registration_forbidden,
    register_datapoint, safe_register_image_at_roboflow)
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}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

class TestRegisterDatapointBasic:
    """Basic tests for register_datapoint function with normal, expected inputs."""

    def test_register_datapoint_with_dict_prediction_and_inference_id(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with a dict prediction containing inference_id.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = {INFERENCE_ID_KEY: 'inference_456', 'top': 'class1'}
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = ['tag1', 'tag2']
        metadata = {'key': 'value'}
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image, patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.annotate_image_at_roboflow') as mock_annotate:
            mock_register_image.return_value = 'roboflow_image_123'
            mock_annotate.return_value = {'success': True}
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointBasic', 'test_register_datapoint_with_dict_prediction_and_inference_id', 'register_datapoint', '9_2', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags, metadata=metadata)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            mock_register_image.assert_called_once()
            mock_annotate.assert_called_once()

    def test_register_datapoint_returns_duplicated_status_when_image_is_duplicate(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that DUPLICATED_STATUS is returned when image is duplicate.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = None
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = None
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointBasic', 'test_register_datapoint_returns_duplicated_status_when_image_is_duplicate', 'register_datapoint', '8_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output

    def test_register_datapoint_with_empty_sv_detections(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with empty sv.Detections object.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = sv.Detections.empty()
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointBasic', 'test_register_datapoint_with_empty_sv_detections', 'register_datapoint', '8_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output

    def test_register_datapoint_with_none_prediction(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with None prediction.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = None
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = ['tag1']
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointBasic', 'test_register_datapoint_with_none_prediction', 'register_datapoint', '8_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output

    def test_register_datapoint_with_valid_sv_detections(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with valid sv.Detections object.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        xyxy = np.array([[10.0, 20.0, 100.0, 200.0]])
        confidence = np.array([0.95])
        class_id = np.array([0])
        tracker_id = np.array([1])
        prediction = sv.Detections(xyxy=xyxy, confidence=confidence, class_id=class_id, tracker_id=tracker_id)
        prediction.data['class_name'] = np.array(['person'])
        prediction.data[INFERENCE_ID_KEY] = np.array(['inf_123'])
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image, patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.annotate_image_at_roboflow') as mock_annotate:
            mock_register_image.return_value = 'roboflow_image_123'
            mock_annotate.return_value = {'success': True}
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointBasic', 'test_register_datapoint_with_valid_sv_detections', 'register_datapoint', '14_2', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            mock_annotate.assert_called_once()

class TestRegisterDatapointEdgeCases:
    """Edge case tests for register_datapoint with extreme or unusual conditions."""

    def test_register_datapoint_with_empty_tags_list(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with empty tags list.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = {'top': 'class1'}
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        metadata = None
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointEdgeCases', 'test_register_datapoint_with_empty_tags_list', 'register_datapoint', '9_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags, metadata=metadata)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

    def test_register_datapoint_with_none_metadata(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with None metadata.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = {'predicted_classes': ['class1', 'class2']}
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = ['tag1']
        metadata = None
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image, patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.annotate_image_at_roboflow') as mock_annotate:
            mock_register_image.return_value = 'roboflow_image_123'
            mock_annotate.return_value = {'success': True}
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointEdgeCases', 'test_register_datapoint_with_none_metadata', 'register_datapoint', '9_2', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags, metadata=metadata)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

    def test_register_datapoint_with_dict_missing_inference_id(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with dict prediction missing inference_id.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = {'top': 'class1'}
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointEdgeCases', 'test_register_datapoint_with_dict_missing_inference_id', 'register_datapoint', '8_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

    def test_register_datapoint_with_very_long_local_image_id(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with extremely long local_image_id.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_' + 'x' * 10000
        prediction = None
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointEdgeCases', 'test_register_datapoint_with_very_long_local_image_id', 'register_datapoint', '8_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

    def test_register_datapoint_with_special_characters_in_batch_name(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with special characters in batch_name.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = None
        api_key = 'test_api_key'
        batch_name = 'test-batch_2024!@#$%'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointEdgeCases', 'test_register_datapoint_with_special_characters_in_batch_name', 'register_datapoint', '8_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

    def test_register_datapoint_with_empty_encoded_image(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with empty encoded_image bytes.'
        target_project = 'test_project'
        encoded_image = b''
        local_image_id = 'local_123'
        prediction = None
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointEdgeCases', 'test_register_datapoint_with_empty_encoded_image', 'register_datapoint', '8_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

    def test_register_datapoint_with_dict_predicted_classes_key(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with dict prediction containing predicted_classes.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = {'predicted_classes': ['cat', 'dog', 'bird']}
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image, patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.annotate_image_at_roboflow') as mock_annotate:
            mock_register_image.return_value = 'roboflow_image_123'
            mock_annotate.return_value = {'success': True}
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointEdgeCases', 'test_register_datapoint_with_dict_predicted_classes_key', 'register_datapoint', '8_2', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_annotate.call_args

    def test_register_datapoint_with_multiple_tags(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with multiple tags.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = None
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = ['tag1', 'tag2', 'tag3', 'tag4', 'tag5']
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointEdgeCases', 'test_register_datapoint_with_multiple_tags', 'register_datapoint', '8_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

    def test_register_datapoint_inference_id_from_sv_detections_with_multiple_detections(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that inference_id is extracted from first detection only.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        xyxy = np.array([[10.0, 20.0, 100.0, 200.0], [50.0, 60.0, 150.0, 250.0]])
        confidence = np.array([0.95, 0.85])
        class_id = np.array([0, 1])
        prediction = sv.Detections(xyxy=xyxy, confidence=confidence, class_id=class_id)
        prediction.data['class_name'] = np.array(['person', 'car'])
        prediction.data[INFERENCE_ID_KEY] = np.array(['inf_123', 'inf_456'])
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image, patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.annotate_image_at_roboflow') as mock_annotate:
            mock_register_image.return_value = 'roboflow_image_123'
            mock_annotate.return_value = {'success': True}
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointEdgeCases', 'test_register_datapoint_inference_id_from_sv_detections_with_multiple_detections', 'register_datapoint', '13_2', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

class TestSafeRegisterImageAtRoboflow:
    """Tests for the safe_register_image_at_roboflow helper function."""

    def test_safe_register_returns_image_id_on_success(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that image ID is returned on successful registration.'
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.register_image_at_roboflow') as mock_register:
            mock_register.return_value = {'id': 'image_xyz', 'duplicate': False, 'success': True}
            result = safe_register_image_at_roboflow(target_project='project1', encoded_image=b'image_data', local_image_id='local_1', api_key='key', batch_name='batch', tags=[], inference_id=None, metadata=None)

    def test_safe_register_returns_none_on_duplicate(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that None is returned when image is duplicate.'
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.register_image_at_roboflow') as mock_register:
            mock_register.return_value = {'duplicate': True, 'success': True}
            result = safe_register_image_at_roboflow(target_project='project1', encoded_image=b'image_data', local_image_id='local_1', api_key='key', batch_name='batch', tags=[], inference_id=None, metadata=None)

class TestIsPredictionRegistrationForbidden:
    """Tests for the is_prediction_registration_forbidden helper function."""

    def test_forbidden_with_none_prediction(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that None prediction is forbidden.'
        result = is_prediction_registration_forbidden(prediction=None)

    def test_forbidden_with_empty_sv_detections(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that empty sv.Detections is forbidden.'
        prediction = sv.Detections.empty()
        result = is_prediction_registration_forbidden(prediction=prediction)

    def test_forbidden_with_dict_missing_top_and_predicted_classes(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that dict without top or predicted_classes is forbidden.'
        prediction = {'some_other_key': 'value'}
        result = is_prediction_registration_forbidden(prediction=prediction)

    def test_allowed_with_dict_containing_top(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that dict with top key is allowed.'
        prediction = {'top': 'class1'}
        result = is_prediction_registration_forbidden(prediction=prediction)

    def test_allowed_with_dict_containing_predicted_classes(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that dict with predicted_classes key is allowed.'
        prediction = {'predicted_classes': ['class1', 'class2']}
        result = is_prediction_registration_forbidden(prediction=prediction)

    def test_allowed_with_valid_sv_detections(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that valid sv.Detections is allowed.'
        xyxy = np.array([[10.0, 20.0, 100.0, 200.0]])
        confidence = np.array([0.95])
        class_id = np.array([0])
        prediction = sv.Detections(xyxy=xyxy, confidence=confidence, class_id=class_id)
        result = is_prediction_registration_forbidden(prediction=prediction)

class TestEncodePrediction:
    """Tests for the encode_prediction helper function."""

    def test_encode_dict_with_predicted_classes(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test encoding dict with predicted_classes.'
        prediction = {'predicted_classes': ['cat', 'dog', 'bird']}
        encoded, file_type = encode_prediction(prediction=prediction)

    def test_encode_dict_with_top(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test encoding dict with top key.'
        prediction = {'top': 'class1'}
        encoded, file_type = encode_prediction(prediction=prediction)

    def test_encode_sv_detections_returns_json(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test encoding sv.Detections returns JSON.'
        xyxy = np.array([[10.0, 20.0, 100.0, 200.0]])
        confidence = np.array([0.95])
        class_id = np.array([0])
        prediction = sv.Detections(xyxy=xyxy, confidence=confidence, class_id=class_id)
        prediction.data['class_name'] = np.array(['person'])
        prediction.data['detection_id'] = np.array(['det_1'])
        encoded, file_type = encode_prediction(prediction=prediction)
        decoded = json.loads(encoded)

class TestRegisterDatapointLargeScale:
    """Large-scale tests for register_datapoint with realistic data volumes."""

    def test_register_datapoint_with_many_tags(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with 100+ tags.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = None
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = [f'tag_{i}' for i in range(150)]
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointLargeScale', 'test_register_datapoint_with_many_tags', 'register_datapoint', '8_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

    def test_register_datapoint_with_large_metadata_dict(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with large metadata dictionary.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = None
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        metadata = {f'key_{i}': f'value_{i}' for i in range(500)}
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointLargeScale', 'test_register_datapoint_with_large_metadata_dict', 'register_datapoint', '9_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags, metadata=metadata)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

    def test_register_datapoint_with_many_predicted_classes(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with dict containing many predicted_classes.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = {'predicted_classes': [f'class_{i}' for i in range(1000)]}
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image, patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.annotate_image_at_roboflow') as mock_annotate:
            mock_register_image.return_value = 'roboflow_image_123'
            mock_annotate.return_value = {'success': True}
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointLargeScale', 'test_register_datapoint_with_many_predicted_classes', 'register_datapoint', '8_2', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_annotate.call_args
            annotation_content = call_args[1]['annotation_content']

    def test_register_datapoint_with_many_sv_detections(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with sv.Detections containing many objects.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        num_detections = 200
        xyxy = np.array([[float(i * 10), float(i * 10), float(i * 10 + 50), float(i * 10 + 50)] for i in range(num_detections)])
        confidence = np.random.rand(num_detections)
        class_id = np.array([i % 10 for i in range(num_detections)])
        prediction = sv.Detections(xyxy=xyxy, confidence=confidence, class_id=class_id)
        prediction.data['class_name'] = np.array([f'class_{i % 10}' for i in range(num_detections)])
        prediction.data['detection_id'] = np.array([f'det_{i}' for i in range(num_detections)])
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image, patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.annotate_image_at_roboflow') as mock_annotate:
            mock_register_image.return_value = 'roboflow_image_123'
            mock_annotate.return_value = {'success': True}
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointLargeScale', 'test_register_datapoint_with_many_sv_detections', 'register_datapoint', '14_2', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_annotate.call_args
            annotation_content = call_args[1]['annotation_content']
            decoded = json.loads(annotation_content)

    def test_register_datapoint_with_large_encoded_image(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with very large encoded_image (>10MB).'
        target_project = 'test_project'
        encoded_image = b'x' * (11 * 1024 * 1024)
        local_image_id = 'local_123'
        prediction = None
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointLargeScale', 'test_register_datapoint_with_large_encoded_image', 'register_datapoint', '8_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args

    def test_register_datapoint_sequential_batch_registration(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registering multiple datapoints in sequence.'
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            for i in range(100):
                codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointLargeScale', 'test_register_datapoint_sequential_batch_registration', 'register_datapoint', '1_1_0', codeflash_loop_index, target_project='test_project', encoded_image=b'image_data', local_image_id=f'local_{i}', prediction=None, api_key='test_api_key', batch_name=f'batch_{i}', tags=[])
                codeflash_output = codeflash_return_value
                result = codeflash_output

    def test_register_datapoint_with_deeply_nested_metadata(self):
        np.random.seed(42)
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test registration with deeply nested metadata structure.'
        target_project = 'test_project'
        encoded_image = b'fake_image_data'
        local_image_id = 'local_123'
        prediction = None
        api_key = 'test_api_key'
        batch_name = 'test_batch'
        tags = []
        metadata = {'level1': {'level2': {'level3': {'level4': {'level5': {'key': 'deep_value'}}}}}}
        with patch('inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1.safe_register_image_at_roboflow') as mock_register_image:
            mock_register_image.return_value = 'roboflow_image_123'
            codeflash_return_value = codeflash_wrap(register_datapoint, 'tests.inference.unit_tests.test_register_datapoint__unit_test_1', 'TestRegisterDatapointLargeScale', 'test_register_datapoint_with_deeply_nested_metadata', 'register_datapoint', '9_1', codeflash_loop_index, target_project=target_project, encoded_image=encoded_image, local_image_id=local_image_id, prediction=prediction, api_key=api_key, batch_name=batch_name, tags=tags, metadata=metadata)
            codeflash_output = codeflash_return_value
            result = codeflash_output
            call_args = mock_register_image.call_args
            passed_metadata = call_args[1]['metadata']