#include <memory.h>
#include <stdlib.h>
#include <stdio.h>
#include <vector>
#include <unordered_map>
#include "bsp.h"
#include "rectpack/finders_interface.h"
#include "ps1types.h"
#include "ps1bsp.h"

static char path[_MAX_PATH];

int process_entities(const world_t *world)
{
	printf("Entities list:\n%s\n", world->entities);
	return 1;
}

int process_textures(const world_t* world)
{
	using spaces_type = rectpack2D::empty_spaces<false>;
	using rect_type = rectpack2D::output_rect_t<spaces_type>;

	auto report_successful = [](rect_type&) {
		return rectpack2D::callback_result::CONTINUE_PACKING;
	};

	auto report_unsuccessful = [](rect_type&) {
		return rectpack2D::callback_result::ABORT_PACKING;
	};

	const auto max_side = 512;	// Max height of PS1 VRAM. 8-bit textures take up half the horizontal space so this is 512x256 in practice, or a quarter of the PS1's VRAM allocation.
	const auto discard_step = -4;

	std::vector<rect_type> rectangles;

	// 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);

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

		if (strcmp(miptex->name, "clip") && strcmp(miptex->name, "trigger"))
			rectangles.emplace_back(rectpack2D::rect_xywh(0, 0, miptex->width >> ps1mip, miptex->height >> ps1mip));
		else
			rectangles.emplace_back(rectpack2D::rect_xywh(0, 0, 0, 0));
	}

	// Automatic atlas packing. Nice but it tries to make a square atlas which is not what we want. (This is solved by hacking the header itself)
	const auto result_size = rectpack2D::find_best_packing<spaces_type>(
		rectangles,
		rectpack2D::make_finder_input(
			max_side,
			discard_step,
			report_successful,
			report_unsuccessful,
			rectpack2D::flipping_option::DISABLED
		)
	);

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

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

	// 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?
			continue;

		char* outName = miptex->name;
		if (*outName == '*' || *outName == '+')
			outName++;

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

			FILE* fraw;
			sprintf_s(path, _MAX_PATH, "textures/%s-%s-mip%d-%dx%d.raw", world->name, outName, mipLevel, miptex->width >> mipLevel, miptex->height >> mipLevel);
			fopen_s(&fraw, path, "wb");
			if (fraw != NULL)
			{
				size_t numBytes = (miptex->width * miptex->height) >> mipLevel;
				fwrite(texBytes, sizeof(unsigned char), numBytes, fraw);
				fclose(fraw);
			}

			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
			{
				//printf("Writing texture %s mip %d to position: (%d, %d) w = %d, h = %d\n", miptex->name, mipLevel, rectangle.x, rectangle.y, rectangle.w, rectangle.h);
				for (int y = 0; y < rectangle.h; ++y)
				{
					memcpy_s(atlas + ((rectangle.y + y) * result_size.w + rectangle.x), rectangle.w * sizeof(unsigned char), texBytes + (y * rectangle.w), rectangle.w * sizeof(unsigned char));
				}
			}
		}
	}

	FILE* fatlas;
	sprintf_s(path, _MAX_PATH, "%s-atlas-%dx%d.raw", world->name, result_size.w, result_size.h);
	fopen_s(&fatlas, path, "wb");
	if (fatlas != NULL)
	{
		fwrite(atlas, sizeof(unsigned char), result_size.w * result_size.h, fatlas);
		fclose(fatlas);
	}

	free(atlas);
	return 1;
}

int process_vertices(const world_t* world)
{
	vec3_t min = { FLT_MAX, FLT_MAX, FLT_MAX }, max = { -FLT_MAX, -FLT_MAX, -FLT_MAX };
	for (int vertIdx = 0; vertIdx < world->numVertices; ++vertIdx)
	{
		vertex_t* vert = &world->vertices[vertIdx];
		if (vert->X > max.x) max.x = vert->X;
		if (vert->Y > max.y) max.y = vert->Y;
		if (vert->Z > max.z) max.z = vert->Z;
		if (vert->X < min.x) min.x = vert->X;
		if (vert->Y < min.y) min.y = vert->Y;
		if (vert->Z < min.z) min.z = vert->Z;
	}

	printf("%d vertices, %d faces, min = (%f, %f, %f), max = (%f, %f, %f)\n", world->numVertices, world->numFaces, min.x, min.y, min.z, max.x, max.y, max.z);

	const int fixedScale = 1 << 14;
	int fixedMin[3] = { (int)(min.x * fixedScale), (int)(min.y * fixedScale), (int)(min.z * fixedScale) };
	int fixedMax[3] = { (int)(max.x * fixedScale), (int)(max.y * fixedScale), (int)(max.z * fixedScale) };
	printf("Fixed point min = (%d, %d, %d), max = (%d, %d, %d)\n", fixedMin[0], fixedMin[1], fixedMin[2], fixedMax[0], fixedMax[1], fixedMax[2]);

	return 1;
}

