/*
 * Copyright (c) 2019 - 2025 Andri Yngvason
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
 * REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
 * AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
 * INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
 * LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
 * OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
 * PERFORMANCE OF THIS SOFTWARE.
 */

#include "neatvnc.h"
#include "rfb-proto.h"
#include "vec.h"
#include "frame.h"
#include "pixels.h"
#include "enc/util.h"
#include "enc/encoder.h"

#include <stdlib.h>
#include <pixman.h>
#include <aml.h>

struct encoder* raw_encoder_new(void);

struct raw_encoder {
	struct encoder encoder;
	struct rfb_pixel_format output_format;
	struct aml_work* work;
};

struct raw_encoder_work {
	struct raw_encoder* parent;
	struct rfb_pixel_format output_format;
	struct nvnc_composite_fb composite_fb;
	struct pixman_region16 damage;
	int n_rects;
	struct encoded_frame *result;
};

struct encoder_impl encoder_impl_raw;

static inline struct raw_encoder* raw_encoder(struct encoder* encoder)
{
	assert(encoder->impl == &encoder_impl_raw);
	return (struct raw_encoder*)encoder;
}

static int raw_encode_box(struct raw_encoder_work* ctx, struct vec* dst,
		const struct rfb_pixel_format* dst_fmt,
		const struct nvnc_frame* fb,
		const struct rfb_pixel_format* src_fmt, int x_start,
		int y_start, int stride, int width, int height)
{
	uint16_t x_pos = fb->x_off;
	uint16_t y_pos = fb->y_off;

	int rc = -1;

	rc = nvnc__encode_rect_head(dst, RFB_ENCODING_RAW, x_pos + x_start,
			y_pos + y_start, width, height);
	if (rc < 0)
		return -1;

	uint8_t* b = fb->buffer->addr;
	int32_t src_bpp = src_fmt->bits_per_pixel / 8;
	int32_t xoff = x_start * src_bpp;
	int32_t src_stride = fb->stride * src_bpp;

	int bpp = dst_fmt->bits_per_pixel / 8;

	rc = vec_reserve(dst, width * height * bpp + dst->len);
	if (rc < 0)
		return -1;

	uint8_t* d = dst->data;

	for (int y = y_start; y < y_start + height; ++y) {
		pixel_to_cpixel(d + dst->len, dst_fmt,
				b + xoff + y * src_stride, src_fmt,
				bpp, width);
		dst->len += width * bpp;
	}

	return 0;
}

static void raw_encoder_do_work(struct aml_work* work)
{
	struct raw_encoder_work* ctx = aml_get_userdata(work);
	int rc __attribute__((unused));

	assert(ctx->composite_fb.n_fbs != 0);

	int n_rects = 0;
	struct pixman_region16 subregions[NVNC_FB_COMPOSITE_MAX] = { 0 };
	for (int i = 0; i < ctx->composite_fb.n_fbs; ++i) {
		struct nvnc_frame* fb = ctx->composite_fb.fbs[i];
		assert(fb);

		pixman_region_init(&subregions[i]);
		pixman_region_intersect_rect(&subregions[i], &ctx->damage,
				fb->x_off, fb->y_off, fb->width, fb->height);
		n_rects += pixman_region_n_rects(&subregions[i]);
	}

	if (n_rects > UINT16_MAX) {
		n_rects = ctx->composite_fb.n_fbs;

		for (int i = 0; i < ctx->composite_fb.n_fbs; ++i) {
			struct nvnc_frame* fb = ctx->composite_fb.fbs[i];
			assert(fb);
			pixman_region_fini(&subregions[i]);
			pixman_region_init_rect(&subregions[i], fb->x_off,
					fb->y_off, fb->width, fb->height);
		}
	}

	size_t bpp = ctx->output_format.bits_per_pixel / 8;
	size_t buffer_size = nvnc__calculate_region_area(&ctx->damage) * bpp
		+ n_rects * sizeof(struct rfb_server_fb_rect);

	struct vec dst;
	rc = vec_init(&dst, buffer_size);
	assert(rc == 0);

	for (int i = 0; i < ctx->composite_fb.n_fbs; ++i) {
		struct nvnc_frame* fb = ctx->composite_fb.fbs[i];
		assert(fb);

		struct rfb_pixel_format src_fmt;
		rc = rfb_pixfmt_from_fourcc(&src_fmt,
				nvnc_frame_get_fourcc_format(fb));
		assert(rc == 0);

		rc = nvnc_frame_map(fb);
		nvnc_assert(rc == 0, "Failed to map framebuffer for encoding");

		int n_subrects = 0;
		struct pixman_box16* box = pixman_region_rectangles(
				&subregions[i], &n_subrects);

		for (int i = 0; i < n_subrects; ++i) {
			int x = box[i].x1;
			int y = box[i].y1;
			int box_width = box[i].x2 - x;
			int box_height = box[i].y2 - y;

			rc = raw_encode_box(ctx, &dst, &ctx->output_format, fb,
					&src_fmt, x - fb->x_off, y - fb->y_off,
					fb->stride, box_width, box_height);
			nvnc_assert(rc == 0, "Failed to encode box");
		}
	}
	ctx->n_rects = n_rects;

	uint16_t width = nvnc_composite_fb_width(&ctx->composite_fb);
	uint16_t height = nvnc_composite_fb_height(&ctx->composite_fb);
	uint64_t pts = nvnc_composite_fb_pts(&ctx->composite_fb);

	ctx->result = nvnc__encoded_frame_new(dst.data, dst.len, n_rects, width,
			height, pts);
	assert(ctx->result);

	for (int i = 0; i < ctx->composite_fb.n_fbs; ++i)
		pixman_region_fini(&subregions[i]);
}

