#include "GL/glew.h"
#include "GLFW/glfw3.h"

#include "pwyazh.h"
#include "pwyazh_GL.h"

#include "xoshiro128plus.c"

#include "common.h"

#include "assert.h"

int main(void)
{
	// glfw init
	glfwSetErrorCallback(error_callback);
	if (glfwInit() == GLFW_FALSE)
		return 1;
	glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
	glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
	glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
	glfwWindowHint(GLFW_RESIZABLE, GLFW_FALSE);
	S32 width = 1600;
	S32 height = 900;
	GLFWwindow *window = glfwCreateWindow(width, height, "ssao", 0, 0);
	if (!window)
		return 1;
	glfwSetInputMode(window, GLFW_CURSOR, GLFW_CURSOR_DISABLED);
	glfwMakeContextCurrent(window);
	if (glewInit() != GLEW_OK)
		return 1;
	glEnable(GL_DEPTH_TEST);

	// program init
	Arena *arena = arena_alloc(Kilobytes(256));
	Input input = input_init();
	Camera camera = {
		.pos = v3f(0.0f, 2.0f, 5.0f),
		.fovx = 90.0f,
		.near = 0.1f,
		.far = 100.0f,
		.yaw = 0.0f,
		.pitch = 0.0f
	};
	light_t light = {{2.0f, 2.0f, 2.0f}, {10.0f, 10.0f, 15.0f}};

	// meshes
	Mesh *quad = mesh_gen_quad(arena);
	Mesh *cube = mesh_load_obj(arena, "../../data/models/cube.obj");

	// shaders
	U32 shader_gbuffer = load_shader("gbuffer.vert", "gbuffer.frag");
	U32 shader_ssao = load_shader("ssao.vert", "ssao.frag");
	U32 shader_ssao_blur = load_shader("ssao.vert", "blur.frag");
	U32 shader_deferred = load_shader("deferred.vert", "deferred.frag");
	U32 shader_screen = load_shader("screen.vert", "screen.frag");
	U32 shader_occlusion = load_shader("screen.vert", "occlusion.frag");

	glUseProgram(shader_ssao);
	shader_set_1i(shader_ssao, "positions", 0);
	shader_set_1i(shader_ssao, "normals", 1);
	shader_set_1i(shader_ssao, "noise", 2);
	glUseProgram(shader_deferred);
	shader_set_1i(shader_deferred, "positions", 0);
	shader_set_1i(shader_deferred, "normals", 1);
	shader_set_1i(shader_deferred, "colors", 2);
	shader_set_1i(shader_deferred, "ssao", 3);
	glUseProgram(shader_ssao_blur);
	shader_set_1i(shader_ssao_blur, "ssao", 0);
	glUseProgram(0);

	// framebuffers
	U32 gbuffer;
	glGenFramebuffers(1, &gbuffer);
	glBindFramebuffer(GL_FRAMEBUFFER, gbuffer);

	U32 gpositions, gnormals, gcolors, rbo; 

	glGenTextures(1, &gpositions);
	glBindTexture(GL_TEXTURE_2D, gpositions);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA16F, width, height, 0, GL_RGBA, GL_FLOAT, 0);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, gpositions, 0);

	glGenTextures(1, &gnormals);
	glBindTexture(GL_TEXTURE_2D, gnormals);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA16F, width, height, 0, GL_RGBA, GL_FLOAT, 0);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT1, GL_TEXTURE_2D, gnormals, 0);

	glGenTextures(1, &gcolors);
	glBindTexture(GL_TEXTURE_2D, gcolors);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, width, height, 0, GL_RGBA, GL_UNSIGNED_BYTE, 0);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT2, GL_TEXTURE_2D, gcolors, 0);

	U32 attachments[3] = {GL_COLOR_ATTACHMENT0, GL_COLOR_ATTACHMENT1, GL_COLOR_ATTACHMENT2};
	glDrawBuffers(3, attachments);

	glGenRenderbuffers(1, &rbo);
	glBindRenderbuffer(GL_RENDERBUFFER, rbo);
	glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH_COMPONENT, width, height);
	glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_RENDERBUFFER, rbo);

	if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
		fprintf(stderr, "error: gbuffer not complete\n");
	fprintf(stdout, "info: gbuffer complete\n");

	S32 show_gbuffer = 0;

	S32 nsamples = 64;
	V3F *samples = malloc(nsamples * sizeof(V3F));
	S32 file = open("/dev/urandom", O_RDONLY);
	if (file < 0) {
		perror("error:");
		return 1;
	}
	U32 seed[4];
	if (read(file, seed, sizeof(seed)) < 0) {
		perror("error");
		return 1;
	}
	close(file);
	for (S32 i = 0; i < nsamples; i++) {
		V3F sample = {
			nextf(seed) * 2.0f - 1.0f,
			nextf(seed) * 2.0f - 1.0f,
			nextf(seed)
		};
		sample = v3f_norm(sample);
		float scale = (F32)i / 64.0f;
		scale = lerp(0.1f, 1.0f, scale * scale);
		sample = v3f_scalef(sample, scale);
		samples[i] = sample;
		assert(-1.0f < sample.x && sample.x < 1.0f);
		assert(-1.0f < sample.y && sample.y < 1.0f);
		assert(0.0f < sample.z && sample.z < 1.0f);
		// printf("color:\t%f\t%f\t%f\n", sample.x, sample.y, sample.z);
	}
	V3F noise[16];
	for (S32 i = 0; i < 16; i++) {
		noise[i] = (V3F){
			nextf(seed) * 2.0f - 1.0f,
			nextf(seed) * 2.0f - 1.0f,
			0.0f,
		};
	}

	U32 noise_texture;
	glGenTextures(1, &noise_texture);
	glBindTexture(GL_TEXTURE_2D, noise_texture);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA16F, 4, 4, 0, GL_RGB, GL_FLOAT, noise);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
	glBindTexture(GL_TEXTURE_2D, 0);

	U32 ssao_fbo;
	glGenFramebuffers(1, &ssao_fbo);
	glBindFramebuffer(GL_FRAMEBUFFER, ssao_fbo);
	U32 ssao_colorbuffer;
	glGenTextures(1, &ssao_colorbuffer);
	glBindTexture(GL_TEXTURE_2D, ssao_colorbuffer);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, width, height, 0, GL_RED, GL_FLOAT, 0);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, ssao_colorbuffer, 0);
	glBindTexture(GL_TEXTURE_2D, 0);
	glBindFramebuffer(GL_FRAMEBUFFER, 0);

	U32 ssao_blur_fbo, ssao_blur_colorbuffer;
	glGenFramebuffers(1, &ssao_blur_fbo);
	glBindFramebuffer(GL_FRAMEBUFFER, ssao_blur_fbo);
	glGenTextures(1, &ssao_blur_colorbuffer);
	glBindTexture(GL_TEXTURE_2D, ssao_blur_colorbuffer);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, width, height, 0, GL_RED, GL_FLOAT, 0);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, ssao_blur_colorbuffer, 0);
	glBindTexture(GL_TEXTURE_2D, 0);
	glBindFramebuffer(GL_FRAMEBUFFER, 0);

	F64 last_time = glfwGetTime();
	while (!glfwWindowShouldClose(window)) {
		F64 time = glfwGetTime();
		F32 dt = time - last_time;
		last_time = time;

		input_update_last_state(&input);

		glfwPollEvents();
		glfwGetFramebufferSize(window, &width, &height);
		process_glfw_keyboard(window, &input);
		process_glfw_mouse_pos(window, &input);

		if (key_first_press(input.exit))
			glfwSetWindowShouldClose(window, GLFW_TRUE);

		F32 speed = 2.0f;
		V3F dv = get_dv_camera_first_person(&input, &camera, speed, dt);
		V3F camera_dp = v3f_add(camera_dp, dv);
		camera_dp = v3f_scalef(camera_dp, 0.8f);
		camera.pos = v3f_add(camera.pos, camera_dp);

		F32 sensitivity = 0.1f;
		input.mouse_offset = v2f_scalef(input.mouse_offset, sensitivity);
		camera.yaw += input.mouse_offset.x;
		camera.pitch += input.mouse_offset.y;
		if (camera.pitch > 89.0f)
			camera.pitch = 89.0f;
		if (camera.pitch < -89.0f)
			camera.pitch = -89.0f;

		// update
		if (key_first_press(input.jump))
			show_gbuffer = !show_gbuffer;

		F32 radius = 3.0f;
		light.position = v3f(cos(time) * radius, light.position.y, -sin(time) * radius);

		// render
		F32 ar = (F32)width / (F32)height;
		MAT4 projection = camera_persp(camera, ar);
		MAT4 view = get_view_matrix(&camera);

		glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

		// render stuff into gbuffer
		glBindFramebuffer(GL_FRAMEBUFFER, gbuffer);
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

		glUseProgram(shader_gbuffer);
		shader_set_mat4fv(shader_gbuffer, "projection", projection);
		shader_set_mat4fv(shader_gbuffer, "view", view);
		shader_set_mat4fv(shader_gbuffer, "model", mat4_identity());
		mesh_draw(cube);
		MAT4 model = mat4_make_scale(v3f(10.0f, 0.2f, 10.0f));
		model = mat4_translate(model, v3f(0.0f, -1.25f, 0.0f));
		shader_set_mat4fv(shader_gbuffer, "model", model);
		mesh_draw(cube);
		glUseProgram(0);

		glBindFramebuffer(GL_FRAMEBUFFER, 0);

		// render ssao texture
		glBindFramebuffer(GL_FRAMEBUFFER, ssao_fbo);
		glClear(GL_COLOR_BUFFER_BIT);
		glUseProgram(shader_ssao);

		glActiveTexture(GL_TEXTURE0);
		glBindTexture(GL_TEXTURE_2D, gpositions);
		glActiveTexture(GL_TEXTURE1);
		glBindTexture(GL_TEXTURE_2D, gnormals);
		glActiveTexture(GL_TEXTURE2);
		glBindTexture(GL_TEXTURE_2D, noise_texture);

		for (S32 i = 0; i < nsamples; i++) {
			char str[512] = {0};
			snprintf(str, 512, "samples[%d]", i);
			shader_set_3fv(shader_ssao, str, samples[i]);
		}
		shader_set_mat4fv(shader_ssao, "projection", projection);
		mesh_draw(quad);

		glUseProgram(0);
		glBindFramebuffer(GL_FRAMEBUFFER, 0);

		glBindFramebuffer(GL_FRAMEBUFFER, ssao_blur_fbo);
		glClear(GL_COLOR_BUFFER_BIT);
		glUseProgram(shader_ssao_blur);
		glActiveTexture(GL_TEXTURE0);
		glBindTexture(GL_TEXTURE_2D, ssao_colorbuffer);
		mesh_draw(quad);
		glUseProgram(0);
		glBindFramebuffer(GL_FRAMEBUFFER, 0);

		if (show_gbuffer) {
			MAT4 scale = mat4_make_scale(v3f(0.5f, 0.5f, 0.5f));
			MAT4 leftup = mat4_translate(scale, v3f(-0.5f, 0.5f, 0.0f));
			MAT4 rightup = mat4_translate(scale, v3f(0.5f, 0.5f, 0.0f));
			MAT4 leftdown = mat4_translate(scale, v3f(-0.5f, -0.5f, 0.0f));
			MAT4 rightdown = mat4_translate(scale, v3f(0.5f, -0.5f, 0.0f));

			glDisable(GL_DEPTH_TEST);

			glUseProgram(shader_screen);
			shader_set_mat4fv(shader_screen, "model", leftup);
			glActiveTexture(GL_TEXTURE0);
			glBindTexture(GL_TEXTURE_2D, gpositions);
			mesh_draw(quad);

			shader_set_mat4fv(shader_screen, "model", rightup);
			glBindTexture(GL_TEXTURE_2D, gnormals);
			mesh_draw(quad);

			shader_set_mat4fv(shader_screen, "model", leftdown);
			glBindTexture(GL_TEXTURE_2D, gcolors);
			mesh_draw(quad);
			glUseProgram(0);

			glUseProgram(shader_occlusion);
			shader_set_mat4fv(shader_occlusion, "model", rightdown);
			glBindTexture(GL_TEXTURE_2D, ssao_blur_colorbuffer);
			mesh_draw(quad);
			glUseProgram(0);

			glEnable(GL_DEPTH_TEST);
		} else {
			// render scene lighting
			glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

			glUseProgram(shader_deferred);

			glActiveTexture(GL_TEXTURE0);
			glBindTexture(GL_TEXTURE_2D, gpositions);

			glActiveTexture(GL_TEXTURE1);
			glBindTexture(GL_TEXTURE_2D, gnormals);

			glActiveTexture(GL_TEXTURE2);
			glBindTexture(GL_TEXTURE_2D, gcolors);

			glActiveTexture(GL_TEXTURE3);
			glBindTexture(GL_TEXTURE_2D, ssao_blur_colorbuffer);

			V4F light_pos_v4 = v4f(light.position.x, light.position.y, light.position.z, 1.0f);
			V4F light_view_pos = mat4_v4f_mul(view, light_pos_v4);
			shader_set_3fv(shader_deferred, "light.position", v3f_from_v4f(light_view_pos));
			shader_set_3fv(shader_deferred, "light.color", light.color);
			shader_set_mat4fv(shader_deferred, "model", mat4_identity());

			mesh_draw(quad);

			glUseProgram(0);
		}

		glfwSwapBuffers(window);
	}

	return 0;
}