/*
 * 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 "pixels.h"
#include "vec.h"
#include "config.h"
#include "enc/util.h"
#include "frame.h"
#include "enc/encoder.h"

#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>
#include <zlib.h>
#include <pixels.h>
#include <pthread.h>
#include <assert.h>
#include <aml.h>
#include <libdrm/drm_fourcc.h>
#ifdef HAVE_JPEG
#include <turbojpeg.h>
#endif

#define UDIV_UP(a, b) (((a) + (b) - 1) / (b))

#define TIGHT_FILL 0x80
#define TIGHT_JPEG 0x90
#define TIGHT_PNG 0xA0
#define TIGHT_BASIC 0x00

#define TIGHT_STREAM(n) ((n) << 4)
#define TIGHT_RESET(n) (1 << (n))

#define TSL 64 /* Tile Side Length */

#define MAX_TILE_SIZE (2 * TSL * TSL * 4)

struct encoder* tight_encoder_new(uint16_t width, uint16_t height);

typedef void (*tight_done_fn)(struct vec* frame, void*);

struct tight_encoder_grid {
	struct tight_tile* grid;
	uint32_t width;
	uint32_t height;
};

struct tight_encoder {
	struct encoder encoder;

	int last_n_fbs;
	uint32_t width;
	uint32_t height;
	int quality;

	struct tight_encoder_grid grid[NVNC_FB_COMPOSITE_MAX];

	z_stream zs[4];
	struct aml_work* zs_worker[4];

	struct rfb_pixel_format dfmt;

	struct rfb_pixel_format sfmt[NVNC_FB_COMPOSITE_MAX];
	struct nvnc_composite_fb composite_fb;

	uint64_t pts;

	uint32_t n_rects;
	uint32_t n_jobs;

	struct vec dst;

	tight_done_fn on_frame_done;
	void* userdata;
};

enum tight_tile_state {
	TIGHT_TILE_READY = 0,
	TIGHT_TILE_DAMAGED,
	TIGHT_TILE_ENCODED,
};

struct tight_tile {
	enum tight_tile_state state;
	size_t size;
	uint8_t type;
	char buffer[MAX_TILE_SIZE];
};

struct tight_zs_worker_ctx {
	struct tight_encoder* encoder;
	int index;
};

struct encoder_impl encoder_impl_tight;

static void do_tight_zs_work(struct aml_work*);
static void on_tight_zs_work_done(struct aml_work*);
static int schedule_tight_finish(struct tight_encoder* self);

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

static int tight_encoder_init_stream(z_stream* zs)
{
	int rc = deflateInit2(zs,
			/* compression level: */ 1,
			/*            method: */ Z_DEFLATED,
			/*       window bits: */ 15,
			/*         mem level: */ 9,
			/*          strategy: */ Z_DEFAULT_STRATEGY);
	return rc == Z_OK ? 0 : -1;
}

static inline struct tight_tile* tight_tile(struct tight_encoder* self,
		int fb_index, uint32_t x, uint32_t y)
{
	struct tight_encoder_grid* grid = &self->grid[fb_index];
	return &grid->grid[x + y * grid->width];
}

static inline uint32_t tight_tile_width(struct tight_encoder* self,
		int fb_index, uint32_t x)
{
	uint32_t width = self->composite_fb.fbs[fb_index]->width;
	return x + TSL > width ? width - x : TSL;
}

static inline uint32_t tight_tile_height(struct tight_encoder* self,
		int fb_index, uint32_t y)
{
	uint32_t height = self->composite_fb.fbs[fb_index]->height;
	return y + TSL > height ? height - y : TSL;
}

static int tight_init_zs_worker(struct tight_encoder* self, int index)
{
	struct tight_zs_worker_ctx* ctx = calloc(1, sizeof(*ctx));
	if (!ctx)
		return -1;

	ctx->encoder = self;
	ctx->index = index;

	self->zs_worker[index] =
		aml_work_new(do_tight_zs_work, on_tight_zs_work_done, ctx, free);
	if (!self->zs_worker[index])
		goto failure;

	return 0;

failure:
	free(ctx);
	return -1;
}

