#include "volume.h"
#include "crypt.h"
#include "compression.h"
#include "hash.h"
#include "path.h"
#include "file.h"
#include "memory_file.h"
#include "util.h"

const uint32_t kGT5_VOLUME_KEY[4] = { 0x2DEE26A7, 0x412D99F5, 0x883C94E9, 0x0F1A7069 };
const uint32_t kGT6_VOLUME_KEY[4] = { 0xAA1B6A59, 0xE70B6FB3, 0x62DC6095, 0x6A594A25 };

const char kGT5_VOLUME_SEED[] = { "KALAHARI-37863889" };
const char kGT6_VOLUME_SEED[] = { "PISCINAS-323419048" };

namespace
{
	const uint32_t kVOLUME_HEADER_MAGIC = 0x5B745162u;
	const uint32_t kVOLUME_HEADER_SIZE = 0xA0;

	const uint32_t kVOLUME_SEGMENT_MAGIC = 0x5B74516Eu;
	const uint32_t kVOLUME_SEGMENT_HEADER_SIZE = 0x08;
	const uint32_t kVOLUME_SEGMENT_SIZE = 0x800;

	static KeyContext s_key_context;

	struct VolumeHeader
	{
		uint32_t magic;
		uint32_t segment_index;
		uint32_t size;
		uint32_t real_size;
		uint64_t unk4;
		uint64_t file_size;
		char title_id[128];

		bool load(IStream* stream)
		{
			if (!stream->read_be(magic) || magic != kVOLUME_HEADER_MAGIC)
				return false;
			if (!stream->read_be(segment_index))
				return false;
			if (!stream->read_be(size))
				return false;
			if (!stream->read_be(real_size))
				return false;
			if (!stream->read_be(unk4))
				return false;
			if (!stream->read_be(file_size))
				return false;
			if (!stream->read(title_id, sizeof(title_id)))
				return false;
			return true;
		}

		bool save(IStream* stream)
		{
			if (!stream->write_be(magic))
				return false;
			if (!stream->write_be(segment_index))
				return false;
			if (!stream->write_be(size))
				return false;
			if (!stream->write_be(real_size))
				return false;
			if (!stream->write_be(unk4))
				return false;
			if (!stream->write_be(file_size))
				return false;
			if (!stream->write(title_id, sizeof(title_id)))
				return false;
			return true;
		}
	};

	struct VolumeSegment
	{
		uint32_t segment_index;
		uint32_t size;
		uint32_t real_size;

		uint32_t unk1;
		uint32_t data_size;
		char* data;

		VolumeSegment(const uint32_t segment_index, const uint32_t size, const uint32_t real_size)
			: segment_index(segment_index)
			, size(size)
			, real_size(real_size)
			, data(nullptr)
		{
		}

		~VolumeSegment()
		{
			delete[] data;
		}

		bool load(IStream* stream)
		{
			if (!stream->read_le(unk1))
				return false;
			if (!stream->read_le(data_size))
				return false;

			size -= kVOLUME_SEGMENT_HEADER_SIZE;
			char* buffer = new char[size];
			if (!stream->read(buffer, size))
				return false;

			data_size = -static_cast<int32_t>(data_size);
			assert(data_size == real_size);
			data = new char[data_size];
			memset(data, 0, data_size);

			uint32_t real_data_size = data_size;
			if (inflate(buffer, size, data, &real_data_size, -15) != 0) {
				delete[] buffer;
				return false;
			}
			assert(data_size == real_data_size);

			delete[] buffer;

			return true;
		}
	};
}

Volume::Volume()
	: game_id_(GameID::UNKNOWN)
{
}

void crypt_header(void* const data, const uint32_t size)
{
	uint32_t key[4];
	data_keygen(s_key_context.seed, s_key_context.key, 1, key);
	data_crypt(key, data, data, size);
	block_crypt(data, data, size);
}

void crypt_segment(void* const data, const uint32_t size, const uint32_t index)
{
	uint32_t key[4];
	data_keygen(s_key_context.seed, s_key_context.key, index, key);
	data_crypt(key, data, data, size);
}

Volume::Volume(const char* const path, const GameID game_id)
	: game_id_(GameID::UNKNOWN)
{
	open(path, game_id);
}

Volume::~Volume()
{
	close();
}

bool Volume::open(const char* const path, const GameID game_id)
{
	game_id_ = game_id;

	const char* seed = nullptr;
	const uint32_t* key = nullptr;
	switch (game_id) {
		case GameID::GT5:
			seed = kGT5_VOLUME_SEED, key = kGT5_VOLUME_KEY;
			break;
		case GameID::GT6:
			seed = kGT6_VOLUME_SEED, key = kGT6_VOLUME_KEY;
			break;
	}
	if (!seed || !key)
		return false;

	memset(&s_key_context, 0, sizeof(s_key_context));
	strncpy(s_key_context.seed, seed, sizeof(s_key_context.seed));
	memcpy(s_key_context.key, key, sizeof(s_key_context.key));

	File volume_file;
	if (!volume_file.open(path, OpenMode::kREAD | OpenMode::kSHARE))
		return false;

	VolumeHeader header;
	{
		char* data = static_cast<char*>(_alloca(kVOLUME_HEADER_SIZE));
		memset(data, 0, kVOLUME_HEADER_SIZE);
		if (!volume_file.read(data, kVOLUME_HEADER_SIZE))
			return false;
		crypt_header(data, kVOLUME_HEADER_SIZE);
		MemoryFile memory_file(data, kVOLUME_HEADER_SIZE);
		if (!header.load(&memory_file))
			return false;
	}

	volume_file.seek(kVOLUME_SEGMENT_SIZE);

	std::unique_ptr<char[]> data(new char[header.size]);
	if (!volume_file.read(data.get(), header.size))
		return false;
	crypt_segment(data.get(), header.size, header.segment_index);

	VolumeSegment segment(header.segment_index, header.size, header.real_size);
	{
		MemoryFile memory_file(data.get(), header.size);
		if (!segment.load(&memory_file))
			return false;
	}

	data.reset();

	packed_file_.reset(new PackedFile(path, segment.data, segment.data_size, segment.size));

	return true;
}

void Volume::close()
{
	packed_file_.reset();

	game_id_ = GameID::UNKNOWN;
}

bool Volume::extract(const char* const output_dir) {
	if (!packed_file_)
		return false;

	packed_file_->extract_entries(output_dir);

	return true;
}