#include <algorithm>
#include <cstdio>
#include <cstring>
#include <vector>

#include "Common/Data/Format/PNGLoad.h"
#include "Common/File/FileUtil.h"
#include "Common/File/Path.h"
#include "Common/Thread/ThreadManager.h"
#include "GPU/Common/TextureReplacer.h"

#include "UnitTest.h"

// The fake hashes in textures.ini are designed to be readable: the address
// part starts with A, the CLUT hash with C, and the contents hash with 1.
static const u64 KEY_A = 0xA0000001C0000001ULL;  // addr 0xA0000001, clut 0xC0000001
static const u32 HASH_A = 0x10000001;
static const u64 KEY_B = 0xA0000002C0000002ULL;  // two mip levels
static const u32 HASH_B = 0x10000002;
static const u64 KEY_C = 0xA0000003C0000003ULL;  // hashrange
static const u32 HASH_C = 0x10000003;
static const u64 KEY_D = 0xA0000004C0000004ULL;  // ignored (empty filename)
static const u32 HASH_D = 0x10000004;
static const u64 KEY_E = 0xA0000005C0000005ULL;  // file missing
static const u32 HASH_E = 0x10000005;
static const u64 KEY_F = 0xA0000006C0000006ULL;  // same file as A, own hashrange
static const u32 HASH_F = 0x10000006;
static const u64 KEY_G = 0xA0000007C0000007ULL;  // ignoreAddress pack, hashrange
static const u32 HASH_G = 0x10000007;

static bool CreateTestPNG(const Path &filename, int w, int h, u32 color) {
	std::vector<u8> buf((size_t)w * h * 4);
	for (int i = 0; i < w * h; i++) {
		buf[i * 4] = (color >> 24) & 0xFF;
		buf[i * 4 + 1] = (color >> 16) & 0xFF;
		buf[i * 4 + 2] = (color >> 8) & 0xFF;
		buf[i * 4 + 3] = color & 0xFF;
	}
	return pngSave(filename, buf.data(), w, h, 4);
}

static bool WriteIni(const Path &packDir, const char *iniContent) {
	FILE *f = File::OpenCFile(packDir / "textures.ini", "w");
	if (!f) {
		return false;
	}
	fwrite(iniContent, 1, strlen(iniContent), f);
	fclose(f);
	return true;
}

static bool CreateTestPack(const Path &packDir) {
	File::DeleteDirRecursively(packDir);
	if (!File::CreateDir(packDir)) {
		return false;
	}

	static const char *iniContent =
		"[options]\n"
		"hash = quick\n"
		"version = 1\n"
		"\n"
		"[hashes]\n"
		"A0000001C000000110000001 = tex_a.png\n"
		"A0000002C000000210000002 = tex_b0.png\n"
		"A0000002C000000210000002_1 = tex_b1.png\n"
		"A0000003C000000310000003 = tex_c.png\n"
		"A0000004C000000410000004 =\n"
		"A0000005C000000510000005 = missing.png\n"
		"A0000006C000000610000006 = tex_a.png\n"
		"\n"
		"[hashranges]\n"
		"A0000003,512,512 = 256,256\n"
		"A0000006,64,64 = 32,32\n"
		"\n"
		"[filtering]\n"
		"A0000001C000000110000001 = nearest\n";

	if (!WriteIni(packDir, iniContent)) {
		return false;
	}

	if (!CreateTestPNG(packDir / "tex_a.png", 64, 64, 0xFF0000FF)) return false;
	if (!CreateTestPNG(packDir / "tex_b0.png", 64, 64, 0x00FF00FF)) return false;
	if (!CreateTestPNG(packDir / "tex_b1.png", 32, 32, 0x00FF00FF)) return false;
	if (!CreateTestPNG(packDir / "tex_c.png", 256, 256, 0x0000FFFF)) return false;

	return true;
}

