/*
 * Copyright (c) 2019 Linaro Limited
 * Copyright 2025 NXP
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include <math.h>

#include <zephyr/device.h>
#include <zephyr/drivers/display.h>
#include <zephyr/drivers/video-controls.h>
#include <zephyr/drivers/video.h>
#include <zephyr/kernel.h>
#include <zephyr/logging/log.h>
#include <zephyr/ztest.h>

LOG_MODULE_REGISTER(test_pattern, LOG_LEVEL_INF);

#define LAB_THRESHOLD ((double)CONFIG_TEST_LAB_THRESHOLD)

#define BARS_NUM   8
#define PIXELS_NUM 5

typedef struct {
	double L;
	double a;
	double b;
} CIELAB;

/*
 * This is measured on a real 8-colorbar pattern generated by an ov5640 camera sensor.
 * For other sensors, it can be slightly different. If it doesn't fit anymore, either
 * this array or the LAB_THRESHOLD can be modified.
 */
static const CIELAB colorbars_target[] = {
	{100.0, 0.0053, -0.0104},     /* White */
	{97.1804, -21.2151, 91.3538}, /* Yellow */
	{90.1352, -58.4675, 6.0570},  /* Cyan */
	{87.7630, -85.9469, 83.2128}, /* Green */
	{56.6641, 95.0182, -66.9129}, /* Magenta */
	{46.6937, 72.7494, 49.5801},  /* Red */
	{27.6487, 71.5662, -97.4712}, /* Blue */
	{1.3726, -2.8040, 2.0043},    /* Black */
};

static inline CIELAB rgb888_to_lab(const uint8_t r, const uint8_t g, const uint8_t b)
{
	CIELAB lab;

	double r_lin = r / 255.0;
	double g_lin = g / 255.0;
	double b_lin = b / 255.0;

	r_lin = r_lin > 0.04045 ? pow((r_lin + 0.055) / 1.055, 2.4) : r_lin / 12.92;
	g_lin = g_lin > 0.04045 ? pow((g_lin + 0.055) / 1.055, 2.4) : g_lin / 12.92;
	b_lin = b_lin > 0.04045 ? pow((b_lin + 0.055) / 1.055, 2.4) : b_lin / 12.92;

	double x = r_lin * 0.4124 + g_lin * 0.3576 + b_lin * 0.1805;
	double y = r_lin * 0.2126 + g_lin * 0.7152 + b_lin * 0.0722;
	double z = r_lin * 0.0193 + g_lin * 0.1192 + b_lin * 0.9505;

	x /= 0.95047;
	z /= 1.08883;

	x = x > 0.008856 ? pow(x, 1.0 / 3.0) : (7.787 * x) + (16.0 / 116.0);
	y = y > 0.008856 ? pow(y, 1.0 / 3.0) : (7.787 * y) + (16.0 / 116.0);
	z = z > 0.008856 ? pow(z, 1.0 / 3.0) : (7.787 * z) + (16.0 / 116.0);

	lab.L = 116.0 * y - 16.0;
	lab.a = 500.0 * (x - y);
	lab.b = 200.0 * (y - z);

	return lab;
}

static inline CIELAB xrgb32_to_lab(const uint32_t color)
{
	uint8_t r = (color >> 16) & 0xFF;
	uint8_t g = (color >> 8) & 0xFF;
	uint8_t b = color & 0xFF;

	return rgb888_to_lab(r, g, b);
}

static inline CIELAB rgb565_to_lab(const uint16_t color)
{
	uint8_t r5 = (color >> 11) & 0x1F;
	uint8_t g6 = (color >> 5) & 0x3F;
	uint8_t b5 = color & 0x1F;

	/* Convert RGB565 to RGB888 */
	uint8_t r = (r5 * 255) / 31;
	uint8_t g = (g6 * 255) / 63;
	uint8_t b = (b5 * 255) / 31;

	return rgb888_to_lab(r, g, b);
}

static inline void sum_lab(CIELAB *sum, const CIELAB lab)
{
	sum->L += lab.L;
	sum->a += lab.a;
	sum->b += lab.b;
}

static inline void average_lab(CIELAB *lab, const uint32_t count)
{
	if (count > 0) {
		lab->L /= count;
		lab->a /= count;
		lab->b /= count;
	}
}

static inline double deltaE(const CIELAB lab1, const CIELAB lab2)
{
	return sqrt(pow(lab1.L - lab2.L, 2) + pow(lab1.a - lab2.a, 2) + pow(lab1.b - lab2.b, 2));
}

/*
 * As color values may vary near the boundary of each bar and also, for computational
 * efficiency, check only a small number of pixels (PIXELS_NUM) in the middle of each bar.
 */
