#include <iostream>

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

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

#include <chrono>

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

constexpr int SWAP_INTERVAL = 1;

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 void ErrorCallback(int error, const char* description) {
	std::cerr << "Error: " << description << std::endl;
}

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

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;
	}

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

	std::vector<glm::mat4> transformations = {
		{
			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,
		},
	};

	glUseProgram(s_ComputeShader);
	glUniformMatrix4fv(1, transformations.size(), true, glm::value_ptr(transformations[0]));


	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;
		}

		// 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);
		glClearColor(0, 0, 0, 1);

		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();
}