static bool TestLookups(TextureReplacer *replacer) {
	// Key A: single mip, found.
	ReplacedTexture *texA = replacer->FindReplacement(ReplacementCacheKey(KEY_A, HASH_A), 64, 64);
	EXPECT_TRUE(texA != nullptr);

	// Key B: two mip levels, found.
	ReplacedTexture *texB = replacer->FindReplacement(ReplacementCacheKey(KEY_B, HASH_B), 64, 64);
	EXPECT_TRUE(texB != nullptr);

	// Key C: found via hashrange.
	ReplacedTexture *texC = replacer->FindReplacement(ReplacementCacheKey(KEY_C, HASH_C), 512, 512);
	EXPECT_TRUE(texC != nullptr);

	// Key D: explicitly ignored (empty filename).
	EXPECT_TRUE(replacer->FindReplacement(ReplacementCacheKey(KEY_D, HASH_D), 16, 16) == nullptr);

	// Key E: file missing, still creates a texture object (loads as NOT_FOUND).
	ReplacedTexture *texE = replacer->FindReplacement(ReplacementCacheKey(KEY_E, HASH_E), 16, 16);
	EXPECT_TRUE(texE != nullptr);

	// Key F: the same file as A, but its own hashrange and no filtering, so it can't share A's texture.
	ReplacedTexture *texF = replacer->FindReplacement(ReplacementCacheKey(KEY_F, HASH_F), 64, 64);
	EXPECT_TRUE(texF != nullptr);
	EXPECT_TRUE(texF != texA);

	// Unknown key: not in the ini at all.
	ReplacementCacheKey unknownKey(0x2000000020000000ULL, 0x20000000);
	EXPECT_TRUE(replacer->FindReplacement(unknownKey, 16, 16) == nullptr);

	if (!texA || !texB || !texC || !texE || !texF) {
		return false;
	}

	// Load the actual textures using the thread manager.
	g_threadManager.Init(1, 1);

	// Key A: 64x64 single level, nearest filtering.
	EXPECT_TRUE(texA->Poll(1.0));
	EXPECT_EQ_INT(texA->NumLevels(), 1);
	int w = 0, h = 0;
	texA->GetSize(0, &w, &h);
	EXPECT_EQ_INT(w, 64);
	EXPECT_EQ_INT(h, 64);
	TextureFiltering filtering = TEX_FILTER_AUTO;
	EXPECT_TRUE(texA->ForceFiltering(&filtering));
	EXPECT_TRUE(filtering == TEX_FILTER_FORCE_NEAREST);

	// Key B: two mip levels, 64x64 and 32x32.
	EXPECT_TRUE(texB->Poll(1.0));
	EXPECT_EQ_INT(texB->NumLevels(), 2);
	texB->GetSize(0, &w, &h);
	EXPECT_EQ_INT(w, 64);
	EXPECT_EQ_INT(h, 64);
	texB->GetSize(1, &w, &h);
	EXPECT_EQ_INT(w, 32);
	EXPECT_EQ_INT(h, 32);

	// Key C: hashrange maps 512x512 -> 256x256, upscaled back to 512x512.
	EXPECT_TRUE(texC->Poll(1.0));
	EXPECT_EQ_INT(texC->NumLevels(), 1);
	texC->GetSize(0, &w, &h);
	EXPECT_EQ_INT(w, 512);
	EXPECT_EQ_INT(h, 512);

	// Key F: hashrange maps 64x64 -> 32x32, so the 64x64 image covers 128x128.
	EXPECT_TRUE(texF->Poll(1.0));
	texF->GetSize(0, &w, &h);
	EXPECT_EQ_INT(w, 128);
	EXPECT_EQ_INT(h, 128);
	EXPECT_FALSE(texF->ForceFiltering(&filtering));

	// Key E: missing file should end up NOT_FOUND.
	EXPECT_TRUE(texE->Poll(1.0));
	EXPECT_TRUE(texE->State() == ReplacementState::NOT_FOUND);

	g_threadManager.Teardown();
	return true;
}

