#include "common.h"
#include "bsp.h"
#include "ps1types.h"
#include "ps1bsp.h"
#include "texture.h"

#include "rectpack/finders_interface.h"
#include "tim.h"

#define PALETTE_SIZE	256

static char path[_MAX_PATH];

/**
tp specifies the color depth for the texture page in the range of 0 to 2 (0:4-bit, 1:8-bit, 2:16-bit).
abr specifies the blend operator for both non-textured and textured semi-transparent primitives which can be ignored for now and lastly,
x,y specifies the X,Y coordinates of the VRAM in 16-bit pixel units.
Keep in mind that the coordinates will be rounded down to the next lowest texture page. 
*/
#define getTPage(tp, abr, x, y) ( \
	(((x)  /  64) & 15) | \
	((((y) / 256) &  1) <<  4) | \
	(((abr)       &  3) <<  5) | \
	(((tp)        &  3) <<  7) | \
	((((y) / 512) &  1) << 11) \
)

#define getTexWindow(x, y, w, h) ( \
	0xe2000000 | \
	((~((w)-1) & 0xFF)) >> 3 | \
	(((~((h)-1) & 0xFF) >> 3) <<  5) | \
	((((x) & 0xFF) >> 3) << 10) | \
	((((y) & 0xFF) >> 3) << 15) \
)

// A straight conversion of the Quake palette colors comes out looking rather over-saturated.
// So we desaturate the palette ahead of time to more closely match the original Quake's look.
// Yoinked from: https://stackoverflow.com/questions/13328029/how-to-desaturate-a-color
static void desaturate(const unsigned char inColor[3], unsigned char outColor[3])
{
	double f = 0.1; // desaturate by 10%
	double L = 0.3 * inColor[0] + 0.6 * inColor[1] + 0.1 * inColor[2];
	outColor[0] = (unsigned char)((double)inColor[0] + f * (L - inColor[0]));
	outColor[1] = (unsigned char)((double)inColor[1] + f * (L - inColor[1]));
	outColor[2] = (unsigned char)((double)inColor[2] + f * (L - inColor[2]));
}

static bool load_palette(const char* paletteFile, Color outPalette[PALETTE_SIZE])
{
	unsigned char palette[PALETTE_SIZE * 3];

	FILE* fp;
	fopen_s(&fp, paletteFile, "rb");
	if (fp == NULL)
		return false;

	fread(palette, sizeof(unsigned char) * 3, PALETTE_SIZE, fp);
	fclose(fp);

	for (int c = 0; c < PALETTE_SIZE; ++c)
	{
		outPalette[c].rgb.r = palette[3 * c + 0];
		outPalette[c].rgb.g = palette[3 * c + 1];
		outPalette[c].rgb.b = palette[3 * c + 2];
		outPalette[c].rgb.a = 0;
	}

	return true;
}

static bool generate_clut(const Color palette[PALETTE_SIZE], tim::PARAM* outTim)
{
	tim::PIX_RGB5* clut = (tim::PIX_RGB5*)malloc(PALETTE_SIZE * sizeof(tim::PIX_RGB5) * 2);
	if (clut == NULL)
		return false;

	memset(clut, 0, PALETTE_SIZE * sizeof(tim::PIX_RGB5) * 2);

	for (int c = 0; c < PALETTE_SIZE - 1; ++c)	// Final palette entry is for alpha masking
	{
		unsigned char color[3];
		desaturate(palette[c].channel, color);

		clut[c].r = color[0] >> 3;
		clut[c].g = color[1] >> 3;
		clut[c].b = color[2] >> 3;
		clut[c].i = 0;

		// Completely black pixels are regarded as transparent by the PS1, so prevent that from happening by making those palette entries *nearly* black
		if (clut[c].r == 0 && clut[c].g == 0 && clut[c].b == 0)
			clut[c].r = clut[c].g = clut[c].b = 1;
	}

	// Create a second palette with the transparency flag set. Used for water surfaces.
	for (int c = 0; c < PALETTE_SIZE; ++c)
	{
		clut[c + PALETTE_SIZE] = clut[c];
		clut[c + PALETTE_SIZE].i = 1;
	}

	outTim->clutData = clut;
	outTim->clutWidth = PALETTE_SIZE;
	outTim->clutHeight = 2;
	return true;
}