static void tight_encoder_resize(struct tight_encoder* self)
{
	struct nvnc_composite_fb *cfb = &self->composite_fb;
	uint16_t width, height;
	width = nvnc_composite_fb_width(cfb);
	height = nvnc_composite_fb_height(cfb);
	int n_fbs = cfb->n_fbs;

	if (self->width == width && self->height == height &&
			self->last_n_fbs == n_fbs)
		return;

	self->width = width;
	self->height = height;
	self->last_n_fbs = n_fbs;

	for (int i = 0; i < NVNC_FB_COMPOSITE_MAX && self->grid[i].grid; ++i) {
		free(self->grid[i].grid);
		self->grid[i].grid = NULL;
	}

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

		struct tight_encoder_grid *grid = &self->grid[i];

		grid->width = UDIV_UP(fb->width, TSL);
		grid->height = UDIV_UP(fb->height, TSL);

		grid->grid = calloc(grid->width * grid->height,
				sizeof(*grid->grid));
		nvnc_assert(grid->grid, "OOM");
	}
}

static int tight_encoder_init(struct tight_encoder* self, uint32_t width,
		uint32_t height)
{
	memset(self, 0, sizeof(*self));

	tight_encoder_init_stream(&self->zs[0]);
	tight_encoder_init_stream(&self->zs[1]);
	tight_encoder_init_stream(&self->zs[2]);
	tight_encoder_init_stream(&self->zs[3]);

	tight_init_zs_worker(self, 0);
	tight_init_zs_worker(self, 1);
	tight_init_zs_worker(self, 2);
	tight_init_zs_worker(self, 3);

	aml_require_workers(aml_get_default(), 1);

	self->pts = NVNC_NO_PTS;

	return 0;
}

static void tight_encoder_destroy(struct tight_encoder* self)
{
	aml_unref(self->zs_worker[3]);
	aml_unref(self->zs_worker[2]);
	aml_unref(self->zs_worker[1]);
	aml_unref(self->zs_worker[0]);

	deflateEnd(&self->zs[3]);
	deflateEnd(&self->zs[2]);
	deflateEnd(&self->zs[1]);
	deflateEnd(&self->zs[0]);

	for (int i = 0; i < NVNC_FB_COMPOSITE_MAX && self->grid[i].grid; ++i)
		free(self->grid[i].grid);
}

static int tight_apply_damage(struct tight_encoder* self,
		struct pixman_region16* damage)
{
	int n_damaged = 0;

	for (int fbi = 0; fbi < self->composite_fb.n_fbs; ++fbi) {
		struct nvnc_frame* fb = self->composite_fb.fbs[fbi];
		struct tight_encoder_grid *grid = &self->grid[fbi];

		for (uint32_t y = 0; y < grid->height; ++y) {
			for (uint32_t x = 0; x < grid->width; ++x) {
				struct pixman_box16 box = {
					.x1 = fb->x_off + x * TSL,
					.y1 = fb->y_off + y * TSL,
				};
				box.x2 = box.x1 + tight_tile_width(self, fbi, x);
				box.y2 = box.y1 + tight_tile_height(self, fbi, y);

				pixman_region_overlap_t overlap
					= pixman_region_contains_rectangle(damage, &box);

				if (overlap != PIXMAN_REGION_OUT) {
					++n_damaged;
					tight_tile(self, fbi, x, y)->state =
						TIGHT_TILE_DAMAGED;
				} else {
					tight_tile(self, fbi, x, y)->state =
						TIGHT_TILE_READY;
				}
			}
		}
	}

	return n_damaged;
}

static void tight_encode_size(struct vec* dst, size_t size)
{
	vec_fast_append_8(dst, (size & 0x7f) | ((size >= 128) << 7));
	if (size >= 128)
		vec_fast_append_8(dst, ((size >> 7) & 0x7f) | ((size >= 16384) << 7));
	if (size >= 16384)
		vec_fast_append_8(dst, (size >> 14) & 0xff);
}