// ignoreAddress matches files without the address, but hash ranges still apply per address.
static bool TestIgnoreAddress(const Path &packDir) {
	static const char *iniContent =
		"[options]\n"
		"hash = xxh32\n"
		"ignoreAddress = true\n"
		"version = 1\n"
		"\n"
		"[hashes]\n"
		"00000000C000000710000007 = tex_c.png\n"
		"\n"
		"[hashranges]\n"
		"A0000007,512,512 = 256,256\n";
	if (!WriteIni(packDir, iniContent)) {
		return false;
	}

	TextureReplacer replacer(nullptr);
	std::string error;
	if (!replacer.LoadPackForTesting(packDir, &error)) {
		return false;
	}

	ReplacementCacheKey key(KEY_G, HASH_G);
	ReplacedTexture *texG = replacer.FindReplacement(key, 512, 512);
	EXPECT_TRUE(texG != nullptr);
	if (!texG) {
		return false;
	}
	EXPECT_TRUE(replacer.FindReplacement(key, 512, 512) == texG);

	g_threadManager.Init(1, 1);
	EXPECT_TRUE(texG->Poll(1.0));
	int w = 0, h = 0;
	texG->GetSize(0, &w, &h);
	EXPECT_EQ_INT(w, 512);
	EXPECT_EQ_INT(h, 512);
	g_threadManager.Teardown();
	return true;
}

static void AppendBE32(std::vector<u8> *out, u32 value) {
	for (int shift = 24; shift >= 0; shift -= 8) {
		out->push_back((value >> shift) & 0xFF);
	}
}

static u32 PNGCrc32(const u8 *data, size_t size) {
	u32 crc = 0xFFFFFFFF;
	for (size_t i = 0; i < size; i++) {
		crc ^= data[i];
		for (int bit = 0; bit < 8; bit++) {
			crc = (crc >> 1) ^ (0xEDB88320 & (0 - (crc & 1)));
		}
	}
	return ~crc;
}

static void AppendPNGChunk(std::vector<u8> *out, const char *type, const std::vector<u8> &data) {
	AppendBE32(out, (u32)data.size());
	std::vector<u8> typeAndData(type, type + 4);
	typeAndData.insert(typeAndData.end(), data.begin(), data.end());
	out->insert(out->end(), typeAndData.begin(), typeAndData.end());
	AppendBE32(out, PNGCrc32(typeAndData.data(), typeAndData.size()));
}

// Writes an 8-bit palette PNG (a palette in reverse gray order, so the index is not the color)
// or an 8-bit grayscale PNG, one byte per pixel. The image data is stored uncompressed, which
// avoids needing libpng or zlib here.
static bool CreateIndexedPNG(const Path &filename, int w, int h, bool palette, const std::vector<u8> &indices) {
	std::vector<u8> raw;
	for (int y = 0; y < h; y++) {
		raw.push_back(0);  // No filter.
		raw.insert(raw.end(), indices.begin() + y * w, indices.begin() + (y + 1) * w);
	}

	// A zlib stream of stored blocks.
	std::vector<u8> zlib = { 0x78, 0x01 };
	for (size_t pos = 0; pos < raw.size() || pos == 0;) {
		const size_t len = std::min<size_t>(raw.size() - pos, 65535);
		const bool last = pos + len >= raw.size();
		zlib.push_back(last ? 1 : 0);
		zlib.push_back(len & 0xFF);
		zlib.push_back(len >> 8);
		zlib.push_back(~len & 0xFF);
		zlib.push_back((~len >> 8) & 0xFF);
		zlib.insert(zlib.end(), raw.begin() + pos, raw.begin() + pos + len);
		pos += len;
		if (last) {
			break;
		}
	}
	u32 a = 1, b = 0;
	for (u8 c : raw) {
		a = (a + c) % 65521;
		b = (b + a) % 65521;
	}
	AppendBE32(&zlib, (b << 16) | a);

	std::vector<u8> file = { 0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A };
	std::vector<u8> ihdr;
	AppendBE32(&ihdr, w);
	AppendBE32(&ihdr, h);
	ihdr.push_back(8);  // Bit depth.
	ihdr.push_back(palette ? 3 : 0);  // Color type: palette or grayscale.
	ihdr.push_back(0);
	ihdr.push_back(0);
	ihdr.push_back(0);
	AppendPNGChunk(&file, "IHDR", ihdr);
	if (palette) {
		std::vector<u8> plte;
		for (int i = 0; i < 256; i++) {
			plte.push_back(255 - i);
			plte.push_back(255 - i);
			plte.push_back(255 - i);
		}
		AppendPNGChunk(&file, "PLTE", plte);
	}
	AppendPNGChunk(&file, "IDAT", zlib);
	AppendPNGChunk(&file, "IEND", {});

	FILE *fp = File::OpenCFile(filename, "wb");
	if (!fp) {
		return false;
	}
	const bool ok = fwrite(file.data(), 1, file.size(), fp) == file.size();
	fclose(fp);
	return ok;
}

