#include <iostream>
#include <stdio.h>
#include <string>
#include <vector>

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

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

// Custom headers
#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"

const int WIDTH = 3840;
const int HEIGHT = 2160;
const float DURATION = 120.0f; // 2 minutes
const int FPS = 30;
const std::string OUTPUT_FOLDER = "./frames";

// ------------------ Offscreen Frame Renderer ------------------
void renderFrame(GLuint shaderProgram, GLuint VAO, GLuint resLoc, GLuint timeLoc, GLuint scaleLoc, GLuint mouseLoc, float time, const std::string &filename) {
	static GLuint fbo = 0;
	static GLuint texture = 0;

	glUniform1f(scaleLoc, 1.0f); // or whatever scale you want
	glUniform2f(mouseLoc, 0.0f, 0.0f);

	// First frame setup
	if (fbo == 0) {
		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";
			exit(-1);
		}
	}

	// Bind FBO
	glBindFramebuffer(GL_FRAMEBUFFER, fbo);
	glViewport(0, 0, WIDTH, HEIGHT);
	glClearColor(0, 0, 0, 1);
	glClear(GL_COLOR_BUFFER_BIT);

	// Render procedural quad
	glUseProgram(shaderProgram);
	glUniform2f(resLoc, float(WIDTH), float(HEIGHT));
	glUniform1f(timeLoc, time);
	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 opp = HEIGHT - y - 1;
		for (int x = 0; x < WIDTH * 3; ++x)
			std::swap(pixels[y * WIDTH * 3 + x], pixels[opp * WIDTH * 3 + x]);
	}

	stbi_write_png(filename.c_str(), WIDTH, HEIGHT, 3, pixels.data(), WIDTH * 3);

	glBindFramebuffer(GL_FRAMEBUFFER, 0);
}

// ------------------ Main ------------------
int main() {
	// Initialize GLFW
	if (!glfwInit()) {
		std::cerr << "Failed to initialize GLFW\n";
		return -1;
	}

	glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
	glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 1);
	glfwWindowHint(GLFW_OPENGL_FORWARD_COMPAT, GL_TRUE);
	glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
	glfwWindowHint(GLFW_VISIBLE, GLFW_FALSE); // hidden

	GLFWwindow *window = glfwCreateWindow(1, 1, "", nullptr, nullptr);
	glfwMakeContextCurrent(window);

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

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

	// ------------------ VAO + VBO ------------------
	float vertices[] = {-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);

	glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(float), (void *)0);
	glEnableVertexAttribArray(0);

	// ------------------ Shader ------------------
	GLuint shaderProgram = createShaderProgram(VERTEX_FILE, FRAGMENT_FILE);
	if (shaderProgram == 0) {
		std::cerr << "Failed to create shader.\n";
		return -1;
	}

	GLuint resLoc = glGetUniformLocation(shaderProgram, "u_resolution");
	GLuint timeLoc = glGetUniformLocation(shaderProgram, "u_time");

	GLuint scaleLoc = glGetUniformLocation(shaderProgram, "u_scale");
	GLuint mouseLoc = glGetUniformLocation(shaderProgram, "u_mouse");

	// In renderFrame(), add:

	// ------------------ Render Loop ------------------
	int totalFrames = int(DURATION * FPS);
	for (int i = 0; i < totalFrames; ++i) {
		float t = i / float(FPS);
		char filename[256];
		sprintf(filename, "%s/frame_%04d.png", OUTPUT_FOLDER.c_str(), i);

		renderFrame(shaderProgram, VAO, resLoc, timeLoc, scaleLoc, mouseLoc, t, filename);
		std::cout << "\rRendered frame " << i + 1 << "/" << totalFrames << std::flush;
	}

	std::cout << "\nDone! Stitch frames into video with ffmpeg:\n";
	std::cout << "ffmpeg -framerate 30 -i ./frames/frame_%04d.png -c:v libx264 -pix_fmt yuv420p noise_video.mp4\n";

	// ------------------ Cleanup ------------------
	glDeleteVertexArrays(1, &VAO);
	glDeleteBuffers(1, &VBO);
	glDeleteBuffers(1, &EBO);
	glDeleteProgram(shaderProgram);
	glfwDestroyWindow(window);
	glfwTerminate();

	return 0;
}