static int tight_deflate(struct tight_tile* tile, void* src,
		size_t len, z_stream* zs, bool flush)
{
	zs->next_in = src;
	zs->avail_in = len;

	do {
		if (tile->size >= MAX_TILE_SIZE)
			return -1;

		zs->next_out = ((Bytef*)tile->buffer) + tile->size;
		zs->avail_out = MAX_TILE_SIZE - tile->size;

		int r = deflate(zs, flush ? Z_SYNC_FLUSH : Z_NO_FLUSH);
		if (r == Z_STREAM_ERROR)
			return -1;

		tile->size = zs->next_out - (Bytef*)tile->buffer;
	} while (zs->avail_out == 0);

	assert(zs->avail_in == 0);

	return 0;
}

static void tight_encode_tile_basic(struct tight_encoder* self,
		struct tight_tile* tile, int fb_index, uint32_t x,
		uint32_t y_start, uint32_t width, uint32_t height, int zs_index)
{
	z_stream* zs = &self->zs[zs_index];
	tile->type = TIGHT_BASIC | TIGHT_STREAM(zs_index);

	int bytes_per_cpixel = nvnc__calc_bytes_per_cpixel(&self->dfmt);
	assert(bytes_per_cpixel <= 4);
	uint8_t row[TSL * 4];

	struct rfb_pixel_format cfmt = { 0 };
	if (bytes_per_cpixel == 3)
		rfb_pixfmt_from_fourcc(&cfmt, DRM_FORMAT_XBGR8888);
	else
		memcpy(&cfmt, &self->dfmt, sizeof(cfmt));

	struct nvnc_frame* fb = self->composite_fb.fbs[fb_index];
	uint8_t* addr = nvnc_frame_get_addr(fb);
	int32_t bpp = self->sfmt[fb_index].bits_per_pixel / 8;
	int32_t byte_stride = nvnc_frame_get_stride(fb) * bpp;
	int32_t xoff = x * bpp;
	// TODO: Limit width and hight to the sides
	for (uint32_t y = y_start; y < y_start + height; ++y) {
		uint8_t* img = addr + xoff + y * byte_stride;
		pixel_to_cpixel(row, &cfmt, img, &self->sfmt[fb_index],
				bytes_per_cpixel, width);

		// TODO What to do if the buffer fills up?
		if (tight_deflate(tile, row, bytes_per_cpixel * width,
				zs, y == y_start + height - 1) < 0)
			abort();
	}

}

#ifdef HAVE_JPEG
static enum TJPF tight_get_jpeg_pixfmt(uint32_t fourcc)
{
	switch (fourcc) {
	case DRM_FORMAT_RGBA8888:
	case DRM_FORMAT_RGBX8888:
		return TJPF_XBGR;
	case DRM_FORMAT_BGRA8888:
	case DRM_FORMAT_BGRX8888:
		return TJPF_XRGB;
	case DRM_FORMAT_ARGB8888:
	case DRM_FORMAT_XRGB8888:
		return TJPF_BGRX;
	case DRM_FORMAT_ABGR8888:
	case DRM_FORMAT_XBGR8888:
		return TJPF_RGBX;
	case DRM_FORMAT_BGR888:
		return TJPF_RGB;
	case DRM_FORMAT_RGB888:
		return TJPF_BGR;
	}

	return TJPF_UNKNOWN;
}

static int tight_encode_tile_jpeg(struct tight_encoder* self,
		struct tight_tile* tile, int fb_index, uint32_t x, uint32_t y,
		uint32_t width, uint32_t height)
{
	tile->type = TIGHT_JPEG;

	unsigned char* buffer = NULL;
	unsigned long size = 0;

	int quality = 11 * self->quality + 1;

	struct nvnc_frame* fb = self->composite_fb.fbs[fb_index];
	uint32_t fourcc = nvnc_frame_get_fourcc_format(fb);
	enum TJPF tjfmt = tight_get_jpeg_pixfmt(fourcc);
	if (tjfmt == TJPF_UNKNOWN)
		return -1;

	tjhandle handle = tjInitCompress();
	if (!handle)
		return -1;

	uint8_t* addr = nvnc_frame_get_addr(fb);
	int32_t bpp = self->sfmt[fb_index].bits_per_pixel / 8;
	int32_t byte_stride = nvnc_frame_get_stride(fb) * bpp;
	int32_t xoff = x * bpp;
	uint8_t* img = addr + xoff + y * byte_stride;

	enum TJSAMP subsampling = (self->quality == 9) ? TJSAMP_444 : TJSAMP_420;

	int rc = -1;
	rc = tjCompress2(handle, img, width, byte_stride, height, tjfmt, &buffer,
			&size, subsampling, quality, TJFLAG_FASTDCT);
	if (rc < 0) {
		nvnc_log(NVNC_LOG_ERROR, "Failed to encode tight JPEG box: %s",
				tjGetErrorStr());
		goto failure;
	}

	if (size > MAX_TILE_SIZE) {
		nvnc_log(NVNC_LOG_ERROR, "Whoops, encoded JPEG was too big for the buffer");
		goto failure;
	}

	memcpy(tile->buffer, buffer, size);
	tile->size = size;

	rc = 0;
	tjFree(buffer);
failure:
	tjDestroy(handle);

	return rc;
}
#endif /* HAVE_JPEG */