static bool ReadFileBytes(const Path &filename, std::vector<u8> *data) {
	std::string str;
	if (!File::ReadBinaryFileToString(filename, &str)) {
		return false;
	}
	data->assign(str.begin(), str.end());
	return true;
}

static bool TestIndexedPNG(const Path &packDir) {
	std::vector<u8> indices(16 * 8);
	for (size_t i = 0; i < indices.size(); i++) {
		indices[i] = (u8)(i * 7 + 3);
	}

	for (int palette = 0; palette < 2; palette++) {
		Path file = packDir / "indexed_test.png";
		if (!CreateIndexedPNG(file, 16, 8, palette != 0, indices)) {
			return false;
		}
		std::vector<u8> png;
		if (!ReadFileBytes(file, &png)) {
			return false;
		}
		int w = 0, h = 0;
		unsigned char *out = nullptr;
		EXPECT_TRUE(pngLoadIndexed(png.data(), png.size(), &w, &h, &out) == 1);
		EXPECT_EQ_INT(w, 16);
		EXPECT_EQ_INT(h, 8);
		if (!out) {
			return false;
		}
		EXPECT_TRUE(memcmp(out, indices.data(), indices.size()) == 0);
		free(out);

		// Too small a limit is refused.
		EXPECT_TRUE(pngLoadIndexed(png.data(), png.size(), &w, &h, &out, 8, 8) == 0);
	}

	// Truecolor is not an index.
	Path rgba = packDir / "indexed_rgba.png";
	if (!CreateTestPNG(rgba, 4, 4, 0x11223344)) {
		return false;
	}
	std::vector<u8> png;
	if (!ReadFileBytes(rgba, &png)) {
		return false;
	}
	int w = 0, h = 0;
	unsigned char *out = nullptr;
	EXPECT_TRUE(pngLoadIndexed(png.data(), png.size(), &w, &h, &out) == 0);
	EXPECT_TRUE(out == nullptr);
	return true;
}

