#include <fstream>
#include <iostream>
#include <sstream>
#include <stdio.h>
#include <stdlib.h>

// glad and GLFW
#include <GLFW/glfw3.h>
#include <glad/glad.h>

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

// HPP files
#include "shader.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

#define VERTEX_FILE "./shaders/vertex.glsl"
#define FRAGMENT_FILE "./shaders/fragment.glsl"
#define COMPUTE_FILE "./shaders/compute.glsl"

const float SCR_WIDTH = 1200.0f;
const float SCR_HEIGHT = 1200.0f;

void renderProceduralVideo(GLuint shaderProgram, GLuint VAO, GLuint resLoc, GLuint timeLoc, int width, int height, float duration, int fps, const std::string &folder) {
	int totalFrames = static_cast<int>(duration * fps);
	for (int i = 0; i < totalFrames; ++i) {
		float t = i / float(fps); // time in seconds
		glUniform1f(timeLoc, t);
		glUniform2f(resLoc, width, height);

		// Setup off-screen FBO
		GLuint fbo, texture;
		glGenFramebuffers(1, &fbo);
		glBindFramebuffer(GL_FRAMEBUFFER, fbo);

		glGenTextures(1, &texture);
		glBindTexture(GL_TEXTURE_2D, texture);
		glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, GL_RGB, GL_UNSIGNED_BYTE, nullptr);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);

		glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, texture, 0);

		if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE) {
			std::cerr << "Framebuffer not complete!\n";
			return;
		}

		glViewport(0, 0, width, height);

		glClearColor(0, 0, 0, 1);
		glClear(GL_COLOR_BUFFER_BIT);

		glUseProgram(shaderProgram);
		glBindVertexArray(VAO);
		glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, 0);

		std::vector<unsigned char> pixels(width * height * 3);
		glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE, pixels.data());

		// Flip vertically
		for (int y = 0; y < height / 2; ++y) {
			int opposite = height - y - 1;
			for (int x = 0; x < width * 3; ++x)
				std::swap(pixels[y * width * 3 + x], pixels[opposite * width * 3 + x]);
		}

		char buffer[256];
		sprintf(buffer, "%s/frame_%04d.png", folder.c_str(), i);
		stbi_write_png(buffer, width, height, 3, pixels.data(), width * 3);

		glBindFramebuffer(GL_FRAMEBUFFER, 0);
		glDeleteFramebuffers(1, &fbo);
		glDeleteTextures(1, &texture);

		std::cout << "\rRendering frame " << i + 1 << "/" << totalFrames << std::flush;
	}
	std::cout << "\nDone rendering video frames!\n";
}

// Call this function with your shader program, VAO, and desired size
void renderProceduralToPNG(GLuint shaderProgram, GLuint VAO, int width, int height, const char *filename) {
	// Setup off-screen framebuffer
	GLuint fbo, texture;
	glGenFramebuffers(1, &fbo);
	glBindFramebuffer(GL_FRAMEBUFFER, fbo);

	glGenTextures(1, &texture);
	glBindTexture(GL_TEXTURE_2D, texture);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, GL_RGB, GL_UNSIGNED_BYTE, nullptr);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);

	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, texture, 0);

	if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE) {
		std::cerr << "Framebuffer not complete!\n";
		return;
	}

	// Save old viewport
	GLint oldViewport[4];
	glGetIntegerv(GL_VIEWPORT, oldViewport);

	glViewport(0, 0, width, height);

	// Clear and render
	glClearColor(0, 0, 0, 1);
	glClear(GL_COLOR_BUFFER_BIT);

	glUseProgram(shaderProgram);
	glBindVertexArray(VAO);
	glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, 0);

	// Read pixels
	std::vector<unsigned char> pixels(width * height * 3);
	glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE, pixels.data());

	// Flip vertically
	for (int y = 0; y < height / 2; ++y) {
		int opposite = height - y - 1;
		for (int x = 0; x < width * 3; ++x) {
			std::swap(pixels[y * width * 3 + x], pixels[opposite * width * 3 + x]);
		}
	}

	stbi_write_png(filename, width, height, 3, pixels.data(), width * 3);

	// Cleanup
	glBindFramebuffer(GL_FRAMEBUFFER, 0);
	glDeleteFramebuffers(1, &fbo);
	glDeleteTextures(1, &texture);

	// Restore old viewport
	glViewport(oldViewport[0], oldViewport[1], oldViewport[2], oldViewport[3]);
}

