#include "sgu_stdafx.h"
#include "sgu_mode_apply_patch.h"
#include "sgu_utils.h"

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

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


//--------------------------- Not Compress-----------------------------------
static int patchFileReader(const struct bspatch_stream* stream, void* buffer, int length)
{
	FILE* fp = (FILE*)stream->opaque;
	size_t n = fread(buffer, 1, (size_t)length, fp);
	if (n != length)
		return -1;

	return 0;
}

//--------------------------- BZ2 -----------------------------------
static int bzPatchFileReader(const struct bspatch_stream* stream, void* buffer, int length)
{
	BZFILE* bz2 = (BZFILE*)stream->opaque;

	int bz2err;
	int n = BZ2_bzRead(&bz2err, bz2, buffer, length);
	if (n != length)
		return -1;

	return 0;
}

//--------------------------- ZSTD -----------------------------------
struct ZSTD_DecompressContext
{
	ZSTD_DCtx* dctx;

	FILE* fpInput;

	void* inputBuf;
	size_t inputBufSize;
	size_t inputBufDataSize;
	size_t inputBufDataReadedPos;

	void* outputBuf;
	size_t outputBufSize;
	size_t outputBufDataSize;
	size_t outputBufDataUploadedPos;
};

ZSTD_DecompressContext* createZSTDDecompressContext(FILE* fpInput)
{
	ZSTD_DecompressContext* ctx = new ZSTD_DecompressContext();
	ctx->dctx = ZSTD_createDCtx();
	ctx->fpInput = fpInput;

	ctx->inputBufSize = ZSTD_DStreamInSize();
	ctx->inputBuf = malloc(ctx->inputBufSize);
	ctx->inputBufDataSize = 0;
	ctx->inputBufDataReadedPos = 0;

	ctx->outputBufSize = ZSTD_DStreamOutSize();
	ctx->outputBuf = malloc(ctx->outputBufSize);
	ctx->outputBufDataSize = 0;
	ctx->outputBufDataUploadedPos = 0;

	return ctx;
}

void destoryZSTDDecompressContext(ZSTD_DecompressContext* ctx)
{
	free(ctx->inputBuf);
	free(ctx->outputBuf);
	delete ctx;
}

