from __future__ import annotations

import base64
import gc
import os
import time

import cv2
import numpy as np
import pytest
import supervision as sv
from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v1 import \
    RoboflowSemanticSegmentationModelBlockV1
from inference.core.workflows.execution_engine.constants import \
    DETECTION_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

def _encode_image_to_b64_png(img: np.ndarray) -> str:
    ok, buffer = cv2.imencode('.png', img)
    b64 = base64.b64encode(buffer.tobytes()).decode('ascii')
    return b64

def test_single_class_detection_basic():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    h, w = (10, 10)
    mask = np.zeros((h, w), dtype=np.uint8)
    mask[2:6, 3:8] = 1
    predictions = {'segmentation_mask': _encode_image_to_b64_png(mask), 'class_map': {'1': 'cat'}}
    codeflash_return_value = codeflash_wrap(RoboflowSemanticSegmentationModelBlockV1._convert_to_sv_detections, 'tests.inference.unit_tests.test__convert_to_sv_detections__unit_test_0', None, 'test_single_class_detection_basic', '_convert_to_sv_detections', '4', codeflash_loop_index, predictions)
    codeflash_output = codeflash_return_value
    detections = codeflash_output
    expected_box = np.array([3, 2, 7, 5], dtype=np.float64)
    class_names = detections.data['class_name']
    det_ids = detections.data[DETECTION_ID_KEY]

def test_multiple_classes_and_fallback_class_name():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    h, w = (12, 12)
    mask = np.zeros((h, w), dtype=np.uint8)
    mask[1:4, 1:5] = 1
    mask[8:11, 7:11] = 2
    predictions = {'segmentation_mask': _encode_image_to_b64_png(mask), 'class_map': {'1': 'alpha'}}
    codeflash_return_value = codeflash_wrap(RoboflowSemanticSegmentationModelBlockV1._convert_to_sv_detections, 'tests.inference.unit_tests.test__convert_to_sv_detections__unit_test_0', None, 'test_multiple_classes_and_fallback_class_name', '_convert_to_sv_detections', '5', codeflash_loop_index, predictions)
    codeflash_output = codeflash_return_value
    detections = codeflash_output
    class_ids = [int(x) for x in detections.class_id.tolist()]
    class_names = detections.data['class_name']

def test_empty_input_returns_empty_detections():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    codeflash_return_value = codeflash_wrap(RoboflowSemanticSegmentationModelBlockV1._convert_to_sv_detections, 'tests.inference.unit_tests.test__convert_to_sv_detections__unit_test_0', None, 'test_empty_input_returns_empty_detections', '_convert_to_sv_detections', '0', codeflash_loop_index, {})
    codeflash_output = codeflash_return_value
    detections = codeflash_output

def test_invalid_image_bytes_returns_empty_detections():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    invalid_b64 = base64.b64encode(b'this is not a png').decode('ascii')
    codeflash_return_value = codeflash_wrap(RoboflowSemanticSegmentationModelBlockV1._convert_to_sv_detections, 'tests.inference.unit_tests.test__convert_to_sv_detections__unit_test_0', None, 'test_invalid_image_bytes_returns_empty_detections', '_convert_to_sv_detections', '1', codeflash_loop_index, {'segmentation_mask': invalid_b64})
    codeflash_output = codeflash_return_value
    detections = codeflash_output

def test_confidence_mask_affects_confidence_values():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    h, w = (8, 8)
    mask = np.zeros((h, w), dtype=np.uint8)
    mask[0:4, 0:4] = 1
    mask[4:8, 4:8] = 2
    conf = np.zeros((h, w), dtype=np.uint8)
    conf[0:4, 0:4] = 127
    conf[4:8, 4:8] = 255
    predictions = {'segmentation_mask': _encode_image_to_b64_png(mask), 'confidence_mask': _encode_image_to_b64_png(conf), 'class_map': {'1': 'one', '2': 'two'}}
    codeflash_return_value = codeflash_wrap(RoboflowSemanticSegmentationModelBlockV1._convert_to_sv_detections, 'tests.inference.unit_tests.test__convert_to_sv_detections__unit_test_0', None, 'test_confidence_mask_affects_confidence_values', '_convert_to_sv_detections', '8', codeflash_loop_index, predictions)
    codeflash_output = codeflash_return_value
    detections = codeflash_output
    confs = detections.confidence.tolist()
    expected_conf1 = float(127) / 255.0
    expected_conf2 = float(255) / 255.0

def test_many_classes_performance_and_correctness():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    num_classes = 255
    w = 3
    h = num_classes
    mask = np.zeros((h, w), dtype=np.uint8)
    for i in range(num_classes):
        mask[i, :] = i + 1
    predictions = {'segmentation_mask': _encode_image_to_b64_png(mask), 'class_map': {}}
    codeflash_return_value = codeflash_wrap(RoboflowSemanticSegmentationModelBlockV1._convert_to_sv_detections, 'tests.inference.unit_tests.test__convert_to_sv_detections__unit_test_0', None, 'test_many_classes_performance_and_correctness', '_convert_to_sv_detections', '6', codeflash_loop_index, predictions)
    codeflash_output = codeflash_return_value
    detections = codeflash_output
    sample_indices = [0, num_classes // 2, num_classes - 1]
    for idx in sample_indices:
        expected_box = np.array([0, idx, w - 1, idx], dtype=np.float64)
    det_ids = detections.data[DETECTION_ID_KEY]