#include <glad/glad.h>
#include <iostream>
#include <string>
#include <fstream>
#include <sstream>
#include <vector>
#include <assimp/Importer.hpp>
#include <assimp/scene.h>
#include <assimp/postprocess.h>
#include <stb_image.h>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtc/type_ptr.hpp>
#include <windows.h>
#include <stdarg.h>
#include "XCEngine/Debug/Logger.h"
#include "XCEngine/Debug/ConsoleLogSink.h"
#include "XCEngine/Debug/FileLogSink.h"
#include "XCEngine/RHI/OpenGL/OpenGLShader.h"
#include "XCEngine/RHI/OpenGL/OpenGLDevice.h"
#include "XCEngine/RHI/OpenGL/OpenGLPipelineState.h"
#include "XCEngine/RHI/OpenGL/OpenGLCommandList.h"

using namespace XCEngine::Debug;
using namespace XCEngine::RHI;

#pragma comment(lib, "opengl32.lib")
#pragma comment(lib, "gdi32.lib")

const unsigned int SCR_WIDTH = 1280;
const unsigned int SCR_HEIGHT = 720;

void Log(const char* format, ...) {
    char buffer[1024];
    va_list args;
    va_start(args, format);
    vsnprintf(buffer, sizeof(buffer), format, args);
    va_end(args);
    
    Logger::Get().Debug(LogCategory::Rendering, buffer);
}

struct Vertex
{
	glm::vec3 Position;
	glm::vec3 Normal;
	glm::vec2 TexCoords;
};

struct Texture
{
	unsigned int id;
	std::string type;
	std::string path;
};

class Mesh
{
public:
	std::vector<Vertex> vertices;
	std::vector<unsigned int> indices;
	std::vector<Texture> textures;
	unsigned int VAO;

	Mesh(std::vector<Vertex> _vertices, std::vector<unsigned int> _indices, std::vector<Texture> _textures)
	{
		vertices = _vertices;
		indices = _indices;
		textures = _textures;
		setupMesh();
	}

	void Draw(OpenGLShader* shader)
	{
		unsigned int diffuseNr = 1;
		unsigned int specularNr = 1;

		for (unsigned int i = 0; i < textures.size(); i++)
		{
			glActiveTexture(GL_TEXTURE0 + i);
			std::string number;
			std::string name = textures[i].type;
			if (name == "texture_diffuse")
				number = std::to_string(diffuseNr++);
			else if (name == "texture_specular")
				number = std::to_string(specularNr++);

			glUniform1i(glGetUniformLocation(shader->GetID(), ("material." + name + number).c_str()), i);
			glBindTexture(GL_TEXTURE_2D, textures[i].id);
		}

		glBindVertexArray(VAO);
		glDrawElements(GL_TRIANGLES, static_cast<unsigned int>(indices.size()), GL_UNSIGNED_INT, 0);
		glActiveTexture(GL_TEXTURE0);
	}

private:
	unsigned int VBO, EBO;

	void setupMesh()
	{
		glGenVertexArrays(1, &VAO);
		glGenBuffers(1, &VBO);
		glGenBuffers(1, &EBO);

		glBindVertexArray(VAO);
		glBindBuffer(GL_ARRAY_BUFFER, VBO);
		glBufferData(GL_ARRAY_BUFFER, vertices.size() * sizeof(Vertex), &vertices[0], GL_STATIC_DRAW);

		glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO);
		glBufferData(GL_ELEMENT_ARRAY_BUFFER, indices.size() * sizeof(unsigned int), &indices[0], GL_STATIC_DRAW);

		glEnableVertexAttribArray(0);
		glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, sizeof(Vertex), (void*)0);
		glEnableVertexAttribArray(1);
		glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, sizeof(Vertex), (void*)offsetof(Vertex, Normal));
		glEnableVertexAttribArray(2);
		glVertexAttribPointer(2, 2, GL_FLOAT, GL_FALSE, sizeof(Vertex), (void*)offsetof(Vertex, TexCoords));

		glBindVertexArray(0);
	}
};

unsigned int TextureFromFile(const char* path, const std::string& directory)
{
	std::string filename = directory + '/' + std::string(path);
	unsigned int textureID;
	glGenTextures(1, &textureID);

	int width, height, nrComponents;
	stbi_set_flip_vertically_on_load(true);
	unsigned char* data = stbi_load(filename.c_str(), &width, &height, &nrComponents, 0);
	if (data)
	{
		GLenum format = (nrComponents == 4) ? GL_RGBA : GL_RGB;
		glBindTexture(GL_TEXTURE_2D, textureID);
		glTexImage2D(GL_TEXTURE_2D, 0, format, width, height, 0, format, GL_UNSIGNED_BYTE, data);
		glGenerateMipmap(GL_TEXTURE_2D);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
		Log("Texture loaded: %s (%dx%d)", filename.c_str(), width, height);
		stbi_image_free(data);
	}
	else
	{
		Log("Failed to load texture: %s", filename.c_str());
		stbi_image_free(data);
	}
	return textureID;
}