static int color_hash(const Color& color)
{
	const int divider = 24;
	int r = (int)roundf((float)color.rgb.r / divider);
	int g = (int)roundf((float)color.rgb.g / divider);
	int b = (int)roundf((float)color.rgb.b / divider);
	return (b << 16) | (g << 8) | r;
}

struct ColorAccumulate
{
	uint32_t r, g, b;
	size_t count;

	ColorAccumulate() : r(0), g(0), b(0), count(0) { }

	ColorAccumulate(const Color& color) : r(color.rgb.r), g(color.rgb.g), b(color.rgb.b), count(1)
	{
	}

	Color getAverage()
	{
		Color result;
		result.rgb.r = (unsigned char)(r / count);
		result.rgb.g = (unsigned char)(g / count);
		result.rgb.b = (unsigned char)(b / count);
		return result;
	}
};

// Dominant color algorithm yoinked from: https://github.com/fast-average-color/fast-average-color
static void analyze_texture(unsigned char* texBytes, int numBytes, const Color palette[PALETTE_SIZE], TextureDescriptor& outTexture)
{
	uint64_t rSum = 0, gSum = 0, bSum = 0;
	std::unordered_map<int, ColorAccumulate> colorCount;

	ColorAccumulate dominant;
	for (int i = 0; i < numBytes; ++i)
	{
		unsigned char b = texBytes[i];
		Color color = palette[b];

		rSum += color.rgb.r;
		gSum += color.rgb.g;
		bSum += color.rgb.b;

		int key = color_hash(color);
		const auto iter = colorCount.find(key);
		if (iter != colorCount.end())
		{
			iter->second.r += color.rgb.r;
			iter->second.g += color.rgb.g;
			iter->second.b += color.rgb.b;
			iter->second.count++;
		}
		else
		{
			colorCount[key] = ColorAccumulate(color);
		}

		if (colorCount[key].count > dominant.count)
			dominant = colorCount[key];
	}

	outTexture.averageColor.rgb.r = (unsigned char)(rSum / numBytes);
	outTexture.averageColor.rgb.g = (unsigned char)(gSum / numBytes);
	outTexture.averageColor.rgb.b = (unsigned char)(bSum / numBytes);

	outTexture.dominantColor = dominant.getAverage();

	desaturate(outTexture.averageColor.channel, outTexture.averageColor.channel);
	desaturate(outTexture.dominantColor.channel, outTexture.dominantColor.channel);
}

using spaces_type = rectpack2D::empty_spaces<false>;
using rect_type = rectpack2D::output_rect_t<spaces_type>;

static rectpack2D::rect_wh pack_textures(const std::map<int, miptex_t*>& textures, const rectpack2D::rect_wh &max_bin, std::vector<rect_type>& outRectangles)
{
	auto report_successful = [](rect_type&) {
		return rectpack2D::callback_result::CONTINUE_PACKING;
	};

	auto report_unsuccessful = [](rect_type&) {
		printf("Failed to fit all textures into atlas!\n");
		return rectpack2D::callback_result::ABORT_PACKING;
	};

	const auto discard_step = -4;

	for (auto texIter = textures.cbegin(); texIter != textures.cend(); ++texIter)
	{
		const miptex_t* miptex = texIter->second;

		// Shrink the larger textures, but keep smaller ones at their original size
		int ps1mip = miptex->width > 64 || miptex->height > 64 ? 1 : 0;

		// Make an exception for the difficulty selection teleporters
		if (!strncmp(miptex->name, "skill", 5))
			ps1mip = 0;

		// Ensure we don't include any textures that are larger than the PS1 can address
		// Texture pages are 64 pixels wide, effectively 128 texels in 8-bit mode, and with U offset added no U value may exceed 255
		// We could solve this more elegantly by page-aligning larger textures, but I'm not sure if that's actually necessary
		while ((miptex->width >> ps1mip) > 128 || (miptex->height >> ps1mip) > 256)
			ps1mip++;

		// TODO: make an exception for the QUAKE sign that's displayed on the start map. It needs to be bold and detailed, but it's too wide for the PS1 to address at full resolution, so it'll need to be broken up.
		// The brush that displays this texture is actually broken up into two parts: one that's 224 texels wide, and one that's 64 texels wide. We can take advantage of that.
		// We just need to turn that second part into a separate texture entry, and patch the texture ID and UVs on the face when we encounter it.
		if (!strcmp(miptex->name, "quake"))
			ps1mip--;

		if (strcmp(miptex->name, "clip") && strcmp(miptex->name, "trigger"))
			outRectangles.emplace_back(rectpack2D::rect_xywh(0, 0, miptex->width >> ps1mip, miptex->height >> ps1mip));
		else
			outRectangles.emplace_back(rectpack2D::rect_xywh(0, 0, miptex->width >> 3, miptex->width >> 3));	// Add the lowest mip level so that it at least gets included in the final texture list, and we don't mess up the texture IDs
	}

	return rectpack2D::find_best_packing<spaces_type>(
		outRectangles,
		rectpack2D::make_finder_input(
			max_bin,
			discard_step,
			report_successful,
			report_unsuccessful,
			rectpack2D::flipping_option::DISABLED
		)
	);
}

