from __future__ import annotations

import base64
import gc
import inspect
import os
import time
from pathlib import Path

import cv2
import dill as pickle
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}######!')
    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 _encode_mask_to_b64_png(mask_array: np.ndarray) -> str:
    arr = mask_array.astype(np.uint8)
    ok, buffer = cv2.imencode('.png', arr)
    b = buffer.tobytes()
    return base64.b64encode(b).decode('ascii')

def _call_convert(preds: dict):
    return RoboflowSemanticSegmentationModelBlockV1._convert_to_sv_detections(preds)

def test_single_class_mask_without_confidence_produces_one_detection():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask = np.zeros((10, 10), dtype=np.uint8)
    mask[2:5, 3:8] = 1
    seg_b64 = _encode_mask_to_b64_png(mask)
    preds = {'segmentation_mask': seg_b64, 'class_map': {'1': 'cat'}}
    result = _call_convert(preds)
    expected_xyxy = np.array([[3.0, 2.0, 7.0, 4.0]], dtype=np.float64)
    class_names = result.data.get('class_name')
    det_ids = result.data.get(DETECTION_ID_KEY)

def test_confidence_mask_affects_confidence_values():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask = np.zeros((8, 8), dtype=np.uint8)
    mask[1:4, 2:6] = 2
    conf = np.zeros_like(mask, dtype=np.uint8)
    conf[1:4, 2:6] = 128
    seg_b64 = _encode_mask_to_b64_png(mask)
    conf_b64 = _encode_mask_to_b64_png(conf)
    preds = {'segmentation_mask': seg_b64, 'confidence_mask': conf_b64, 'class_map': {'2': 'dog'}}
    result = _call_convert(preds)
    expected_conf = 128.0 / 255.0
    stored_conf = result['confidence_mask']

def test_empty_or_invalid_segmentation_returns_empty_detections():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    preds = {'segmentation_mask': '', 'class_map': {}}
    result = _call_convert(preds)

def test_all_zero_segmentation_returns_empty_detections():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask = np.zeros((5, 5), dtype=np.uint8)
    seg_b64 = _encode_mask_to_b64_png(mask)
    preds = {'segmentation_mask': seg_b64, 'class_map': {}}
    result = _call_convert(preds)

def test_class_map_missing_uses_string_of_id():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask = np.zeros((6, 6), dtype=np.uint8)
    mask[0:2, 0:3] = 3
    seg_b64 = _encode_mask_to_b64_png(mask)
    preds = {'segmentation_mask': seg_b64}
    result = _call_convert(preds)
    class_names = result.data.get('class_name')

def test_confidence_mask_shape_mismatch_raises_error():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    seg_mask = np.zeros((5, 5), dtype=np.uint8)
    seg_mask[1:4, 1:4] = 1
    seg_b64 = _encode_mask_to_b64_png(seg_mask)
    conf_mask = np.zeros((4, 4), dtype=np.uint8)
    conf_b64 = _encode_mask_to_b64_png(conf_mask)
    preds = {'segmentation_mask': seg_b64, 'confidence_mask': conf_b64, 'class_map': {'1': 'mismatch'}}
    with pytest.raises(Exception):
        _call_convert(preds)

def test_many_classes_large_mask_performance_and_correctness():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    num_classes = 200
    height = num_classes
    width = 50
    mask = np.zeros((height, width), dtype=np.uint8)
    for i in range(num_classes):
        mask[i, :] = i + 1
    class_map = {str(i + 1): f'class_{i + 1}' for i in range(num_classes)}
    seg_b64 = _encode_mask_to_b64_png(mask)
    preds = {'segmentation_mask': seg_b64, 'class_map': class_map}
    result = _call_convert(preds)
    expected_class_ids = np.arange(1, num_classes + 1)
    samples = [0, num_classes // 2, num_classes - 1]
    for idx in samples:
        expected_name = f'class_{idx + 1}'
        expected_xy = np.array([0.0, float(idx), float(width - 1), float(idx)], dtype=np.float64)