// The [indexed] section marks an image as palette indices. It is looked up without the CLUT hash
// so that every palette of the texture shares the one image.
static bool TestIndexedReplacement(const Path &packDir) {
	static const char *iniContent =
		"[options]\n"
		"hash = quick\n"
		"version = 1\n"
		"\n"
		"[hashes]\n"
		"A00000080000000010000008 = indexed.png\n"
		"A00000090000000010000009 = indexed_rgba.png\n"
		"A000000A000000001000000A = indexed.png\n"
		"\n"
		"[indexed]\n"
		"A00000080000000010000008 = true\n"
		"A00000090000000010000009 = true\n";
	if (!WriteIni(packDir, iniContent)) {
		return false;
	}

	std::vector<u8> indices(16 * 8);
	for (size_t i = 0; i < indices.size(); i++) {
		indices[i] = (u8)(255 - i * 2);
	}
	if (!CreateIndexedPNG(packDir / "indexed.png", 16, 8, true, indices)) {
		return false;
	}
	if (!CreateTestPNG(packDir / "indexed_rgba.png", 16, 8, 0x11223344)) {
		return false;
	}

	TextureReplacer replacer(nullptr);
	std::string error;
	if (!replacer.LoadPackForTesting(packDir, &error)) {
		return false;
	}

	// Two palettes of the same texture.
	ReplacedTexture *pal1 = replacer.FindReplacement(ReplacementCacheKey(0xA0000008C0000001ULL, 0x10000008), 16, 8);
	ReplacedTexture *pal2 = replacer.FindReplacement(ReplacementCacheKey(0xA0000008C0000002ULL, 0x10000008), 16, 8);
	EXPECT_TRUE(pal1 != nullptr);
	EXPECT_TRUE(pal1 == pal2);

	// The same image without the [indexed] entry is a regular replacement, which can't share the texture.
	ReplacedTexture *plain = replacer.FindReplacement(ReplacementCacheKey(0xA000000AC0000001ULL, 0x1000000A), 16, 8);
	EXPECT_TRUE(plain != nullptr);
	EXPECT_TRUE(plain != pal1);

	ReplacedTexture *bad = replacer.FindReplacement(ReplacementCacheKey(0xA0000009C0000001ULL, 0x10000009), 16, 8);
	EXPECT_TRUE(bad != nullptr);
	if (!pal1 || !plain || !bad) {
		return false;
	}
	EXPECT_TRUE(pal1->IsIndexed());
	EXPECT_FALSE(plain->IsIndexed());

	g_threadManager.Init(1, 1);
	EXPECT_TRUE(pal1->Poll(1.0));
	EXPECT_TRUE(pal1->State() == ReplacementState::ACTIVE);
	EXPECT_EQ_INT(pal1->NumLevels(), 1);
	int w = 0, h = 0;
	pal1->GetIndexDataSize(0, &w, &h);
	EXPECT_EQ_INT(w, 16);
	EXPECT_EQ_INT(h, 8);
	pal1->GetSize(0, &w, &h);
	EXPECT_EQ_INT(w, 16);
	EXPECT_EQ_INT(h, 8);
	EXPECT_TRUE(pal1->Format() == Draw::DataFormat::R8G8B8A8_UNORM);

	std::vector<u8> copy(indices.size());
	EXPECT_TRUE(pal1->CopyIndexLevelTo(0, copy.data(), copy.size()));
	EXPECT_TRUE(copy == indices);

	// A truecolor image can't serve as indices.
	EXPECT_TRUE(bad->Poll(1.0));
	EXPECT_TRUE(bad->State() == ReplacementState::NOT_FOUND);
	g_threadManager.Teardown();
	return true;
}

bool TestTextureReplacer() {
	Path packDir = Path("unittest_texture_pack");
	if (!CreateTestPack(packDir)) {
		return false;
	}

	TextureReplacer replacer(nullptr);
	std::string error;
	if (!replacer.LoadPackForTesting(packDir, &error)) {
		ERROR_LOG(Log::G3D, "Failed to load test texture pack: %s", error.c_str());
		File::DeleteDirRecursively(packDir);
		return false;
	}

	if (!TestLookups(&replacer)) {
		File::DeleteDirRecursively(packDir);
		return false;
	}

	// Reloading the ini forgets what the old one said about each texture.
	EXPECT_TRUE(replacer.GetNumTrackedTextures() != 0);
	EXPECT_TRUE(replacer.LoadPackForTesting(packDir, &error));
	EXPECT_EQ_INT(replacer.GetNumTrackedTextures(), 0);

	if (!TestIgnoreAddress(packDir)) {
		File::DeleteDirRecursively(packDir);
		return false;
	}

	if (!TestIndexedPNG(packDir) || !TestIndexedReplacement(packDir)) {
		File::DeleteDirRecursively(packDir);
		return false;
	}

	File::DeleteDirRecursively(packDir);
	return true;
}