static void build_atlas(const std::vector<miptex_t*>& textures, const std::vector<rect_type>& rectangles, const rectpack2D::rect_wh& size, unsigned char* outAtlas, TextureList& outTextures)
{
}

bool process_textures(const world_t* world, TextureList& outTextures)
{
	// Construct palette CLUT and start building the output TIM structure
	Color palette[PALETTE_SIZE];
	load_palette("palette.lmp", palette);

	tim::PARAM outTim = { 0 };
	outTim.format = 1;	// 8-bit per pixel, all Quake textures use this
	outTim.imgXoffs = 512;
	outTim.imgYoffs = 256;
	outTim.clutXoffs = 512;
	outTim.clutYoffs = 0;
	generate_clut(palette, &outTim);

	auto max_bin = rectpack2D::rect_wh(1024, 256);	// 8-bit textures take up half the horizontal space so this is 512x256 in practice, or a quarter of the PS1's VRAM allocation.
	
	printf("Finding best texture packing...\n");

	std::map<int, miptex_t*> textures, repeatables;

	// Try some texture packing and see if we fit inside the PS1's VRAM
	for (int texNum = 0; texNum < world->mipheader.numtex; ++texNum)
	{
		miptex_t* miptex = &world->miptexes[texNum];
		if (miptex->name[0] == '\0')	// Weird edge case on N64START.bsp, corrupt data perhaps?
			miptex->width = miptex->height = 0;

		//printf("Texture %d (%dx%d): %.16s\n", texNum, miptex->width, miptex->height, miptex->name);

		if (texture_isRepeatable(miptex))
			repeatables[texNum] = miptex;
		else
			textures[texNum] = miptex;
	}

	// Pack the repeatable textures first, so that they get placed at predictable locations
	// rectpack2D is set up to pack larger textures first, so this actually works out with power-of-two textures
	std::vector<rect_type> rRects;
	const auto rSize = pack_textures(repeatables, max_bin, rRects);

	max_bin.w -= rSize.w;

	// Pack the remaining textures into the remaining space
	std::vector<rect_type> tRects;
	const auto tSize = pack_textures(textures, max_bin, tRects);

	auto result_size = rectpack2D::rect_wh(rSize.w + tSize.w, max(rSize.h, tSize.h));

	// Collect all of the texture rectangles back into a single in-order list again
	std::vector<rect_type> rectangles;
	auto rIter = rRects.begin();
	auto tIter = tRects.begin();
	for (int texNum = 0; texNum < world->mipheader.numtex; ++texNum)
	{
		if (repeatables.find(texNum) != repeatables.end())
		{
			rectangles.push_back(*rIter++);
			continue;
		}
		
		auto rect = *tIter++;
		rect.x += rSize.w;
		rectangles.push_back(rect);
	}

	printf("%d textures. Packed texture atlas size: %d x %d\n", world->mipheader.numtex, result_size.w, result_size.h);
	
	outTim.imgWidth = result_size.w;
	outTim.imgHeight = result_size.h;
	outTim.imgData = malloc(result_size.w * result_size.h * sizeof(unsigned char));
	if (outTim.imgData == NULL)
		return false;

	memset(outTim.imgData, 0, result_size.w * result_size.h * sizeof(unsigned char));

	printf("Constructing texture atlas...\n");

	// Try to construct the texture atlas, see what we get
	for (int texNum = 0; texNum < world->mipheader.numtex; ++texNum)
	{
		miptex_t* miptex = &world->miptexes[texNum];
		if (miptex->name[0] == '\0')	// Weird edge case on N64START.bsp, corrupt data perhaps?
		{
			outTextures.push_back(TextureDescriptor{ 0 });	// We have to add something, otherwise the texture IDs get messed up
			continue;
		}

		TextureDescriptor tex = { 0 };

		for (int mipLevel = 0; mipLevel < 4; ++mipLevel)
		{
			unsigned char* texBytes = world->textures[texNum * 4 + mipLevel];

			const auto& rectangle = rectangles[texNum];
			if (miptex->width >> mipLevel == rectangle.w)	// This is the mip level we've previously decided we want for our PS1 atlas
			{
				// Copy the source texture line by line into the atlas at the desired position
				for (int y = 0; y < rectangle.h; ++y)
				{
					memcpy_s((unsigned char*)outTim.imgData + ((rectangle.y + y) * result_size.w + rectangle.x), rectangle.w * sizeof(unsigned char), texBytes + (y * rectangle.w), rectangle.w * sizeof(unsigned char));
				}

				tex.w = (u_char)rectangle.w;
				tex.h = (u_char)rectangle.h;

				u_short x = (rectangle.x / 2) + outTim.imgXoffs;	// Divide by 2 to get the coordinate in 16-bit pixel units
				u_short y = rectangle.y + outTim.imgYoffs;

				tex.ps1tex.tpage = getTPage(outTim.format, 0, x, y);
				tex.uoffs = (u_char)((x % 64) << (2 - outTim.format));
				tex.voffs = (u_char)(y & 0xFF);
				if (texture_isRepeatable(miptex))
					tex.ps1tex.twin = getTexWindow(tex.uoffs, tex.voffs, tex.w, tex.h);	// TODO: figure out the right offsets that are multiples of w and h; NOTE: if uoffs has to >> 1, then w probably has to as well
				else
					tex.ps1tex.twin = getTexWindow(0, 0, 0, 0);
			}
		}

		// Use the smallest mip level for analysis, should be good enough
		analyze_texture(world->textures[texNum * 4 + 3], (miptex->width >> 3) * (miptex->height >> 3), palette, tex);

		outTextures.push_back(tex);
	}

	// Identify animated textures
	std::unordered_map<std::string, std::unordered_map<int, int>> animationFrames;
	for (int texNum = 0; texNum < world->mipheader.numtex; ++texNum)
	{
		miptex_t* miptex = &world->miptexes[texNum];
		if (miptex->name[0] == '+')
		{
			// Animated texture
			int frameNum = miptex->name[1] - '0';
			if (frameNum >= 0 && frameNum <= 9)
			{
				std::string animName = &miptex->name[2];
				animationFrames[animName][frameNum] = texNum;
			}
			else
			{
				frameNum = miptex->name[1] - 'a';
				if (frameNum >= 0 && frameNum <= 9)
				{
					std::string animName = std::string(&miptex->name[2]) + "_alt";
					animationFrames[animName][frameNum] = texNum;
				}
			}
		}
	}

	// Link animated texture frames together
	for (auto animIter = animationFrames.cbegin(); animIter != animationFrames.cend(); ++animIter)
	{
		const auto& frames = animIter->second;
		for (auto frameIter = frames.cbegin(); frameIter != frames.cend(); ++frameIter)
		{
			int frameNum = frameIter->first;
			int texNum = frameIter->second;
			auto& texture = outTextures[texNum];

			const auto& nextFrameIter = frames.find(frameNum + 1);
			if (nextFrameIter != frames.cend())
				texture.ps1tex.nextframe = nextFrameIter->second;
			else
				texture.ps1tex.nextframe = frames.find(0)->second;
		}
	}

	sprintf_s(path, _MAX_PATH, "atlas-%s.tim", world->name);
	tim::ExportFile(path, &outTim);

	free(outTim.imgData);
	free(outTim.clutData);
	return true;
}

bool texture_isRepeatable(const miptex_t* miptex)
{
	// Check if the texture is square
	if (miptex->width != miptex->height)
		return false;

	// Check if width and height are powers of two (and a multiple of 8)
	unsigned int pot = 8;
	while (pot < miptex->width)
		pot <<= 1;

	return pot == miptex->width;
}