typedef struct
{
	vertex_t* vertex;
	unsigned short index;
} vertexref_t;

// Determines a floating origin for the given leaf
static void leaf_zone(const dleaf_t* leaf, short zone[3])
{
	const unsigned short mask = 0xFE00;	// Zero out the 9 least significant bits

	short midX = (leaf->bound.min[0] + leaf->bound.max[0]) / 2;
	short midY = (leaf->bound.min[1] + leaf->bound.max[1]) / 2;
	short midZ = (leaf->bound.min[2] + leaf->bound.max[2]) / 2;
	
	zone[0] = midX & mask;
	zone[1] = midY & mask;
	zone[2] = midZ & mask;
}

static unsigned char sample_lightmap(const world_t* world, const face_t *face, const BoundBox& bounds, const Vec3& point)
{
	if (face->lightmap < 0)
		return 0;

	const unsigned char* lightmap = &world->lightmap[face->lightmap];
	const plane_t* plane = &world->planes[face->plane_id];

	int width, height;
	float u, v;
	switch (plane->type)
	{
	case 0:
	case 3:
		// Towards X
		width = (int)(ceil(bounds.max.y / 16) - floor(bounds.min.y / 16)) * 16;
		height = (int)(ceil(bounds.max.z / 16) - floor(bounds.min.z / 16)) * 16;
		u = (point.y - bounds.min.y) / (bounds.max.y - bounds.min.y);
		v = (point.z - bounds.min.z) / (bounds.max.z - bounds.min.z);
		break;
	case 1:
	case 4:
		// Towards Y
		width = (int)(ceil(bounds.max.x / 16) - floor(bounds.min.x / 16)) * 16;
		height = (int)(ceil(bounds.max.z / 16) - floor(bounds.min.z / 16)) * 16;
		u = (point.x - bounds.min.x) / (bounds.max.x - bounds.min.x);
		v = (point.z - bounds.min.z) / (bounds.max.z - bounds.min.z);
		break;
	case 2:
	case 5:
		// Towards Z
		width = (int)(ceil(bounds.max.x / 16) - floor(bounds.min.x / 16)) * 16;
		height = (int)(ceil(bounds.max.y / 16) - floor(bounds.min.y / 16)) * 16;
		u = (point.x - bounds.min.x) / (bounds.max.x - bounds.min.x);
		v = (point.y - bounds.min.y) / (bounds.max.y - bounds.min.y);
		break;
	default:
		printf("Error: unknown plane type %d\n", plane->type);
		return 0;
	}

	height >>= 4;
	width >>= 4;

	return lightmap[(int)(v * (width + 1) + u)];
}

static void export_lightmap(const world_t* world, const face_t* face, const BoundBox& bounds, int faceIdx)
{
	if (face->lightmap < 0)
		return;

	const unsigned char* lightmap = &world->lightmap[face->lightmap];
	const plane_t* plane = &world->planes[face->plane_id];

	int width, height;
	switch (plane->type)
	{
	case 0:
	case 3:
		// Towards X
		width = (int)(ceil(bounds.max.y / 16) - floor(bounds.min.y / 16));
		height = (int)(ceil(bounds.max.z / 16) - floor(bounds.min.z / 16));
		break;
	case 1:
	case 4:
		// Towards Y
		width = (int)(ceil(bounds.max.x / 16) - floor(bounds.min.x / 16));
		height = (int)(ceil(bounds.max.z / 16) - floor(bounds.min.z / 16));
		break;
	case 2:
	case 5:
		// Towards Z
		width = (int)(ceil(bounds.max.x / 16) - floor(bounds.min.x / 16));
		height = (int)(ceil(bounds.max.y / 16) - floor(bounds.min.y / 16));
		break;
	default:
		printf("Error: unknown plane type %d\n", plane->type);
		return;
	}

	width += 1;

	char path[_MAX_PATH];
	sprintf_s(path, _MAX_PATH, "lightmap_face%d_e%d_PT%d_%dx%d.raw", faceIdx, face->ledge_num, plane->type, width, height);
	FILE* flm;
	fopen_s(&flm, path, "wb");
	if (!flm)
		return;

	for (int y = 0; y < height; ++y)
	{
		fwrite(&lightmap[y * width], sizeof(unsigned char), width, flm);
	}

	fclose(flm);
}