static void raw_encoder_on_done(struct aml_work* work)
{
	struct raw_encoder_work* ctx = aml_get_userdata(work);
	struct raw_encoder* self = ctx->parent;

	assert(ctx->result);

	aml_unref(self->work);
	self->work = NULL;

	ctx->result->metadata = ctx->composite_fb.metadata;

	encoder_finish_frame(&self->encoder, ctx->result);
}

struct encoder* raw_encoder_new(void)
{
	struct raw_encoder* self = calloc(1, sizeof(*self));
	if (!self)
		return NULL;

	encoder_init(&self->encoder, &encoder_impl_raw);

	return (struct encoder*)self;
}

static void raw_encoder_destroy(struct encoder* encoder)
{
	struct raw_encoder* self = raw_encoder(encoder);
	if (self->work) {
		aml_stop(aml_get_default(), self->work);
		aml_unref(self->work);
	}
	free(self);
}

static void raw_encoder_set_output_format(struct encoder* encoder,
		const struct rfb_pixel_format* pixfmt)
{
	struct raw_encoder* self = raw_encoder(encoder);
	memcpy(&self->output_format, pixfmt, sizeof(self->output_format));
}

static void raw_encoder_work_destroy(void* obj)
{
	struct raw_encoder_work* ctx = obj;
	nvnc_composite_fb_unref(&ctx->composite_fb);
	pixman_region_fini(&ctx->damage);
	if (ctx->result)
		encoded_frame_unref(ctx->result);
	encoder_unref(&ctx->parent->encoder);
	free(ctx);
}

static int raw_encoder_encode(struct encoder* encoder,
		struct nvnc_composite_fb* fb,
		struct pixman_region16* damage)
{
	struct raw_encoder* self = raw_encoder(encoder);

	struct raw_encoder_work* ctx = calloc(1, sizeof(*ctx));
	if (!ctx)
		return -1;

	self->work = aml_work_new(raw_encoder_do_work, raw_encoder_on_done,
			ctx, raw_encoder_work_destroy);
	if (!self->work) {
		free(ctx);
		return -1;
	}

	encoder_ref(encoder);

	ctx->parent = self;
	nvnc_composite_fb_copy(&ctx->composite_fb, fb);
	memcpy(&ctx->output_format, &self->output_format,
			sizeof(ctx->output_format));
	pixman_region_copy(&ctx->damage, damage);

	int rc = aml_start(aml_get_default(), self->work);
	if (rc < 0) {
		aml_unref(self->work);
		self->work = NULL;
	}

	return rc;
}

struct encoder_impl encoder_impl_raw = {
	.destroy = raw_encoder_destroy,
	.set_output_format = raw_encoder_set_output_format,
	.encode = raw_encoder_encode,
};
