// Code based on DOSBox-0.65

#include "ZMBVEncoder.hh"

#include "FrameSource.hh"
#include "PixelOperations.hh"

#include "cstd.hh"
#include "endian.hh"
#include "narrow.hh"
#include "unreachable.hh"

#include <algorithm>
#include <array>
#include <bit>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <tuple>

namespace openmsx {

static constexpr uint8_t DBZV_VERSION_HIGH = 0;
static constexpr uint8_t DBZV_VERSION_LOW = 1;
static constexpr uint8_t COMPRESSION_ZLIB = 1;
static constexpr unsigned MAX_VECTOR = 16;
static constexpr unsigned BLOCK_WIDTH  = MAX_VECTOR;
static constexpr unsigned BLOCK_HEIGHT = MAX_VECTOR;
static constexpr unsigned FLAG_KEYFRAME = 0x01;

struct CodecVector {
	int8_t x;
	int8_t y;
};

static constexpr unsigned VECTOR_TAB_SIZE =
	1 +                                       // center
	8 * MAX_VECTOR +                          // horizontal, vertical, diagonal
	MAX_VECTOR * MAX_VECTOR - 2 * MAX_VECTOR; // rest (only MAX_VECTOR/2)

static constexpr auto vectorTable = [] {
	std::array<CodecVector, VECTOR_TAB_SIZE> result = {};

	unsigned p = 0;
	// center
	result[p] = {.x = 0, .y = 0};
	p += 1;
	// horizontal, vertical, diagonal
	for (int i = 1; i <= int(MAX_VECTOR); ++i, p += 8) {
		result[p + 0] = {.x = int8_t( i), .y = int8_t( 0)};
		result[p + 1] = {.x = int8_t(-i), .y = int8_t( 0)};
		result[p + 2] = {.x = int8_t( 0), .y = int8_t( i)};
		result[p + 3] = {.x = int8_t( 0), .y = int8_t(-i)};
		result[p + 4] = {.x = int8_t( i), .y = int8_t( i)};
		result[p + 5] = {.x = int8_t(-i), .y = int8_t( i)};
		result[p + 6] = {.x = int8_t( i), .y = int8_t(-i)};
		result[p + 7] = {.x = int8_t(-i), .y = int8_t(-i)};
	}
	// rest
	for (int y = 1; y <= int(MAX_VECTOR / 2); ++y) {
		for (int x = 1; x <= int(MAX_VECTOR / 2); ++x) {
			if (x == y) continue; // already have diagonal
			result[p + 0] = {.x = int8_t( x), .y = int8_t( y)};
			result[p + 1] = {.x = int8_t(-x), .y = int8_t( y)};
			result[p + 2] = {.x = int8_t( x), .y = int8_t(-y)};
			result[p + 3] = {.x = int8_t(-x), .y = int8_t(-y)};
			p += 4;
		}
	}
	assert(p == VECTOR_TAB_SIZE);

	// sort
	auto compare = [](const CodecVector& l, const CodecVector& r) {
		auto cost = [](const CodecVector& v) {
			auto c = cstd::sqrt(double(v.x * v.x + v.y * v.y));
			if ((v.x == 0) || (v.y == 0)) {
				// no penalty for purely horizontal/vertical offset
				c *= 1.0;
			} else if (cstd::abs(v.x) == cstd::abs(v.y)) {
				// small penalty for pure diagonal
				c *= 2.0;
			} else {
				// bigger penalty for 'random' direction
				c *= 4.0;
			}
			return c;
		};
		return std::tuple(cost(l), l.x, l.y) <
		       std::tuple(cost(r), r.x, r.y);
	};
	std::ranges::sort(result, compare);

	return result;
}();

struct KeyframeHeader {
	uint8_t high_version;
	uint8_t low_version;
	uint8_t compression;
	uint8_t format;
	uint8_t blockWidth;
	uint8_t blockHeight;
};


static inline void writePixel(
	unsigned pixel, Endian::L32& dest)
{
	PixelOperations pixelOps;
	unsigned r = pixelOps.red(pixel);
	unsigned g = pixelOps.green(pixel);
	unsigned b = pixelOps.blue(pixel);
	dest = (r << 16) | (g <<  8) |  b;
}


ZMBVEncoder::ZMBVEncoder(unsigned width_, unsigned height_)
	: width(width_)
	, height(height_)
{
	setupBuffers();
	memset(&zstream, 0, sizeof(zstream));
	deflateInit(&zstream, 6); // compression level

	// I did a small test: compression level vs compression speed
	//  (recorded Space Manbow intro, video only)
	//
	// level |  time  | size
	// ------+--------+----------
	//   0   | 1m12.6 | 139442594
	//   1   | 1m12.1 |   5217288
	//   2   | 1m10.8 |   4887258
	//   3   | 1m11.8 |   4610668
	//   4   | 1m13.1 |   3791932  <-- old default
	//   5   | 1m14.2 |   3602078
	//   6   | 1m14.5 |   3363766  <-- current default
	//   7   | 1m15.8 |   3333938
	//   8   | 1m25.0 |   3301168
	//   9   | 2m04.1 |   3253706
	//
	// Level 6 seems a good compromise between size/speed for THIS test.
}

void ZMBVEncoder::setupBuffers()
{
	static constexpr size_t pixelSize = sizeof(Pixel);

	pitch = width + 2 * MAX_VECTOR;
	auto bufSize = (height + 2 * MAX_VECTOR) * pitch * pixelSize + 2048;

	oldFrame.resize(bufSize);
	newFrame.resize(bufSize);
	std::ranges::fill(std::span{oldFrame}, 0);
	std::ranges::fill(std::span{newFrame}, 0);
	work.resize(bufSize);
	outputSize = neededSize();
	output.resize(outputSize);

	assert((width  % BLOCK_WIDTH ) == 0);
	assert((height % BLOCK_HEIGHT) == 0);
	size_t xBlocks = width / BLOCK_WIDTH;
	size_t yBlocks = height / BLOCK_HEIGHT;
	blockOffsets.resize(xBlocks * yBlocks);
	for (auto y : xrange(yBlocks)) {
		for (auto x : xrange(xBlocks)) {
			blockOffsets[y * xBlocks + x] =
				((y * BLOCK_HEIGHT) + MAX_VECTOR) * pitch +
				(x * BLOCK_WIDTH) + MAX_VECTOR;
		}
	}
}

unsigned ZMBVEncoder::neededSize() const
{
	static constexpr unsigned pixelSize = sizeof(Pixel);
	unsigned f = pixelSize * width * height + 2 * (1 + (width / 8)) * (1 + (height / 8)) + 1024;
	return f + f / 1000;
}

unsigned ZMBVEncoder::possibleBlock(int vx, int vy, size_t offset)
{
	int ret = 0;
	const auto* pOld = &(std::bit_cast<const Pixel*>(oldFrame.data()))[offset + (vy * pitch) + vx];
	const auto* pNew = &(std::bit_cast<const Pixel*>(newFrame.data()))[offset];
	for (unsigned y = 0; y < BLOCK_HEIGHT; y += 4) {
		for (unsigned x = 0; x < BLOCK_WIDTH; x += 4) {
			if (pOld[x] != pNew[x]) ++ret;
		}
		pOld += pitch * 4;
		pNew += pitch * 4;
	}
	return ret;
}

unsigned ZMBVEncoder::compareBlock(int vx, int vy, size_t offset)
{
	int ret = 0;
	const auto* pOld = &(std::bit_cast<const Pixel*>(oldFrame.data()))[offset + (vy * pitch) + vx];
	const auto* pNew = &(std::bit_cast<const Pixel*>(newFrame.data()))[offset];
	repeat(BLOCK_HEIGHT, [&] {
		for (auto x : xrange(BLOCK_WIDTH)) {
			if (pOld[x] != pNew[x]) ++ret;
		}
		pOld += pitch;
		pNew += pitch;
	});
	return ret;
}

void ZMBVEncoder::addXorBlock(int vx, int vy, size_t offset, unsigned& workUsed)
{
	using LE_P = typename Endian::Little<Pixel>::type;

	const auto* pOld = &(std::bit_cast<const Pixel*>(oldFrame.data()))[offset + (vy * pitch) + vx];
	const auto* pNew = &(std::bit_cast<const Pixel*>(newFrame.data()))[offset];
	repeat(BLOCK_HEIGHT, [&] {
		for (auto x : xrange(BLOCK_WIDTH)) {
			auto pXor = pNew[x] ^ pOld[x];
			writePixel(pXor, *std::bit_cast<LE_P*>(&work[workUsed]));
			workUsed += sizeof(Pixel);
		}
		pOld += pitch;
		pNew += pitch;
	});
}

void ZMBVEncoder::addXorFrame(unsigned& workUsed)
{
	auto* vectors = std::bit_cast<int8_t*>(&work[workUsed]);

	unsigned xBlocks = width / BLOCK_WIDTH;
	unsigned yBlocks = height / BLOCK_HEIGHT;
	unsigned blockCount = xBlocks * yBlocks;

	// Align the following xor data on 4 byte boundary
	workUsed = (workUsed + blockCount * 2 + 3) & ~3;

	int bestVx = 0;
	int bestVy = 0;
	for (auto b : xrange(blockCount)) {
		auto offset = blockOffsets[b];
		// first try best vector of previous block
		unsigned bestChange = compareBlock(bestVx, bestVy, offset);
		if (bestChange >= 4) {
			int possibles = 64;
			for (const auto& v : vectorTable) {
				if (possibleBlock(v.x, v.y, offset) < 4) {
					if (auto testChange = compareBlock(v.x, v.y, offset);
					    testChange < bestChange) {
						bestChange = testChange;
						bestVx = narrow<int>(v.x);
						bestVy = narrow<int>(v.y);
						if (bestChange < 4) break;
					}
					--possibles;
					if (possibles == 0) break;
				}
			}
		}
		vectors[b * 2 + 0] = narrow<int8_t>(bestVx << 1);
		vectors[b * 2 + 1] = narrow<int8_t>(bestVy << 1);
		if (bestChange) {
			vectors[b * 2 + 0] |= 1;
			addXorBlock(bestVx, bestVy, offset, workUsed);
		}
	}
}

void ZMBVEncoder::addFullFrame(unsigned& workUsed)
{
	using LE_P = typename Endian::Little<Pixel>::type;
	static constexpr size_t pixelSize = sizeof(Pixel);

	auto* readFrame =
		&newFrame[pixelSize * (MAX_VECTOR + MAX_VECTOR * pitch)];
	repeat(height, [&] {
		const auto* pixelsIn = std::bit_cast<const Pixel*>(readFrame);
		auto* pixelsOut = std::bit_cast<LE_P*>(&work[workUsed]);
		for (auto x : xrange(width)) {
			writePixel(pixelsIn[x], pixelsOut[x]);
		}
		readFrame += pitch * sizeof(Pixel);
		workUsed += narrow<unsigned>(width * sizeof(Pixel));
	});
}

const ZMBVEncoder::Pixel* ZMBVEncoder::getScaledLine(const FrameSource* frame, unsigned y, Pixel* workBuf) const
{
	switch (height) {
	case 240:
		return frame->getLinePtr320_240(y, std::span<uint32_t, 320>(workBuf, 320)).data();
	case 480:
		return frame->getLinePtr640_480(y, std::span<uint32_t, 640>(workBuf, 640)).data();
	case 720:
		return frame->getLinePtr960_720(y, std::span<uint32_t, 960>(workBuf, 960)).data();
	default:
		UNREACHABLE;
	}
}

std::span<const uint8_t> ZMBVEncoder::compressFrame(bool keyFrame, const FrameSource* frame)
{
	std::swap(newFrame, oldFrame); // replace oldFrame with newFrame

	// Reset the work buffer
	unsigned workUsed = 0;
	unsigned writeDone = 1;
	uint8_t* writeBuf = output.data();

	output[0] = 0; // first byte contains info about this frame
	if (keyFrame) {
		static const uint8_t ZMBV_FORMAT_32BPP = 8;

		output[0] |= FLAG_KEYFRAME;
		auto* header = std::bit_cast<KeyframeHeader*>(
			writeBuf + writeDone);
		header->high_version = DBZV_VERSION_HIGH;
		header->low_version = DBZV_VERSION_LOW;
		header->compression = COMPRESSION_ZLIB;
		header->format = ZMBV_FORMAT_32BPP;
		header->blockWidth = BLOCK_WIDTH;
		header->blockHeight = BLOCK_HEIGHT;
		writeDone += sizeof(KeyframeHeader);
		deflateReset(&zstream); // restart deflate
	}

	// copy lines (to add black border)
	static constexpr size_t pixelSize = sizeof(Pixel);
	auto linePitch = pitch * pixelSize;
	auto lineWidth = size_t(width) * pixelSize;
	uint8_t* dest =
		&newFrame[pixelSize * (MAX_VECTOR + MAX_VECTOR * pitch)];
	for (auto i : xrange(height)) {
		const auto* scaled = std::bit_cast<const uint8_t*>(
			getScaledLine(frame, i, std::bit_cast<Pixel*>(dest)));
		if (scaled != dest) memcpy(dest, scaled, lineWidth);
		dest += linePitch;
	}

	// Add the frame data.
	if (keyFrame) {
		// Key frame: full frame data.
		addFullFrame(workUsed);
	} else {
		// Non-key frame: delta frame data.
		addXorFrame(workUsed);
	}
	// Compress the frame data with zlib.
	zstream.next_in = work.data();
	zstream.avail_in = workUsed;
	zstream.total_in = 0;

	zstream.next_out = std::bit_cast<Bytef*>(writeBuf + writeDone);
	zstream.avail_out = outputSize - writeDone;
	zstream.total_out = 0;
	auto r = deflate(&zstream, Z_SYNC_FLUSH);
	assert(r == Z_OK); (void)r;

	return {output.data(), writeDone + zstream.total_out};
}

} // namespace openmsx