template<class TData> size_t writeMapData(const std::vector<TData>& data, ps1bsp_dentry_t& dentry, FILE* f)
{
	dentry.offset = (unsigned int)ftell(f);
	dentry.size = sizeof(TData) * data.size();
	return fwrite(data.data(), sizeof(TData), data.size(), f);
}

int process_faces(const world_t* world)
{
	// Write some data to a file
	FILE* fbsp;
	fopen_s(&fbsp, "test.ps1bsp", "wb");
	if (!fbsp)
		return 0;

	ps1bsp_header_t outHeader = { 0 };	// Write an empty placeholder header first
	outHeader.version = 1;
	fwrite(&outHeader, sizeof(ps1bsp_header_t), 1, fbsp);

	// Write vertex data to a file (no vertex splitting yet)
	std::vector<ps1bsp_vertex_t> outVertices;
	for (unsigned short i = 0; i < world->numVertices; ++i)
	{
		vertex_t* inVertex = &world->vertices[i];

		ps1bsp_vertex_t outVertex = { 0 };
		// Ensure we don't overflow 16-bit short values. Most Quake maps will stay within these bounds so it *should* be fine (for now).
		if (inVertex->X > -8192 && inVertex->X < 8192 && inVertex->Y > -8192 && inVertex->Y < 8192 && inVertex->Z > -8192 && inVertex->Z < 8192)
		{
			outVertex.x = (short)(inVertex->X * 4);
			outVertex.y = (short)(inVertex->Y * 4);
			outVertex.z = (short)(inVertex->Z * 4);
		}
		else
		{
			printf("Error: vertices found outside of acceptable range: (%f, %f, %f)\n", inVertex->X, inVertex->Y, inVertex->Z);
			fclose(fbsp);
			return 0;
		}

		outVertices.push_back(outVertex);
	}
	
	std::vector<ps1bsp_face_t> outFaces;
	std::vector<ps1bsp_facevertex_t> outFaceVertices;
	for (int faceIdx = 0; faceIdx < world->numFaces; ++faceIdx)
	{
		face_t* face = &world->faces[faceIdx];

		ps1bsp_face_t outFace = { 0 };
		outFace.firstFaceVertex = (unsigned short)outFaceVertices.size();

		// Traverse the list of face edges to collect all of the face's vertices
		BoundBox bounds;
		for (int edgeListIdx = 0; edgeListIdx < face->ledge_num; ++edgeListIdx)
		{
			int edgeIdx = world->edgeList[face->ledge_id + edgeListIdx];

			unsigned short vertIndex = edgeIdx > 0 ?
				world->edges[edgeIdx].vertex0 :
				world->edges[-edgeIdx].vertex1;

			ps1bsp_facevertex_t faceVertex;
			faceVertex.index = vertIndex;
			faceVertex.light = 0;
			outFaceVertices.push_back(faceVertex);

			// Calculate bounding box of this face
			const vertex_t* vertex = &world->vertices[vertIndex];
			if (edgeListIdx == 0)
				bounds.init(vertex->toVec());
			else
				bounds.includePoint(vertex->toVec());
		}

		// Sample lightmap contribution of this face on each vertex
		for (size_t faceVertIdx = outFace.firstFaceVertex; faceVertIdx < outFaceVertices.size(); ++faceVertIdx)
		{
			ps1bsp_facevertex_t& faceVertex = outFaceVertices[faceVertIdx];
			unsigned char lightmap = sample_lightmap(world, face, bounds, world->vertices[faceVertex.index].toVec());
			faceVertex.light = lightmap + (0xFF - face->baselight);

			if (face->lightmap >= 0)
			{
				ps1bsp_vertex_t& vertex = outVertices[faceVertex.index];
				*(unsigned short*)(&vertex.baseLight) += faceVertex.light;
				vertex.r++;
			}
		}

		// For visualizing and debugging lightmaps
		//if (face->ledge_num >= 10)
		//	export_lightmap(world, face, bounds, faceIdx);

		outFace.numFaceVertices = (unsigned char)(outFaceVertices.size() - outFace.firstFaceVertex);
		outFaces.push_back(outFace);
	}

	// Average the lightmap values for each vertex
	for (auto iter = outVertices.begin(); iter != outVertices.end(); ++iter)
	{
		unsigned char count = (*iter).r;
		if (count == 0)
			continue;

		unsigned short accumulate = *(unsigned short*)(&(*iter).baseLight);
		(*iter).baseLight = accumulate / count;
		(*iter).r = 0;
	}

	// Write collected data to file and update header info
	writeMapData(outVertices, outHeader.vertices, fbsp);
	writeMapData(outFaces, outHeader.faces, fbsp);
	writeMapData(outFaceVertices, outHeader.faceVertices, fbsp);

	// Write final header
	fseek(fbsp, 0, SEEK_SET);
	fwrite(&outHeader, sizeof(ps1bsp_header_t), 1, fbsp);
	fclose(fbsp);

	printf("PS1BSP: wrote %d vertices, %d indices, %d faces\n", outVertices.size(), outFaceVertices.size(), outFaces.size());

	return 1;
}