class Model
{
public:
	std::vector<Texture> textures_loaded;
	std::vector<Mesh> meshes;
	std::string directory;

	Model(std::string const& path)
	{
		loadModel(path);
	}

	void Draw(OpenGLShader* shader)
	{
		for (auto& mesh : meshes)
			mesh.Draw(shader);
	}

private:
	void loadModel(std::string const& path)
	{
		Assimp::Importer importer;
		const aiScene* scene = importer.ReadFile(path, aiProcess_Triangulate | aiProcess_GenSmoothNormals | aiProcess_FlipUVs);
		if (!scene || scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE || !scene->mRootNode)
		{
			Log("Failed to load model: %s", path.c_str());
			return;
		}
		Log("Model loaded: %s", path.c_str());
		directory = path.substr(0, path.find_last_of('/'));
		processNode(scene->mRootNode, scene);
	}

	void processNode(aiNode* node, const aiScene* scene)
	{
		for (unsigned int i = 0; i < node->mNumMeshes; i++)
		{
			aiMesh* mesh = scene->mMeshes[node->mMeshes[i]];
			meshes.push_back(processMesh(mesh, scene));
		}
		for (unsigned int i = 0; i < node->mNumChildren; i++)
			processNode(node->mChildren[i], scene);
	}

	Mesh processMesh(aiMesh* mesh, const aiScene* scene)
	{
		std::vector<Vertex> vertices;
		std::vector<unsigned int> indices;
		std::vector<Texture> textures;

		for (unsigned int i = 0; i < mesh->mNumVertices; i++)
		{
			Vertex vertex;
			vertex.Position = glm::vec3(mesh->mVertices[i].x, mesh->mVertices[i].y, mesh->mVertices[i].z);
			vertex.Normal = glm::vec3(mesh->mNormals[i].x, mesh->mNormals[i].y, mesh->mNormals[i].z);
			if (mesh->mTextureCoords[0])
				vertex.TexCoords = glm::vec2(mesh->mTextureCoords[0][i].x, mesh->mTextureCoords[0][i].y);
			else
				vertex.TexCoords = glm::vec2(0.0f, 0.0f);
			vertices.push_back(vertex);
		}

		for (unsigned int i = 0; i < mesh->mNumFaces; i++)
		{
			aiFace face = mesh->mFaces[i];
			for (unsigned int j = 0; j < face.mNumIndices; j++)
				indices.push_back(face.mIndices[j]);
		}

		aiMaterial* material = scene->mMaterials[mesh->mMaterialIndex];
		std::vector<Texture> diffuseMaps = loadMaterialTextures(material, aiTextureType_DIFFUSE, "texture_diffuse");
		textures.insert(textures.end(), diffuseMaps.begin(), diffuseMaps.end());
		std::vector<Texture> specularMaps = loadMaterialTextures(material, aiTextureType_SPECULAR, "texture_specular");
		textures.insert(textures.end(), specularMaps.begin(), specularMaps.end());

		return Mesh(vertices, indices, textures);
	}

	std::vector<Texture> loadMaterialTextures(aiMaterial* mat, aiTextureType type, std::string typeName)
	{
		std::vector<Texture> textures;
		for (unsigned int i = 0; i < mat->GetTextureCount(type); i++)
		{
			aiString str;
			mat->GetTexture(type, i, &str);
			bool skip = false;
			for (auto& tex : textures_loaded)
			{
				if (tex.path == str.C_Str())
				{
					textures.push_back(tex);
					skip = true;
					break;
				}
			}
			if (!skip)
			{
				Texture texture;
				texture.id = TextureFromFile(str.C_Str(), directory);
				texture.type = typeName;
				texture.path = str.C_Str();
				textures.push_back(texture);
				textures_loaded.push_back(texture);
			}
		}
		return textures;
	}
};

Model* model = nullptr;
OpenGLShader* shader = nullptr;
OpenGLDevice* device = nullptr;
OpenGLPipelineState* pipeline = nullptr;
OpenGLCommandList* cmdList = nullptr;
int frameCount = 0;