static void tight_encode_tile(struct tight_encoder* self, int fb_index,
		uint32_t gx, uint32_t gy)
{
	struct tight_tile* tile = tight_tile(self, fb_index, gx, gy);

	uint32_t x = gx * TSL;
	uint32_t y = gy * TSL;

	uint32_t width = tight_tile_width(self, fb_index, x);
	uint32_t height = tight_tile_height(self, fb_index, y);

	tile->size = 0;

#ifdef HAVE_JPEG
	if (self->quality >= 10) {
		tight_encode_tile_basic(self, tile, fb_index, x, y, width,
				height, gx % 4);
	} else {
		tight_encode_tile_jpeg(self, tile, fb_index, x, y, width,
				height);
	}
#else
	tight_encode_tile_basic(self, tile, fb_index, x, y, width, height, gx % 4);
#endif

	tile->state = TIGHT_TILE_ENCODED;
}

static void do_tight_zs_work(struct aml_work* work)
{
	struct tight_zs_worker_ctx* ctx = aml_get_userdata(work);
	struct tight_encoder* self = ctx->encoder;
	int index = ctx->index;

	for (int fbi = 0; fbi < self->composite_fb.n_fbs; ++fbi)
		for (uint32_t y = 0; y < self->grid[fbi].height; ++y)
			for (uint32_t x = index; x < self->grid[fbi].width; x += 4)
				if (tight_tile(self, fbi, x, y)->state == TIGHT_TILE_DAMAGED)
					tight_encode_tile(self, fbi, x, y);
}

static void on_tight_zs_work_done(struct aml_work* obj)
{
	struct tight_zs_worker_ctx* ctx = aml_get_userdata(obj);
	struct tight_encoder* self = ctx->encoder;

	if (--self->n_jobs == 0) {
		schedule_tight_finish(self);
	}

	encoder_unref(&self->encoder);
}

static int tight_schedule_zs_work(struct tight_encoder* self, int index)
{
	encoder_ref(&self->encoder);

	int rc = aml_start(aml_get_default(), self->zs_worker[index]);
	if (rc >= 0)
		++self->n_jobs;
	else
		encoder_unref(&self->encoder);

	return rc;
}

static int tight_schedule_encoding_jobs(struct tight_encoder* self)
{
	for (int i = 0; i < 4; ++i)
		if (tight_schedule_zs_work(self, i) < 0)
			return -1;

	return 0;
}

static void tight_finish_tile(struct tight_encoder* self,
		int fb_index, uint32_t gx, uint32_t gy)
{
	struct tight_tile* tile = tight_tile(self, fb_index, gx, gy);

	struct nvnc_frame* fb = self->composite_fb.fbs[fb_index];
	uint16_t x_pos = fb->x_off;
	uint16_t y_pos = fb->y_off;

	uint32_t x = gx * TSL;
	uint32_t y = gy * TSL;

	uint32_t width = tight_tile_width(self, fb_index, x);
	uint32_t height = tight_tile_height(self, fb_index, y);

	nvnc__encode_rect_head(&self->dst, RFB_ENCODING_TIGHT, x_pos + x, y_pos + y,
			width, height);

	vec_append(&self->dst, &tile->type, sizeof(tile->type));
	tight_encode_size(&self->dst, tile->size);
	vec_append(&self->dst, tile->buffer, tile->size);

	tile->state = TIGHT_TILE_READY;
}

