from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path
from typing import Any, Callable, Optional

import cv2
import dill as pickle
import numpy as np
import pytest
from inference.core.workflows.core_steps.visualizations.text_display.utils import \
    draw_background_with_alpha


def codeflash_wrap(
    codeflash_wrapped: Callable[..., Any],
    codeflash_test_module_name: str,
    codeflash_test_class_name: str | None,
    codeflash_test_name: str,
    codeflash_function_name: str,
    codeflash_line_id: str,
    codeflash_loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> Any:
    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 blank_img(width, height, color=(0, 0, 0)):
    """Create a blank image of the given color (BGR)."""
    arr = np.zeros((height, width, 3), dtype=np.uint8)
    arr[:, :] = color
    return arr

def images_equal(img1, img2):
    return np.array_equal(img1, img2)

def test_basic_full_alpha_rectangle():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test drawing a solid rectangle with alpha=1 (should fully overwrite region).'
    img = blank_img(10, 10, color=(0, 0, 0))
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (2, 2), (7, 7), (10, 20, 30), alpha=1.0, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_basic_full_alpha_rectangle', 'draw_background_with_alpha', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    roi = img[2:7, 2:7]

def test_basic_zero_alpha_rectangle():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test drawing with alpha=0 (should not change the image).'
    img = blank_img(10, 10, color=(50, 60, 70))
    img_before = img.copy()
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (2, 2), (7, 7), (100, 110, 120), alpha=0.0, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_basic_zero_alpha_rectangle', 'draw_background_with_alpha', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_basic_half_alpha_rectangle():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test drawing with alpha=0.5 (should blend colors equally).'
    img = blank_img(10, 10, color=(100, 100, 100))
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (0, 0), (10, 10), (200, 0, 0), alpha=0.5, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_basic_half_alpha_rectangle', 'draw_background_with_alpha', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    expected = np.array([150, 50, 50], dtype=np.uint8)

def test_basic_rounded_rectangle():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that drawing with border_radius>0 does not raise and modifies image.'
    img = blank_img(20, 20, color=(0, 0, 0))
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (5, 5), (15, 15), (0, 255, 0), alpha=1.0, border_radius=4)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_basic_rounded_rectangle', 'draw_background_with_alpha', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_basic_non_square_rectangle():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test drawing a non-square rectangle.'
    img = blank_img(20, 10, color=(0, 0, 0))
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (2, 1), (18, 8), (123, 222, 111), alpha=1.0, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_basic_non_square_rectangle', 'draw_background_with_alpha', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    roi = img[1:8, 2:18]

def test_edge_rectangle_outside_image():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Rectangle completely outside image should not change the image.'
    img = blank_img(10, 10, color=(10, 20, 30))
    img_before = img.copy()
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (20, 20), (30, 30), (255, 0, 0), alpha=1.0, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_edge_rectangle_outside_image', 'draw_background_with_alpha', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_rectangle_partially_outside_image():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Rectangle partially outside image should be clamped to image bounds.'
    img = blank_img(10, 10, color=(0, 0, 0))
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (-5, -5), (5, 5), (255, 255, 255), alpha=1.0, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_edge_rectangle_partially_outside_image', 'draw_background_with_alpha', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    roi = img[0:5, 0:5]

def test_edge_zero_area_rectangle():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Rectangle with zero area should not modify the image.'
    img = blank_img(10, 10, color=(1, 2, 3))
    img_before = img.copy()
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (5, 5), (5, 10), (10, 20, 30), alpha=1.0, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_edge_zero_area_rectangle', 'draw_background_with_alpha', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (5, 5), (10, 5), (10, 20, 30), alpha=1.0, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_edge_zero_area_rectangle', 'draw_background_with_alpha', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_negative_border_radius():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Negative border_radius should be treated as 0 (rectangle).'
    img = blank_img(10, 10, color=(0, 0, 0))
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (2, 2), (8, 8), (50, 100, 150), alpha=1.0, border_radius=-5)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_edge_negative_border_radius', 'draw_background_with_alpha', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    roi = img[2:8, 2:8]

def test_edge_large_border_radius():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'border_radius larger than half the rect min side should be clamped.'
    img = blank_img(10, 10, color=(0, 0, 0))
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (2, 2), (9, 9), (100, 200, 50), alpha=1.0, border_radius=100)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_edge_large_border_radius', 'draw_background_with_alpha', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_alpha_out_of_bounds():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Alpha < 0 should act as 0, alpha > 1 as 1 (cv2.addWeighted clamps).'
    img = blank_img(10, 10, color=(10, 10, 10))
    img_copy = img.copy()
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (0, 0), (10, 10), (200, 0, 0), alpha=-0.5, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_edge_alpha_out_of_bounds', 'draw_background_with_alpha', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (0, 0), (10, 10), (200, 0, 0), alpha=2.0, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_edge_alpha_out_of_bounds', 'draw_background_with_alpha', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_single_pixel_rectangle():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test drawing a 1x1 rectangle.'
    img = blank_img(5, 5, color=(0, 0, 0))
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (2, 2), (3, 3), (255, 100, 50), alpha=1.0, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_edge_single_pixel_rectangle', 'draw_background_with_alpha', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_edge_rectangle_touching_image_border():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test rectangle exactly on the image border.'
    img = blank_img(5, 5, color=(0, 0, 0))
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (0, 0), (5, 5), (11, 22, 33), alpha=1.0, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_edge_rectangle_touching_image_border', 'draw_background_with_alpha', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_full_image_rectangle():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Draw a rectangle covering the whole image (500x500).'
    img = blank_img(500, 500, color=(10, 20, 30))
    _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (0, 0), (500, 500), (100, 150, 200), alpha=0.7, border_radius=0)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_large_full_image_rectangle', 'draw_background_with_alpha', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    expected = (0.7 * np.array([100, 150, 200]) + 0.3 * np.array([10, 20, 30])).astype(np.uint8)
    for y in [0, 250, 499]:
        for x in [0, 250, 499]:
            pass

def test_large_many_small_rectangles():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Draw many small rectangles over a large image.'
    img = blank_img(100, 100, color=(0, 0, 0))
    for i in range(0, 100, 10):
        for j in range(0, 100, 10):
            _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (i, j), (i + 10, j + 10), (i, j, 255), alpha=1.0, border_radius=3)
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_large_many_small_rectangles', 'draw_background_with_alpha', '2_0_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    for i in range(5, 100, 10):
        for j in range(5, 100, 10):
            pass

def test_large_performance_multiple_calls():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test many sequential calls do not crash or slow down.'
    img = blank_img(50, 50, color=(0, 0, 0))
    for i in range(50):
        _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (i, 0), (min(i + 10, 50), 50), (i * 5 % 256, i * 7 % 256, i * 11 % 256), alpha=0.3 + 0.01 * i, border_radius=i % 6)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_large_performance_multiple_calls', 'draw_background_with_alpha', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_large_alpha_gradient():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Draw rectangles with increasing alpha to form a gradient.'
    img = blank_img(100, 10, color=(0, 0, 0))
    for i in range(10):
        _call__bound__arguments = inspect.signature(draw_background_with_alpha).bind(img, (i * 10, 0), ((i + 1) * 10, 10), (255, 0, 0), alpha=i / 9, border_radius=0)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(draw_background_with_alpha, 'tests.inference.unit_tests.test_draw_background_with_alpha__unit_test_0', None, 'test_large_alpha_gradient', 'draw_background_with_alpha', '2_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    mid = img[5, 45]