void Initialize()
{
	char exePath[MAX_PATH];
	GetModuleFileNameA(NULL, exePath, MAX_PATH);
	std::string exeDir = std::string(exePath);
	exeDir = exeDir.substr(0, exeDir.find_last_of("\\/"));
	SetCurrentDirectoryA(exeDir.c_str());

	Log("OpenGL Test Application Started");

	pipeline = new OpenGLPipelineState();
	ViewportState viewport;
	viewport.width = SCR_WIDTH;
	viewport.height = SCR_HEIGHT;
	pipeline->SetViewport(viewport);
	pipeline->SetClearColor(0.1f, 0.1f, 0.1f, 1.0f);
	pipeline->Apply();

	model = new Model("res/models/backpack/backpack.obj");
	shader = new OpenGLShader();
	shader->CompileFromFile("Shaders/vertexshader.glsl", "Shaders/fragmentshader.glsl");

	cmdList = new OpenGLCommandList();

	Log("Initialization complete");
}

bool SaveScreenshot(const char* filename)
{
	Log("Saving screenshot: %s", filename);
	
	int width = SCR_WIDTH;
	int height = SCR_HEIGHT;
	
	std::vector<unsigned char> pixels(width * height * 3);
	glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE, pixels.data());
	
	FILE* fp = fopen(filename, "wb");
	if (!fp) {
		Log("Failed to open file for screenshot: %s", filename);
		return false;
	}
	
	fprintf(fp, "P6\n%d %d\n255\n", width, height);
	
	std::vector<unsigned char> flippedPixels(width * height * 3);
	for (int y = 0; y < height; y++) {
		memcpy(&flippedPixels[y * width * 3], &pixels[(height - 1 - y) * width * 3], width * 3);
	}
	
	fwrite(flippedPixels.data(), 1, width * height * 3, fp);
	fclose(fp);
	
	Log("Screenshot saved successfully");
	return true;
}

void Render()
{
	pipeline->Clear(3); // COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT

	shader->Use();

	glm::mat4 projection = glm::perspective(glm::radians(45.0f), (float)SCR_WIDTH / (float)SCR_HEIGHT, 0.1f, 100.0f);
	glm::mat4 modelMat = glm::mat4(1.0f);
	glm::mat4 view = glm::lookAt(glm::vec3(0, 0, 3), glm::vec3(0, 0, 0), glm::vec3(0, 1, 0));

	cmdList->SetShader(shader);
	cmdList->SetUniformMat4("view", &view[0][0]);
	cmdList->SetUniformMat4("model", &modelMat[0][0]);
	cmdList->SetUniformMat4("projection", &projection[0][0]);
	cmdList->SetUniformVec3("viewPos", 0.0f, 0.0f, 3.0f);
	cmdList->SetUniformFloat("material.shininess", 32.0f);
	cmdList->SetUniformVec3("dirLight.direction", 0.0f, -1.0f, 0.0f);
	cmdList->SetUniformVec3("dirLight.ambient", 0.3f, 0.3f, 0.3f);
	cmdList->SetUniformVec3("dirLight.diffuse", 0.8f, 0.8f, 0.8f);
	cmdList->SetUniformVec3("dirLight.specular", 0.5f, 0.5f, 0.5f);

	model->Draw(shader);
}

int main()
{
    Logger::Get().AddSink(std::make_unique<FileLogSink>("OpenGL_engine_log.txt"));
    Logger::Get().SetMinimumLevel(LogLevel::Debug);
    
    AllocConsole();
    freopen("CONOUT$", "w", stdout);

	device = new OpenGLDevice();
	if (!device->CreateRenderWindow(SCR_WIDTH, SCR_HEIGHT, "XCRender", false))
	{
		Log("Failed to create window");
		return -1;
	}

	const RHIDeviceInfo& info = device->GetDeviceInfo();
	Log("OpenGL Version: %ls", info.version.c_str());
	Log("Renderer: %ls", info.renderer.c_str());

	Initialize();

	device->PollEvents();
	
	MSG msg = {};
	bool shouldClose = false;
	while (!shouldClose)
	{
		while (PeekMessageW(&msg, NULL, 0, 0, PM_REMOVE))
		{
			if (msg.message == WM_QUIT)
			{
				shouldClose = true;
				break;
			}
			TranslateMessage(&msg);
			DispatchMessageW(&msg);
		}

		if (shouldClose)
			break;

		Render();
		device->SwapBuffers();
		device->PollEvents();
		
		frameCount++;
		
		if (frameCount == 30) {
			Log("Saving screenshot at frame %d", frameCount);
			SaveScreenshot("screenshot.ppm");
			PostMessageW(device->GetWindow(), WM_CLOSE, 0, 0);
		}
	}

	Log("Application closed");

	delete cmdList;
	delete pipeline;
	delete device;
	delete shader;
	delete model;
	
    Logger::Get().Shutdown();
	return 0;
}