#include <iostream>

#include <glm/glm.hpp>
#include <glm/gtc/type_ptr.hpp>

#include "Renderer.h"
#include "Shader.h"

#include <chrono>
#include <random>

constexpr int WORK_GROUP_SIZE = 64;
constexpr int PARTICLE_COUNT = WORK_GROUP_SIZE * 15625;

constexpr int SWAP_INTERVAL = 1;

constexpr int TRANSFORMATION_COUNT = 3;

static const std::vector<glm::mat4> SIERPINSKI_TRIANGLE = {
	{
		0.5f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		0.5f,
		0.0f,
		0.36f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		1.0f,
	},
	{
		0.5f,
		0.0f,
		0.0f,
		-0.5f,
		0.0f,
		0.5f,
		0.0f,
		-0.5f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		1.0f,
	},
	{
		0.5f,
		0.0f,
		0.0f,
		0.5f,
		0.0f,
		0.5f,
		0.0f,
		-0.5f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		0.0f,
		1.0f,
	},
};

class Timer {
  public:
	Timer() {
		Reset();
	}
	void Reset() {
		m_Start = std::chrono::high_resolution_clock::now();
	}
	float Elapsed() const {
		return std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::high_resolution_clock::now() - m_Start).count() *
			   0.001f * 0.001f;
	}
	float ElapsedMillis() const {
		return std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::high_resolution_clock::now() - m_Start).count() *
			   0.001f;
	}

  private:
	std::chrono::time_point<std::chrono::high_resolution_clock> m_Start;
};

static uint32_t s_ComputeShader = -1;
static uint32_t s_GraphicsShader = -1;
static const std::filesystem::path s_ComputeShaderPath = "Shaders/Compute.glsl";
static const std::filesystem::path s_VertexShaderPath = "Shaders/Vertex.glsl";
static const std::filesystem::path s_FragmentShaderPath = "Shaders/Fragment.glsl";

static std::random_device s_RandomDevice;
static std::mt19937 s_Generator(s_RandomDevice());
static std::uniform_real_distribution<float> s_Distrib(0, 1);

static void ErrorCallback(int error, const char* description) {
	std::cerr << "Error: " << description << std::endl;
}

static void ApplyTransforms(const std::vector<glm::mat4>& transformations) {
	glUseProgram(s_ComputeShader);
	glUniformMatrix4fv(1, transformations.size(), false, glm::value_ptr(transformations[0]));
}

struct Transform {
	float m_ScaleX;
	float m_ScaleY;
	float m_ScaleZ;

	float m_RotationX;
	float m_RotationY;
	float m_RotationZ;

	float m_ShearXY;
	float m_ShearXZ;
	float m_ShearYX;
	float m_ShearYZ;
	float m_ShearZX;
	float m_ShearZY;

	float m_TranslateX;
	float m_TranslateY;
	float m_TranslateZ;

	glm::mat4 ToMatrix() {
		auto scale = glm::scale(glm::mat4(1), {m_ScaleX, m_ScaleY, m_ScaleZ});
		auto rotateX = glm::rotate(scale, m_RotationX, {1, 0, 0});
		auto rotateY = glm::rotate(rotateX, m_RotationY, {0, 1, 0});
		auto rotateZ = glm::rotate(rotateY, m_RotationZ, {0, 0, 1});
		auto shear = glm::shear(rotateZ, {0, 0, 0}, {m_ShearXY, m_ShearXZ}, {m_ShearYX, m_ShearYZ}, {m_ShearZX, m_ShearZY});
		auto translate = glm::translate(rotateZ, {m_TranslateX, m_TranslateY, m_TranslateZ});

		return translate;
	}
};

static void GenNewFractal() {
	// scale, rotation, shear, translation

	std::vector<glm::mat4> transformations(TRANSFORMATION_COUNT);

	for (std::size_t i = 0; i < transformations.size(); i++) {
		Transform transform;
		transform.m_ScaleX = s_Distrib(s_Generator) * 0.4 + 0.4;
		transform.m_ScaleY = s_Distrib(s_Generator) * 0.4 + 0.4;
		transform.m_ScaleZ = s_Distrib(s_Generator) * 0.4 + 0.4;

		transform.m_RotationX = s_Distrib(s_Generator) * 2 * 3.14;
		transform.m_RotationY = s_Distrib(s_Generator) * 2 * 3.14;
		transform.m_RotationZ = s_Distrib(s_Generator) * 2 * 3.14;

		transform.m_ShearXY = s_Distrib(s_Generator) * 0.2f - 0.1f;
		transform.m_ShearXZ = s_Distrib(s_Generator) * 0.2f - 0.1f;
		transform.m_ShearYX = s_Distrib(s_Generator) * 0.2f - 0.1f;
		transform.m_ShearYZ = s_Distrib(s_Generator) * 0.2f - 0.1f;
		transform.m_ShearZX = s_Distrib(s_Generator) * 0.2f - 0.1f;
		transform.m_ShearZY = s_Distrib(s_Generator) * 0.2f - 0.1f;

		transform.m_TranslateX = s_Distrib(s_Generator) * 1.2f - 0.6f;
		transform.m_TranslateY = s_Distrib(s_Generator) * 1.2f - 0.6f;
		transform.m_TranslateZ = s_Distrib(s_Generator) * 1.2f - 0.6f;

		transformations[i] = transform.ToMatrix();
	}

	ApplyTransforms(transformations);
}