static inline bool is_colorbar_ok(const uint8_t *const buf, const struct video_format *fmt)
{
	int i;
	int bw = fmt->width / BARS_NUM;
	CIELAB colorbars[BARS_NUM] = {0};

	for (int h = 0; h < fmt->height; h++) {
		for (i = 0; i < BARS_NUM; i++) {
			if (fmt->pixelformat == VIDEO_PIX_FMT_XRGB32) {
				uint32_t *pixel =
					(uint32_t *)&buf[4 * (h * fmt->width + bw / 2 + i * bw)];

				for (int j = -PIXELS_NUM / 2; j <= PIXELS_NUM / 2; j++) {
					sum_lab(&colorbars[i], xrgb32_to_lab(*(pixel + j)));
				}
			} else if (fmt->pixelformat == VIDEO_PIX_FMT_RGB565) {
				uint16_t *pixel =
					(uint16_t *)&buf[2 * (h * fmt->width + bw / 2 + i * bw)];

				for (int j = -PIXELS_NUM / 2; j <= PIXELS_NUM / 2; j++) {
					sum_lab(&colorbars[i], rgb565_to_lab(*(pixel + j)));
				}
			} else {
				printk("Format %d is not supported", fmt->pixelformat);
				return false;
			}
		}
	}

	for (i = 0; i < BARS_NUM; i++) {
		average_lab(&colorbars[i], PIXELS_NUM * fmt->height);
		if (deltaE(colorbars[i], colorbars_target[i]) > LAB_THRESHOLD) {
			return false;
		}
	}

	return true;
}

static const struct device *const video_dev = DEVICE_DT_GET(DT_CHOSEN(zephyr_camera));
struct video_format fmt;

static void *test_pattern_setup(void)
{
	struct video_buffer *vbuf = &(struct video_buffer){};
	struct video_caps caps = {
		.type = VIDEO_BUF_TYPE_OUTPUT,
	};
	struct video_control ctrl = {
		.id = VIDEO_CID_TEST_PATTERN, .val = CONFIG_TEST_PATTERN_CTRL,
	};
	int ret;

	zassert(device_is_ready(video_dev),  "device initialization failed");

	ret = video_get_caps(video_dev, &caps);
	zassert_ok(ret, "getting video capabilities failed");

	fmt.type = VIDEO_BUF_TYPE_OUTPUT;
	ret = video_get_format(video_dev, &fmt);
	zassert_ok(ret, "getting default video format failed");

	if (CONFIG_TEST_FRAME_HEIGHT > 0) {
		fmt.height = CONFIG_TEST_FRAME_HEIGHT;
	}
	if (CONFIG_TEST_FRAME_WIDTH > 0) {
		fmt.width = CONFIG_TEST_FRAME_WIDTH;
	}
	if (strcmp(CONFIG_TEST_PIXEL_FORMAT, "") != 0) {
		fmt.pixelformat = VIDEO_FOURCC_FROM_STR(CONFIG_TEST_PIXEL_FORMAT);
	}

	LOG_INF("Video format: %s %ux%u",
		VIDEO_FOURCC_TO_STR(fmt.pixelformat), fmt.width, fmt.height);

	ret = video_set_format(video_dev, &fmt);
	zassert_ok(ret, "setting video format failed");

	ret = video_set_ctrl(video_dev, &ctrl);
	zassert_ok(ret, "setting test pattern");

	/* Alloc video buffers and enqueue for capture */
	zassert(caps.min_vbuf_count <= CONFIG_VIDEO_BUFFER_POOL_NUM_MAX,
		"not enough buffers");
	zassert(fmt.size <= (CONFIG_VIDEO_BUFFER_POOL_HEAP_SIZE / CONFIG_VIDEO_BUFFER_POOL_NUM_MAX),
		"buffers too large");

	for (int i = 0; i < CONFIG_VIDEO_BUFFER_POOL_NUM_MAX; i++) {
		vbuf = video_buffer_aligned_alloc(fmt.size, CONFIG_VIDEO_BUFFER_POOL_ALIGN,
						  K_NO_WAIT);
		zassert_not_null(vbuf);

		vbuf->type = VIDEO_BUF_TYPE_OUTPUT;

		ret = video_enqueue(video_dev, vbuf);
		zassert_ok(ret);
	}

	LOG_INF("Device %s configured starting capture", video_dev->name);

	ret = video_stream_start(video_dev, VIDEO_BUF_TYPE_OUTPUT);
	zassert_ok(ret);

	return NULL;
}

void test_pattern_after(void *)
{
	int ret;

	ret = video_stream_stop(video_dev, VIDEO_BUF_TYPE_OUTPUT);
	zassert_ok(ret);
}

ZTEST(test_pattern, test_pattern_frames)
{
	struct video_buffer *vbuf = &(struct video_buffer){
		.type = VIDEO_BUF_TYPE_OUTPUT
	};
	size_t valid = 0;
	int ret;

	for (size_t i = 0; i < CONFIG_TEST_FRAMES_TOTAL; i++) {
		ret = video_dequeue(video_dev, &vbuf, K_FOREVER);
		zassert_ok(ret);

		LOG_INF("Got frame, testing color bars");

		valid += is_colorbar_ok(vbuf->buffer, &fmt);
		if (valid >= CONFIG_TEST_FRAMES_VALID) {
			LOG_INF("Got %u valid frames out of %u, stopping the test", valid, i + 1);
			break;
		}

		ret = video_enqueue(video_dev, vbuf);
		zassert_ok(ret);
	}

	zassert_equal(valid, CONFIG_TEST_FRAMES_VALID,
		      "there should be at least %u valid frames out of %u",
		      CONFIG_TEST_FRAMES_VALID, CONFIG_TEST_FRAMES_TOTAL);
}

ZTEST_SUITE(test_pattern, NULL, test_pattern_setup, NULL, test_pattern_after, NULL);