int main(void) {
	// Initialize GLFW
	if (!glfwInit())
		return -1;

	glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
	glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 1);
	glfwWindowHint(GLFW_OPENGL_FORWARD_COMPAT, GL_TRUE); // REQUIRED on macOS
	glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
	glfwWindowHint(GLFW_SAMPLES, 8); // 8x MSAA

	GLFWwindow *window = glfwCreateWindow((int)SCR_WIDTH, (int)SCR_HEIGHT, "RASTER", nullptr, nullptr);
	if (!window) {
		glfwTerminate();
		return -1;
	}
	glfwMakeContextCurrent(window);

	if (!gladLoadGLLoader((GLADloadproc)glfwGetProcAddress)) {
		std::cerr << "Failed to initialize GLAD\n";
		return -1;
	}

	glViewport(0, 0, (int)SCR_WIDTH, (int)SCR_HEIGHT);

	// DEBUG
	std::cout << "OpenGL Version: " << glGetString(GL_VERSION) << std::endl;

	// ERROR: macOS does not support compute shaders in OpenGL < 4.3.
	// GLuint computeProgram = createComputeProgram(COMPUTE_FILE);

	// ----- Vertex Data -----
	// Cube
	// Positions
	float vertices[] = {// positions
	                    -1.0f, -1.0f, 0.0f, 1.0f, -1.0f, 0.0f, 1.0f, 1.0f, 0.0f, -1.0f, 1.0f, 0.0f};

	unsigned int indices[] = {0, 1, 2, 2, 3, 0};

	GLuint VAO, VBO, EBO;
	glGenVertexArrays(1, &VAO);
	glGenBuffers(1, &VBO);
	glGenBuffers(1, &EBO);

	glBindVertexArray(VAO);

	glBindBuffer(GL_ARRAY_BUFFER, VBO);
	glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);

	glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO);
	glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(indices), indices, GL_STATIC_DRAW);

	// position - VERTEX SHADER
	glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(float), (void *)0);
	glEnableVertexAttribArray(0);

	// color - FRAGMENT SHADER
	glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(float), (void *)(3 * sizeof(float)));
	glEnableVertexAttribArray(1);

	glBindBuffer(GL_ARRAY_BUFFER, 0);
	// glBindVertexArray(0);

	// Load shaders
	GLuint shaderProgram = createShaderProgram(VERTEX_FILE, FRAGMENT_FILE);

	// Check if shader program was created successfully
	if (shaderProgram == 0) {
		std::cerr << "Failed to create shader programs.\n";
		return -1;
	}

	// Caulculate model, view, projection matrices
	glm::mat4 model = glm::mat4(1.0f);
	glm::mat4 view = glm::translate(glm::mat4(1.0f), glm::vec3(0.0f, 0.0f, -1.0f));
	glm::mat4 projection = glm::perspective(glm::radians(45.0f), (SCR_WIDTH / SCR_HEIGHT) * 2, 0.1f, 100.0f);

	GLuint modelLoc = (unsigned int)glGetUniformLocation(shaderProgram, "model");
	GLuint viewLoc = (unsigned int)glGetUniformLocation(shaderProgram, "view");
	GLuint projLoc = (unsigned int)glGetUniformLocation(shaderProgram, "projection");

	glUseProgram(shaderProgram);
	glUniformMatrix4fv((int)modelLoc, 1, GL_FALSE, glm::value_ptr(model));
	glUniformMatrix4fv((int)viewLoc, 1, GL_FALSE, glm::value_ptr(view));
	glUniformMatrix4fv((int)projLoc, 1, GL_FALSE, glm::value_ptr(projection));

	GLint resLoc = glGetUniformLocation(shaderProgram, "u_resolution");
	glUniform2f(resLoc, 3840, 2160); // hardcode for now

	GLint timeLoc = glGetUniformLocation(shaderProgram, "u_time");

	glEnable(GL_MULTISAMPLE);

	float scale = 1.0f;
	int frameCount = 0;

	// Render loop
	float timeValue = (float)glfwGetTime();

	// --- FPS Counter ---
	double currentTime = glfwGetTime();

	glUniform1f(timeLoc, timeValue);
	frameCount++;

	std::string outputFolder = "./frames";
	int width = 3840;
	int height = 2160;
	float duration = 120.0f; // seconds
	int fps = 30;

	renderProceduralVideo(shaderProgram, VAO, resLoc, timeLoc, width, height, duration, fps, outputFolder);

	glDeleteVertexArrays(1, &VAO);
	glDeleteBuffers(1, &VBO);
	glDeleteBuffers(1, &EBO);
	glDeleteProgram(shaderProgram);

	glfwDestroyWindow(window);
	glfwTerminate();
	return 0;
}