static int zstdPatchFileReader(const struct bspatch_stream* stream, void* buffer, int length)
{
	ZSTD_DecompressContext* ctx = (ZSTD_DecompressContext*)(stream->opaque);
	if (length == 0)
	{
		return 0;
	}

	//total uploaded size
	size_t uploadedSize = 0;

	do {
		//upload all decompressed data
		if (ctx->outputBufDataSize > ctx->outputBufDataUploadedPos)
		{
			size_t remainDecompressedDataSize = ctx->outputBufDataSize - ctx->outputBufDataUploadedPos;
			size_t remainNotUploadedSize = (size_t)length - uploadedSize;

			size_t uploadSize = remainNotUploadedSize < remainDecompressedDataSize ? remainNotUploadedSize : remainDecompressedDataSize;
			if (uploadSize > 0)
			{
				memcpy((uint8_t*)buffer + uploadedSize, ((uint8_t*)ctx->outputBuf) + ctx->outputBufDataUploadedPos, uploadSize);

				uploadedSize += uploadSize;
				ctx->outputBufDataUploadedPos += uploadSize;
			}
		}

		if (uploadedSize >= length)
		{
			return 0;
		}
		
		//compressed data empty
		if (ctx->inputBufDataSize == ctx->inputBufDataReadedPos)
		{
			//read file
			ctx->inputBufDataSize = fread(ctx->inputBuf, 1, ctx->inputBufSize, ctx->fpInput);
			ctx->inputBufDataReadedPos = 0;

			if (ctx->inputBufDataSize == 0)
			{
				//TODO:
				return -1;
			}
		}

		//decompress data
		ZSTD_inBuffer input = { ctx->inputBuf, ctx->inputBufDataSize, ctx->inputBufDataReadedPos };
		ZSTD_outBuffer output = { ctx->outputBuf, ctx->outputBufSize, 0 };

		const size_t ret = ZSTD_decompressStream(ctx->dctx, &output, &input);
		if (ZSTD_isError(ret))
		{
			//TODO:
			return -1;
		}
		ctx->inputBufDataReadedPos = input.pos;

		ctx->outputBufDataUploadedPos = 0;
		ctx->outputBufDataSize = output.pos;
	} while (true);

	return 0;
}

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

	uint8_t* pInputBuffer = nullptr;
	FILE* patchFile = nullptr;
	FILE* outputFile = nullptr;
	uint8_t* pOutputBuffer = nullptr;

	bool result = false;
	do {
		// Open the input file
		int64_t inputFileSize;
		pInputBuffer = readFileToBuffer(inputFilename, inputFileSize);
		if (pInputBuffer == nullptr)
		{
			printf("Error: Failed to read input file: %s\n", inputFilename);
			break;
		}

		// Open the patch file
		fopen_s(&patchFile, patchFilename, "rb");
		if (!patchFile)
		{
			printf("Error: Could not open patch file %s\n", patchFilename);
			break;
		}

		// Open the output file
		fopen_s(&outputFile, outputFilename, "wb");
		if (!outputFile)
		{
			printf("Error: Could not open output file %s\n", outputFilename);
			break;
		}

		//read the patch file header and check for appropriate magic
		uint8_t header[24];
		if (fread(header, 1, 24, patchFile) != 24 || memcmp(header, "ENDSLEY/BSDIFF43", 16) != 0)
		{
			printf("Error: Corrupt patch file %s\n", patchFilename);
			break;
		}

		// Read output file lengths from header
		int64_t outputFileSize = *(int64_t*)(header + 16);
		pOutputBuffer = new uint8_t[outputFileSize + 1];
		if (pOutputBuffer == nullptr)
		{
			printf("Error: Failed to allocate memory[%lld] for output buffer\n", outputFileSize + 1);
			break;
		}

		BZFILE* bzHandle = nullptr;
		ZSTD_DecompressContext* zstdCtx = nullptr;
		if (compressMethod == 1)
		{
			int bz2err;
			bzHandle = BZ2_bzReadOpen(&bz2err, patchFile, 0, 0, NULL, 0);
			if (bzHandle == nullptr)
			{
				printf("BZ2_bzReadOpen, bz2err=%d", bz2err);
				break;
			}
		}
		else if(compressMethod==2)
		{
			zstdCtx = createZSTDDecompressContext(patchFile);
			if (zstdCtx == 0)
			{
				printf("create zstd context failed!");
				break;
			}
		}

		// Apply the patch
		struct bspatch_stream stream;
		if (compressMethod == 1)
		{
			stream.read = bzPatchFileReader;
			stream.opaque = bzHandle;
		}
		else if (compressMethod == 2)
		{
			stream.read = zstdPatchFileReader;
			stream.opaque = zstdCtx;
		}
		else
		{
			stream.read = patchFileReader;
			stream.opaque = patchFile;
		}

		if (bspatch(pInputBuffer, inputFileSize, pOutputBuffer, outputFileSize, &stream) !=0 )
		{
			printf("Error: Failed to apply patch\n");
			if (compressMethod == 1)
			{
				// Clean up the bzip2 reads
				int bz2err;
				BZ2_bzReadClose(&bz2err, bzHandle);
				bzHandle = nullptr;
			}
			else if (compressMethod == 2)
			{
				destoryZSTDDecompressContext(zstdCtx);
				zstdCtx = nullptr;
			}
			break;
		}

		if (compressMethod == 1)
		{
			// Clean up the bzip2 reads
			int bz2err;
			BZ2_bzReadClose(&bz2err, bzHandle);
			bzHandle = nullptr;
		}
		else if (compressMethod == 2)
		{
			destoryZSTDDecompressContext(zstdCtx);
			zstdCtx = nullptr;
		}

		// Write the output buffer to the output file
		if (outputFileSize != fwrite(pOutputBuffer, 1, outputFileSize, outputFile))
		{
			printf("Error: Failed to write output file %s\n", outputFilename);
			break;
		}

		result = true;
		printf("Patch applied successfully. Output file: %s\n", outputFilename);
	}while (false);

	// Close all files and free buffers
	freeFileBuffer(pInputBuffer);
	pInputBuffer = nullptr;

	if (pOutputBuffer != nullptr)
	{
		delete[] pOutputBuffer;
		pOutputBuffer = nullptr;
	}

	if (patchFile != nullptr)
	{
		fclose(patchFile);
		patchFile = nullptr;
	}
	if (outputFile != nullptr)
	{
		fclose(outputFile);
		outputFile = nullptr;
	}
	
	return result;
}