from __future__ import annotations

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

import cv2 as cv
import dill as pickle
import numpy as np
import pytest
import supervision as sv
from inference.core.workflows.core_steps.classical_cv.mask_area_measurement.v1 import (
    OUTPUT_KEY, MaskAreaMeasurementBlockV1, get_detection_area)


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_bbox_array(list_of_boxes):
    """Helper to convert a list of [x1,y1,x2,y2] boxes into an np.ndarray shape (N,4)."""
    return np.array(list_of_boxes, dtype=float)

def _make_mask_stack(list_of_2d_masks):
    """Helper to stack a list of 2D masks (H x W) into a (N, H, W) array as expected by the code."""
    stack = np.stack([np.array(m, dtype=np.uint8) for m in list_of_2d_masks], axis=0)
    return stack

def test_single_bbox_without_mask_uses_bbox_area():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    xyxy = _make_bbox_array([[0, 0, 10, 5]])
    detections = sv.Detections(xyxy=xyxy)
    block = MaskAreaMeasurementBlockV1()
    _call__bound__arguments = inspect.signature(block.run).bind(detections)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_single_bbox_without_mask_uses_bbox_area', 'run', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_mask_with_contour_uses_contour_area():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask = np.zeros((10, 10), dtype=np.uint8)
    mask[2:6, 3:7] = 255
    stacked = _make_mask_stack([mask])
    xyxy = _make_bbox_array([[0, 0, 100, 100]])
    detections = sv.Detections(xyxy=xyxy, mask=stacked)
    area_via_helper = get_detection_area(detections, 0)
    block = MaskAreaMeasurementBlockV1()
    _call__bound__arguments = inspect.signature(block.run).bind(detections)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_single_mask_with_contour_uses_contour_area', 'run', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_mask_without_contours_falls_back_to_bbox():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask = np.zeros((8, 8), dtype=np.uint8)
    stacked = _make_mask_stack([mask])
    xyxy = _make_bbox_array([[1, 2, 6, 7]])
    detections = sv.Detections(xyxy=xyxy, mask=stacked)
    block = MaskAreaMeasurementBlockV1()
    _call__bound__arguments = inspect.signature(block.run).bind(detections)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_mask_without_contours_falls_back_to_bbox', 'run', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_zero_sized_bbox_returns_zero_area():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    xyxy = _make_bbox_array([[5, 5, 5, 10]])
    detections = sv.Detections(xyxy=xyxy)
    block = MaskAreaMeasurementBlockV1()
    _call__bound__arguments = inspect.signature(block.run).bind(detections)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_zero_sized_bbox_returns_zero_area', 'run', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_reversed_coordinates_can_produce_negative_area():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    xyxy = _make_bbox_array([[10, 0, 5, 10]])
    detections = sv.Detections(xyxy=xyxy)
    block = MaskAreaMeasurementBlockV1()
    _call__bound__arguments = inspect.signature(block.run).bind(detections)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_reversed_coordinates_can_produce_negative_area', 'run', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_multiple_mixed_detections_combined_behavior():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask1 = np.zeros((6, 6), dtype=np.uint8)
    mask1[1:3, 1:3] = 255
    mask2 = np.zeros((6, 6), dtype=np.uint8)
    masks = _make_mask_stack([mask1, mask2])
    xyxy_all = _make_bbox_array([[0, 0, 10, 10], [1, 1, 7, 7], [2, 2, 5, 9]])
    mask3 = np.zeros((6, 6), dtype=np.uint8)
    combined_masks = _make_mask_stack([mask1, mask2, mask3])
    detections = sv.Detections(xyxy=xyxy_all, mask=combined_masks)
    block = MaskAreaMeasurementBlockV1()
    _call__bound__arguments = inspect.signature(block.run).bind(detections)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_multiple_mixed_detections_combined_behavior', 'run', '9', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = [pytest.approx(4.0, rel=1e-06), 36.0, 21.0]

def test_large_scale_bbox_only_many_detections():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    boxes = [[0, 0, float(i), float(i)] for i in range(n)]
    xyxy = _make_bbox_array(boxes)
    detections = sv.Detections(xyxy=xyxy)
    block = MaskAreaMeasurementBlockV1()
    _call__bound__arguments = inspect.signature(block.run).bind(detections)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_large_scale_bbox_only_many_detections', 'run', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    mid_index = n // 2

def test_large_scale_masks_many_detections_performance_and_correctness():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 200
    masks = []
    boxes = []
    for i in range(n):
        size = 1 + i % 7
        m = np.zeros((8, 8), dtype=np.uint8)
        m[0:size, 0:size] = 255
        masks.append(m)
        boxes.append([0, 0, 100.0, 100.0])
    xyxy = _make_bbox_array(boxes)
    stacked_masks = _make_mask_stack(masks)
    detections = sv.Detections(xyxy=xyxy, mask=stacked_masks)
    block = MaskAreaMeasurementBlockV1()
    _call__bound__arguments = inspect.signature(block.run).bind(detections)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(block.run, 'tests.inference.unit_tests.test_run__unit_test_0', None, 'test_large_scale_masks_many_detections_performance_and_correctness', 'run', '8', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for i in [0, 1, 6, 7, 13, n - 1]:
        expected_size = 1 + i % 7
        expected_area = float(expected_size * expected_size)