#include "sgu_stdafx.h"
#include "sgu_mode_make_diff.h"
#include "sgu_utils.h"

extern "C"
{
#include "bsdiff.h"
}

#include "bzlib.h"
#include "zstd.h"

static int bzPatchFileWriter(struct bsdiff_stream* stream, const void* buffer, int size)
{
	BZFILE* bz2 = (BZFILE*)stream->opaque;

	int bz2err;
	BZ2_bzWrite(&bz2err, bz2, (void*)buffer, size);

	if (bz2err != BZ_STREAM_END && bz2err != BZ_OK)
		return -1;

	return 0;
}

struct ZSTD_CompressContext
{
	ZSTD_CCtx* cctx;

	void* inputBuf;
	size_t inputBufSize;
	size_t inputBufDataSize;

	void* outputBuf;
	size_t outputBufSize;

	FILE* fpOutput;

	FILE* fpDebugOutput;
	FILE* fpDebugOutput2;
};

ZSTD_CompressContext* createZSTDCompressContext(FILE* fpOutput)
{
	ZSTD_CompressContext* ctx = new ZSTD_CompressContext();
	if (ctx == nullptr)
	{
		return nullptr;
	}
	ctx->cctx = ZSTD_createCCtx();
	if (ctx->cctx == nullptr)
	{
		return nullptr;
	}

	ZSTD_CCtx_setParameter(ctx->cctx, ZSTD_c_compressionLevel, ZSTD_CLEVEL_DEFAULT);
	ZSTD_CCtx_setParameter(ctx->cctx, ZSTD_c_checksumFlag, 1);

	ctx->inputBufSize = ZSTD_CStreamInSize();
	ctx->inputBuf = malloc(ctx->inputBufSize);
	if (ctx->inputBuf == nullptr)
	{
		return nullptr;
	}
	ctx->inputBufDataSize = 0;

	ctx->outputBufSize = ZSTD_CStreamOutSize();
	ctx->outputBuf = malloc(ctx->outputBufSize);
	if (ctx->outputBuf == nullptr)
	{
		return nullptr;
	}

	ctx->fpOutput = fpOutput;

	ctx->fpDebugOutput = fopen("org.patch", "wb");
	ctx->fpDebugOutput2 = fopen("org.patch.2", "wb");
	return ctx;
}

void destoryZSTDCompressContext(ZSTD_CompressContext* ctx)
{
	fclose(ctx->fpDebugOutput);
	fclose(ctx->fpDebugOutput2);

	free(ctx->inputBuf);
	free(ctx->outputBuf);
	delete ctx;
}

void zstdCompressData(ZSTD_CompressContext* ctx, ZSTD_EndDirective mode)
{
	/* Set the input buffer to what we just read.
	 * We compress until the input buffer is empty, each time flushing the
	 * output.
	 */
	ZSTD_inBuffer input = { ctx->inputBuf, ctx->inputBufDataSize, 0 };
	bool finished;
	do {
		/* Compress into the output buffer and write all of the output to
		 * the file so we can reuse the buffer next iteration.
		 */
		ZSTD_outBuffer output = { ctx->outputBuf, ctx->outputBufSize, 0 };
		const size_t remaining = ZSTD_compressStream2(ctx->cctx, &output, &input, mode);

		//write to output file
		fwrite(ctx->outputBuf, 1, output.pos, ctx->fpOutput);

		fwrite(ctx->outputBuf, 1, output.pos, ctx->fpDebugOutput2);

		/* If we're on the last chunk we're finished when zstd returns 0,
		 * which means its consumed all the input AND finished the frame.
		 * Otherwise, we're finished when we've consumed all the input.
		 */
		finished = (mode== ZSTD_e_end) ? (remaining == 0) : (input.pos == input.size);
	} while (!finished);
	ctx->inputBufDataSize = 0;
}

void zstdPatchFileClose(ZSTD_CompressContext* ctx)
{
	zstdCompressData(ctx, ZSTD_e_end);
}

static int zstdPatchFileWriter(struct bsdiff_stream* stream, const void* inputData, int inputSize)
{
	ZSTD_CompressContext* ctx = (ZSTD_CompressContext*)(stream->opaque);
	
	fwrite(inputData, 1, inputSize, ctx->fpDebugOutput);

	size_t readedSize = 0;

	for (;;) {
		size_t inputBufRemainSize = ctx->inputBufSize - ctx->inputBufDataSize;
		size_t remainDataSize = (size_t)inputSize - readedSize;
		size_t readSize = inputBufRemainSize<remainDataSize ? inputBufRemainSize : remainDataSize;
		memcpy(((uint8_t*)ctx->inputBuf)+ctx->inputBufDataSize, (const uint8_t*)inputData + readedSize, readSize);
		readedSize += readSize;
		ctx->inputBufDataSize += readSize;

		if (ctx->inputBufDataSize >= ctx->inputBufSize)
		{
			zstdCompressData(ctx, ZSTD_e_continue);
		}

		if (readedSize >= (size_t)inputSize)
		{
			break;
		}
	}

	return 0;
}