int process_bsp(const world_t *world)
{
	// Test reading the entity string data
	if (!process_entities(world))
	{
		return 0;
	}

	// Test exporting texture data
	if (!process_textures(world))
	{
		return 0;
	}

	// Inspect vertex data
	if (!process_vertices(world))
	{
		return 0;
	}

	// Inspect faces/edges data
	if (!process_faces(world))
	{
		return 0;
	}

	return 1;
}

int load_bsp(const char* bspname, world_t* world)
{
	FILE* f;
	dheader_t* header = &world->header;

	world->name = bspname;

	sprintf_s(path, _MAX_PATH, "%s.bsp", bspname);
	fopen_s(&f, path, "rb");
	if (f == NULL)
		return 0;

	fread(header, sizeof(dheader_t), 1, f);
	printf("Header model version: %d\n", header->version);

	// Load entities
	fseek(f, header->entities.offset, SEEK_SET);
	world->entitiesLength = header->entities.size + 1;
	world->entities = (char*)malloc(world->entitiesLength * sizeof(char));
	if (world->entities == NULL)
		return 0;

	memset(world->entities, 0, world->entitiesLength * sizeof(char));
	fread(world->entities, sizeof(char), world->entitiesLength, f);

	// Load textures
	mipheader_t* mipheader = &world->mipheader;
	fseek(f, header->miptex.offset, SEEK_SET);
	fread(&mipheader->numtex, sizeof(long), 1, f);
	mipheader->offset = (long*)malloc(mipheader->numtex * sizeof(long));
	if (mipheader->offset == NULL)
		return 0;

	fread(mipheader->offset, sizeof(long), mipheader->numtex, f);

	world->miptexes = (miptex_t*)malloc(mipheader->numtex * sizeof(miptex_t));
	if (world->miptexes == NULL)
		return 0;

	const int numMipLevels = 4;
	world->textures = (unsigned char**)malloc(mipheader->numtex * numMipLevels * sizeof(unsigned char*));
	if (world->textures == NULL)
		return 0;

	memset(world->textures, 0, mipheader->numtex * numMipLevels * sizeof(unsigned char*));

	for (int texNum = 0; texNum < mipheader->numtex; ++texNum)
	{
		miptex_t* miptex = &world->miptexes[texNum];

		unsigned long miptexOffset = header->miptex.offset + mipheader->offset[texNum];
		fseek(f, miptexOffset, SEEK_SET);
		fread(miptex, sizeof(miptex_t), 1, f);

		for (int mipLevel = 0; mipLevel < numMipLevels; ++mipLevel)
		{
			unsigned long mipOffset = *(&miptex->offset1 + mipLevel);
			fseek(f, miptexOffset + mipOffset, SEEK_SET);

			size_t numBytes = (miptex->width * miptex->height) >> mipLevel;
			unsigned char* texBytes = (unsigned char*)malloc(sizeof(unsigned char) * numBytes);
			if (texBytes == NULL)
				return 0;

			fread(texBytes, sizeof(unsigned char), numBytes, f);
			world->textures[texNum * numMipLevels + mipLevel] = texBytes;
		}
	}

	// Load planes
	world->numPlanes = header->planes.size / sizeof(plane_t);
	world->planes = (plane_t*)malloc(header->planes.size);
	if (world->planes == NULL)
		return 0;

	fseek(f, header->planes.offset, SEEK_SET);
	fread(world->planes, sizeof(plane_t), world->numPlanes, f);
	
	// Load vertices
	world->numVertices = header->vertices.size / sizeof(vertex_t);
	world->vertices = (vertex_t*)malloc(header->vertices.size);
	if (world->vertices == NULL)
		return 0;

	fseek(f, header->vertices.offset, SEEK_SET);
	fread(world->vertices, sizeof(vertex_t), world->numVertices, f);

	// Load edges
	world->numEdges = header->edges.size / sizeof(edge_t);
	world->edges = (edge_t*)malloc(header->edges.size);
	if (world->edges == NULL)
		return 0;

	fseek(f, header->edges.offset, SEEK_SET);
	fread(world->edges, sizeof(edge_t), world->numEdges, f);

	world->edgeListLength = header->ledges.size / sizeof(int);
	world->edgeList = (int*)malloc(header->ledges.size);
	if (world->edgeList == NULL)
		return 0;

	fseek(f, header->ledges.offset, SEEK_SET);
	fread(world->edgeList, sizeof(int), world->edgeListLength, f);

	// Load faces
	world->numFaces = header->faces.size / sizeof(face_t);
	world->faces = (face_t*)malloc(header->faces.size);
	if (world->faces == NULL)
		return 0;

	fseek(f, header->faces.offset, SEEK_SET);
	fread(world->faces, sizeof(face_t), world->numFaces, f);

	world->faceListLength = header->lface.size / sizeof(unsigned short);
	world->faceList = (unsigned short*)malloc(header->lface.size);
	if (world->faceList == NULL)
		return 0;

	// Load visibility list
	fseek(f, header->lface.offset, SEEK_SET);
	fread(world->faceList, sizeof(unsigned short), world->faceListLength, f);

	world->visListLength = header->visilist.size / sizeof(unsigned char);
	world->visList = (unsigned char*)malloc(header->visilist.size);
	if (world->visList == NULL)
		return 0;

	fseek(f, header->visilist.offset, SEEK_SET);
	fread(world->visList, sizeof(unsigned char), world->visListLength, f);

	// Load nodes
	world->numNodes = header->nodes.size / sizeof(node_t);
	world->nodes = (node_t*)malloc(header->nodes.size);
	if (world->nodes == NULL)
		return 0;

	fseek(f, header->nodes.offset, SEEK_SET);
	fread(world->nodes, sizeof(node_t), world->numNodes, f);

	// Load leaves
	world->numLeaves = header->leaves.size / sizeof(dleaf_t);
	world->leaves = (dleaf_t*)malloc(header->leaves.size);
	if (world->leaves == NULL)
		return 0;

	fseek(f, header->leaves.offset, SEEK_SET);
	fread(world->leaves, sizeof(dleaf_t), world->numLeaves, f);

	// Load lightmaps
	world->lightmapLength = header->lightmaps.size / sizeof(unsigned char);
	world->lightmap = (unsigned char*)malloc(header->lightmaps.size);
	if (world->lightmap == NULL)
		return 0;

	fseek(f, header->lightmaps.offset, SEEK_SET);
	fread(world->lightmap, sizeof(unsigned char), world->lightmapLength, f);

	fclose(f);
	return 1;
}

void free_bsp(world_t* world)
{
	free(world->leaves);
	free(world->nodes);
	free(world->visList);

	free(world->faces);
	free(world->faceList);
	free(world->edges);
	free(world->edgeList);
	free(world->vertices);
	free(world->planes);

	for (int i = 0; i < world->mipheader.numtex; ++i)
	{
		free(world->textures[i]);
	}

	free(world->textures);
	free(world->miptexes);
	free(world->mipheader.offset);

	free(world->entities);
}

int main(int argc, char** argv)
{
	world_t world = { 0 };
	if (!load_bsp(argv[1], &world))
		return 1;

	int result = process_bsp(&world);

	free_bsp(&world);
	return !result;
}