#include "stdafx.h"
#include <assert.h>
#include <cstring>
#include "UpsPatcher.h"
#include "CRC32.h"

uint64_t UpsPatcher::ReadBase128Number(std::istream &file)
{
	uint64_t result = 0;
	int shift = 0;
	uint8_t buffer;
	while(true) {
		file.read((char*)&buffer, 1);
		if(file.eof()) {
			return -1;
		}
		result += (buffer & 0x7F) << shift;
		shift += 7;
		if(buffer & 0x80) {
			break;
		}
		result += (uint64_t)1 << shift;
	}

	return result;
}

bool UpsPatcher::PatchBuffer(string upsFilepath, vector<uint8_t> &input, vector<uint8_t> &output)
{
	ifstream upsFile(upsFilepath, std::ios::in | std::ios::binary);
	if(upsFile) {
		return PatchBuffer(upsFile, input, output);
	}
	return false;
}

bool UpsPatcher::PatchBuffer(std::istream &upsFile, vector<uint8_t> &input, vector<uint8_t> &output)
{
	upsFile.seekg(0, std::ios::end);
	size_t fileSize = (size_t)upsFile.tellg();
	upsFile.seekg(0, std::ios::beg);

	char header[4];
	upsFile.read((char*)&header, 4);
	if(memcmp((char*)&header, "UPS1", 4) != 0) {
		//Invalid UPS file
		return false;
	}

	uint64_t inputFileSize = ReadBase128Number(upsFile);
	uint64_t outputFileSize = ReadBase128Number(upsFile);
	if(inputFileSize == -1 || outputFileSize == -1) {
		//Invalid file
		return false;
	}

	output.resize((size_t)outputFileSize);
	std::copy(input.begin(), input.end(), output.begin());

	uint32_t pos = 0;
	while((size_t)upsFile.tellg() < fileSize - 12) {
		uint32_t offset = (uint32_t)ReadBase128Number(upsFile);
		if(offset == -1) {
			//Invalid file
			return false;
		}

		pos += offset;

		while(true) {
			uint8_t xorValue = 0;
			upsFile.read((char*)&xorValue, 1);
			if((size_t)upsFile.tellg() > fileSize - 12) {
				//Invalid file
				return false;
			}

			output[pos] ^= xorValue;
			pos++;

			if(!xorValue) {
				break;
			}
		}
	}

	uint8_t inputChecksum[4];
	uint8_t outputChecksum[4];
	upsFile.read((char*)inputChecksum, 4);
	upsFile.read((char*)outputChecksum, 4);
	uint32_t patchInputCrc = inputChecksum[0] | (inputChecksum[1] << 8) | (inputChecksum[2] << 16) | (inputChecksum[3] << 24);
	uint32_t patchOutputCrc = outputChecksum[0] | (outputChecksum[1] << 8) | (outputChecksum[2] << 16) | (outputChecksum[3] << 24);
	uint32_t inputCrc = CRC32::GetCRC(input.data(), input.size());
	uint32_t outputCrc = CRC32::GetCRC(output.data(), output.size());

	if(patchInputCrc != inputCrc || patchOutputCrc != outputCrc) {
		return false;
	}
	return true;
}