static int patchFileWriter(struct bsdiff_stream* stream, const void* inputData, int inputSize)
{
	FILE* fpOutput = (FILE*)(stream->opaque);
	fwrite(inputData, 1, inputSize, fpOutput);
	return 0;
}

bool createDiff(const char* inputFilename1, const char* inputFilename2, const char* outputFilename, int32_t compressMethod)
{
	if (!inputFilename1 || inputFilename1[0] == 0 || !inputFilename2 || inputFilename2[0] == 0)
	{
		printf("Error: Input filename is not specified.\n");
		return false;
	}
	if (!outputFilename || outputFilename[0] == 0)
	{
		printf("Error: Output filename is not specified.\n");
		return false;
	}

	bool result = false;
	uint8_t* pBuffer1 = nullptr;
	uint8_t* pBuffer2 = nullptr;
	FILE* fpOutput = nullptr;
	do {
		int64_t fileSize1;
		pBuffer1 = readFileToBuffer(inputFilename1, fileSize1);
		if (pBuffer1 == nullptr)
		{
			printf("Error: Failed to read input file 1: %s\n", inputFilename1);
			break;
		}

		int64_t fileSize2;
		pBuffer2 = readFileToBuffer(inputFilename2, fileSize2);
		if (pBuffer2 == nullptr)
		{
			printf("Error: Failed to read input file 2: %s\n", inputFilename2);
			break;
		}

		// Open the output file
		fopen_s(&fpOutput, outputFilename, "wb");
		if (fpOutput == nullptr)
		{
			printf("Error: Failed to open output file: %s\n", outputFilename);
			break;
		}

		// Write the header
		if (fwrite("ENDSLEY/BSDIFF43", 16, 1, fpOutput) != 1 ||
			fwrite(&fileSize2, sizeof(fileSize2), 1, fpOutput) != 1)
		{
			printf("Failed to write header\n");
			break;
		}

		BZFILE* bzHandle = nullptr;
		ZSTD_CompressContext* zstdContext = nullptr;
		if (compressMethod == 1)
		{
			int bz2err;
			bzHandle = BZ2_bzWriteOpen(&bz2err, fpOutput, 9, 0, 0);
			if (bzHandle == nullptr || bz2err != BZ_OK)
			{
				printf("BZ2_bzWriteOpen failed, bz2err=%d\n", bz2err);
				break;
			}
		}
		else if (compressMethod == 2)
		{
			zstdContext = createZSTDCompressContext(fpOutput);
			if (zstdContext == nullptr)
			{
				printf("Failed create zstd context!");
				break;
			}
		}
		else
		{

		}
		//generate diff
		struct bsdiff_stream stream;
		stream.malloc = malloc;
		stream.free = free;
		if (compressMethod == 1)
		{
			stream.write = bzPatchFileWriter;
			stream.opaque = bzHandle;
		}
		else if (compressMethod == 2)
		{
			stream.write = zstdPatchFileWriter;
			stream.opaque = zstdContext;
		}
		else
		{
			stream.write = patchFileWriter;
			stream.opaque = fpOutput;
		}

		if (bsdiff(pBuffer1, fileSize1, pBuffer2, fileSize2, &stream) !=0 )
		{
			printf("bsdiff error!\n");
			break;
		}

		if (compressMethod == 1)
		{
			int bz2err;
			BZ2_bzWriteClose(&bz2err, bzHandle, 0, NULL, NULL);
			if (bz2err != BZ_OK)
			{
				printf("BZ2_bzWriteClose failed, bz2err=%d", bz2err);
				break;
			}
			bzHandle = nullptr;
		}
		else if (compressMethod == 2)
		{
			zstdPatchFileClose(zstdContext);
			destoryZSTDCompressContext(zstdContext);
			zstdContext = nullptr;
		}

		result = true;
		printf("Diff file created successfully: %s\n", outputFilename);
	} while (false);

	if (fpOutput != nullptr)
	{
		fclose(fpOutput);
		fpOutput = nullptr;
	}
	freeFileBuffer(pBuffer1);
	pBuffer1 = nullptr;

	freeFileBuffer(pBuffer2);
	pBuffer2 = nullptr;

	return result;
}