static void tight_finish(struct tight_encoder* self)
{
	for (int fbi = 0; fbi < self->composite_fb.n_fbs; ++fbi)
		for (uint32_t y = 0; y < self->grid[fbi].height; ++y)
			for (uint32_t x = 0; x < self->grid[fbi].width; ++x)
				if (tight_tile(self, fbi, x, y)->state ==
						TIGHT_TILE_ENCODED)
					tight_finish_tile(self, fbi, x, y);
}

static void do_tight_finish(struct aml_work* work)
{
	struct tight_encoder* self = aml_get_userdata(work);
	tight_finish(self);
}

static void on_tight_finished(struct aml_work* work)
{
	struct tight_encoder* self = aml_get_userdata(work);

	struct nvnc_frame_metadata* metadata = self->composite_fb.metadata;
	nvnc_frame_metadata_ref(metadata);

	nvnc_composite_fb_unref(&self->composite_fb);
	memset(&self->composite_fb, 0, sizeof(self->composite_fb));

	struct encoded_frame* result;
	result = nvnc__encoded_frame_new(self->dst.data, self->dst.len,
			self->n_rects, self->width, self->height, self->pts);
	assert(result);

	result->metadata = metadata;

	encoder_finish_frame(&self->encoder, result);

	self->pts = NVNC_NO_PTS;
	nvnc_frame_metadata_unref(metadata);
	encoded_frame_unref(result);
	encoder_unref(&self->encoder);
}

static int schedule_tight_finish(struct tight_encoder* self)
{
	encoder_ref(&self->encoder);

	struct aml_work* work = aml_work_new(do_tight_finish, on_tight_finished,
			self, NULL);
	if (!work) {
		encoder_unref(&self->encoder);
		return -1;
	}

	int rc = aml_start(aml_get_default(), work);
	aml_unref(work);
	return rc;
}

struct encoder* tight_encoder_new(uint16_t width, uint16_t height)
{
	struct tight_encoder* self = calloc(1, sizeof(*self));
	if (!self)
		return NULL;

	if (tight_encoder_init(self, width, height) < 0) {
		free(self);
		return NULL;
	}

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

	return (struct encoder*)self;
}

static void tight_encoder_destroy_wrapper(struct encoder* encoder)
{
	tight_encoder_destroy(tight_encoder(encoder));
	free(encoder);
}

static void tight_encoder_set_output_format(struct encoder* encoder,
		const struct rfb_pixel_format* pixfmt)
{
	struct tight_encoder* self = tight_encoder(encoder);
	memcpy(&self->dfmt, pixfmt, sizeof(self->dfmt));
}

static void tight_encoder_set_quality(struct encoder* encoder, int value)
{
	struct tight_encoder* self = tight_encoder(encoder);
	self->quality = value;
}

static int tight_encoder_encode(struct encoder* encoder,
		struct nvnc_composite_fb* composite_fb,
		struct pixman_region16* damage)
{
	struct tight_encoder* self = tight_encoder(encoder);
	int rc;

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

		rc = rfb_pixfmt_from_fourcc(&self->sfmt[i],
				nvnc_frame_get_fourcc_format(fb));
		nvnc_assert(rc == 0, "Unhandled pixel format for input buffer");
	}

	nvnc_composite_fb_copy(&self->composite_fb, composite_fb);
	self->pts = nvnc_composite_fb_pts(composite_fb);

	tight_encoder_resize(self);

	rc = nvnc_composite_fb_map(composite_fb);
	nvnc_assert(rc == 0, "Failed to map input buffer");

	// TODO: Estimate a better buffer size
	rc = vec_init(&self->dst, self->width * self->height * 4);
	if (rc < 0)
		return -1;

	self->n_rects = tight_apply_damage(self, damage);
	assert(self->n_rects > 0);

	rc = tight_schedule_encoding_jobs(self);
	nvnc_assert(rc == 0, "Failed to schedule encoding jobs");

	return 0;
}

struct encoder_impl encoder_impl_tight = {
	.destroy = tight_encoder_destroy_wrapper,
	.set_output_format = tight_encoder_set_output_format,
	.set_quality = tight_encoder_set_quality,
	.encode = tight_encoder_encode,
};