static void KeyCallback(GLFWwindow* window, int key, int scancode, int action, int mods) {
	if (action != GLFW_PRESS)
		return;

	if (key == GLFW_KEY_ESCAPE && action == GLFW_PRESS)
		glfwSetWindowShouldClose(window, GLFW_TRUE);

	if (key == GLFW_KEY_R) {
		s_ComputeShader = ReloadComputeShader(s_ComputeShader, s_ComputeShaderPath);
		s_GraphicsShader = ReloadGraphicsShader(s_GraphicsShader, s_VertexShaderPath, s_FragmentShaderPath);
		glBufferData(GL_SHADER_STORAGE_BUFFER, sizeof(float) * PARTICLE_COUNT * 3, nullptr, GL_DYNAMIC_COPY);
		ApplyTransforms(SIERPINSKI_TRIANGLE);
	}

	if (key == GLFW_KEY_T) {
		GenNewFractal();
	}
}

static GLuint CreateDummyVAO() {
	GLuint vertexArray;
	glCreateVertexArrays(1, &vertexArray);

	GLuint vertexBuffer;
	glCreateBuffers(1, &vertexBuffer);

	// Buffer with just one point
	float vertices[] = {0.0f, 0.0f, 0.0f};
	glNamedBufferData(vertexBuffer, sizeof(vertices), vertices, GL_STATIC_DRAW);
	glVertexArrayVertexBuffer(vertexArray, 0, vertexBuffer, 0, sizeof(float) * 3);
	glEnableVertexArrayAttrib(vertexArray, 0);
	glVertexArrayAttribFormat(vertexArray, 0, 3, GL_FLOAT, GL_FALSE, 0);
	glVertexArrayAttribBinding(vertexArray, 0, 0);

	return vertexArray;
}

static void CreateGpuBuffer() {
	GLuint ssbo;
	glGenBuffers(1, &ssbo);
	glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, ssbo);
	glBufferData(GL_SHADER_STORAGE_BUFFER, sizeof(float) * 3 * PARTICLE_COUNT, nullptr,
		GL_DYNAMIC_COPY);  // sizeof(data) only works for statically sized C/C++ arrays.
}

static GLFWwindow* InitWindow() {
	glfwSetErrorCallback(ErrorCallback);

	if (!glfwInit())
		exit(EXIT_FAILURE);

	glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
	glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 6);

	int width = 1280;
	int height = 720;

	GLFWwindow* window = glfwCreateWindow(width, height, "", NULL, NULL);
	if (!window) {
		glfwTerminate();
		exit(EXIT_FAILURE);
	}

	glfwSetKeyCallback(window, KeyCallback);

	glfwMakeContextCurrent(window);
	gladLoadGL();
	glfwSwapInterval(SWAP_INTERVAL);

	s_ComputeShader = CreateComputeShader(s_ComputeShaderPath);
	if (s_ComputeShader == -1) {
		std::cerr << "Compute shader failed\n";
		return nullptr;
	}

	s_GraphicsShader = CreateGraphicsShader(s_VertexShaderPath, s_FragmentShaderPath);
	if (s_GraphicsShader == -1) {
		std::cerr << "Graphics shader failed\n";
		return nullptr;
	}

	auto viewMatrix = glm::lookAt(glm::vec3{1, 2, 2}, {0, 0, 0}, {0, 1, 0});
	auto projectionMatrix = glm::perspective(70.0f, 2.0f, 0.01f, 10.0f);

	glUseProgram(s_GraphicsShader);
	glUniformMatrix4fv(0, 1, false, glm::value_ptr(viewMatrix));
	glUniformMatrix4fv(1, 1, false, glm::value_ptr(projectionMatrix));

	return window;
}

int main() {
	GLFWwindow* window = InitWindow();
	if (!window)
		return -1;

	auto vertexArray = CreateDummyVAO();

	CreateGpuBuffer();

	float lastTime = (float)glfwGetTime();

	int fps = 0;
	float secondsTimer = 0.0f;

	glBindVertexArray(vertexArray);

	// ApplyTransforms(SIERPINSKI_TRIANGLE);
	GenNewFractal();

	glClearColor(0.4f, 0.4f, 0.4f, 1.0f);

	while (!glfwWindowShouldClose(window)) {
		// ScopedTimer timer("Main Loop");

		float currentTime = (float)glfwGetTime();
		float dt = currentTime - lastTime;
		lastTime = currentTime;

		secondsTimer += dt;
		if (secondsTimer >= 1.0f) {
			std::string title = "FPS : " + std::to_string(fps);
			glfwSetWindowTitle(window, title.c_str());

			secondsTimer = 0.0f;
			fps = 0;

			GenNewFractal();
		}

		// Compute
		glUseProgram(s_ComputeShader);
		glDispatchCompute(PARTICLE_COUNT / WORK_GROUP_SIZE, 1, 1);

		// Ensure all writes to the image are complete
		glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

		// Graphics
		glUseProgram(s_GraphicsShader);

		glDrawArraysInstanced(GL_POINTS, 0, 1, PARTICLE_COUNT);

		glfwSwapBuffers(window);
		glfwPollEvents();

		glClear(GL_COLOR_BUFFER_BIT);

		int width, height;
		glfwGetWindowSize(window, &width, &height);
		glViewport(0, 0, width, height);

		// float positions[3 * PARTICLE_COUNT];
		// glGetBufferSubData(GL_SHADER_STORAGE_BUFFER, 0, sizeof(float) * 3 * PARTICLE_COUNT, positions);

		// std::cout << "Positions :\n";
		// for (size_t i = 0; i < PARTICLE_COUNT; i++)
		// {
		// 	std::cout << "\t" << positions[i * 3] << " " << positions[i * 3 + 1] << " " << positions[i * 3 + 2] << "\n";
		// }

		fps++;
	}

	glfwDestroyWindow(window);

	glfwTerminate();
}