#include "render.h"
#include "state.h"
#include "../platform.h"
#include "../lib/math.h"
#include "../debug/logger.h"
#include "gl_helpers.h"
#include "string.h"
#include "../camera.h" // @Cleanup: remove dependency from camera

r_state r_render_state;


void r_init()
{
	// OpenGL debug messages
	glEnable(GL_DEBUG_OUTPUT);
	glEnable(GL_DEBUG_OUTPUT_SYNCHRONOUS);
	glDebugMessageCallback(glDebugOutput, nullptr);
	glDebugMessageControl(GL_DONT_CARE, GL_DONT_CARE, GL_DONT_CARE, 0, nullptr, GL_TRUE);

	// Depth test
	glEnable(GL_DEPTH_TEST);
	// Face culling
	glEnable(GL_CULL_FACE);
	// Texture pixel alignment (default is 4, so rgb textures should have every pixel padded to 4 bytes. With this we can have 3 byte pixels with no padding between eachother)
	glPixelStorei(GL_UNPACK_ALIGNMENT, 1);

	// Blending
	glEnable(GL_BLEND);
	glBlendEquation(GL_FUNC_ADD);
	glBlendFuncSeparate(GL_ONE, GL_ONE_MINUS_SRC_ALPHA, GL_ONE, GL_ONE_MINUS_SRC_ALPHA);

	// Cubemap
	glEnable(GL_TEXTURE_CUBE_MAP_SEAMLESS);


	// Load shaders
	if(!r_shader_from_files(&r_render_state.shader_2d, "shaders/2d.vert", "shaders/2d.frag"))
	{
		LOG(LOG_ERROR, "Cannot load 2D shader.");
	}
	if(!r_shader_from_files(&r_render_state.shader_postprocessing, "shaders/postprocessing.vert", "shaders/postprocessing.frag"))
	{
		LOG(LOG_ERROR, "Cannot load postprocessing shader.");
	}
	if(!r_shader_from_files(&r_render_state.shader_pbr, "shaders/pbr.vert", "shaders/pbr.frag"))
	{
		LOG(LOG_ERROR, "Cannot load pbr shader.");
	}
	if(!r_shader_from_files(&r_render_state.shader_shadow_map, "shaders/shadow_map.vert", "shaders/shadow_map.frag"))
	{
		LOG(LOG_ERROR, "Cannot load shadow_map shader.");
	}
	if(!r_shader_from_files(&r_render_state.shader_environment_map, "shaders/environment_map.vert", "shaders/environment_map.frag"))
	{
		LOG(LOG_ERROR, "Cannot load environment_map shader.");
	}

	// Screen size
	u32 width, height;
	p_window_dimensions(&width, &height);
	glViewport(0, 0, width, height);


	/* We have to apply some post-processing effects only to the 3D world
	 * and some effects to both the 3D world and the HUD.
	 *
	 * So we make 2 framebuffers:
	 *     1. target for 3D world
	 *     2. target for HUD
	 * After rendering, we apply the correct post-processing effects and
	 * merge them toghether.
	 */
	// Init framebuffers
	v2s size = {width, height};
	r_render_state.framebuffer_SCREEN = {.gl_id = 0, .size = size, .flags = 0}; // This is a special framebuffer, because it's already there. We don't need to create it
	r_render_state.framebuffer_HUD = r_framebuffer_create(size, 0);
	r_render_state.framebuffer_3D  = r_framebuffer_create(size, R_FRAMEBUFFER_DEPTH);

	// Default framebuffer
	r_framebuffer_select(&r_render_state.framebuffer_SCREEN);

	// Init fullscreen quad
	glGenVertexArrays(1, &r_render_state.gl_screen_quad_VAO);
	glBindVertexArray(r_render_state.gl_screen_quad_VAO);
	glGenBuffers(1, &r_render_state.gl_screen_quad_VBO);
	glBindBuffer(GL_ARRAY_BUFFER, r_render_state.gl_screen_quad_VBO);
	v2 quad_vertices[2*6] =
	{
	//  position       UV coords
		v2{-1,  1},    v2{0, 1},
		v2{-1, -1},    v2{0, 0},
		v2{ 1, -1},    v2{1, 0},
		v2{-1,  1},    v2{0, 1},
		v2{ 1, -1},    v2{1, 0},
		v2{ 1,  1},    v2{1, 1}
	};
	glBufferData(GL_ARRAY_BUFFER, 2*6 * sizeof(v2), quad_vertices, GL_STATIC_DRAW);
	glEnableVertexAttribArray(0);
	glEnableVertexAttribArray(1);
	glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 2*sizeof(v2), (void*)(0));
	glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 2*sizeof(v2), (void*)(sizeof(v2)));

	// Render state
	r_size_update(width, height);
}

void r_deinit()
{
	r_framebuffer_destroy(&r_render_state.framebuffer_HUD);
	r_framebuffer_destroy(&r_render_state.framebuffer_3D);

	glDeleteBuffers(1, &r_render_state.gl_screen_quad_VBO);
	glDeleteVertexArrays(1, &r_render_state.gl_screen_quad_VAO);
}

void r_size_update(u32 width, u32 height)
{
	// Screen size
	r_render_state.width  = width;
	r_render_state.height = height;
	// Update shaders
	glUseProgram(r_render_state.shader_2d.id);
	glUniform1i(r_render_state.shader_2d.width , width);
	glUniform1i(r_render_state.shader_2d.height, height);

	// Update framebuffers size
	v2s size = {width, height};
	r_render_state.framebuffer_SCREEN.size = size;
	r_framebuffer_update_size(&r_render_state.framebuffer_HUD, size);
	r_framebuffer_update_size(&r_render_state.framebuffer_3D , size);
}

void r_time_update(f64 time)
{
	r_render_state.time = time;
}




void r_clear(v4 color)
{
	glClearColor(color.r, color.g, color.b, color.a);
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
}

void r_swap()
{
	p_graphics_swap();
}



r_scene r_scene_create()
{
	r_scene scene;

	scene.view_matrix         = m4_identity;
	scene.view_matrix_inverse = m4_identity;

	scene.projection_matrix         = m4_identity;
	scene.projection_matrix_inverse = m4_identity;

	scene.objects = NULL;
	scene.objects_count = 0;


	// Init lights memory
	memset(&scene.lights, 0, sizeof(r_light_container));
	glGenBuffers(1, &scene.gl_lights);
	glBindBuffer(GL_UNIFORM_BUFFER, scene.gl_lights);
	glBufferData(GL_UNIFORM_BUFFER, sizeof(r_light_container), &scene.lights, GL_STATIC_DRAW);
	glBindBufferBase(GL_UNIFORM_BUFFER, 0, scene.gl_lights);

	return scene;
}

void r_scene_destroy(r_scene *scene)
{
	glDeleteBuffers(1, &scene->gl_lights);
	r_cubemap_destroy(&scene->environment_map);
}

void r_scene_set_view(r_scene *scene, m4 m)
{
	scene->view_matrix         = m;
	scene->view_matrix_inverse = inverse(m);
}

void r_scene_set_projection(r_scene *scene, m4 m)
{
	scene->projection_matrix         = m;
	scene->projection_matrix_inverse = inverse(m);
}

void r_update_lights(r_scene *scene)
{
	glBindBuffer(GL_UNIFORM_BUFFER, scene->gl_lights);
	glBufferData(GL_UNIFORM_BUFFER, sizeof(r_light_container), &scene->lights, GL_STATIC_DRAW);
}



r_object build_basic_cube()
{
	r_object object = r_object_allocate(1);

	object.meshes[0] = (r_mesh*)p_alloc(sizeof(r_mesh));
	*object.meshes[0] = r_mesh_build_cube();

	object.mesh_material[0] = (r_material*)p_alloc(sizeof(r_material));
	memset(object.mesh_material[0], 0, sizeof(r_material));
	object.mesh_material[0]->shader = &r_render_state.shader_pbr;
	object.mesh_material[0]->albedo_factor  = {1,1,1,1};
	object.mesh_material[0]->emissive_factor  = {1,1,1,0};

	return object;
}

void r_render_scene(r_scene *scene)
{
	// Render shadow maps
	scene->shadow_map = r_build_shadow_map_sun(scene, &scene->lights.sun_lights[0]);

	// Render environment map
	glDisable(GL_CULL_FACE);
	glDisable(GL_DEPTH_TEST);
	static r_object cube = build_basic_cube();
	cube.position = extract_column(scene->view_matrix_inverse, 3).xyz;
	glUseProgram(r_render_state.shader_environment_map.id);
	glActiveTexture(GL_TEXTURE0);
	glBindTexture(GL_TEXTURE_CUBE_MAP, scene->environment_map.gl_id);
	glUniform1i(r_render_state.shader_environment_map.has_environment_map, 1);
	glUniform1i(r_render_state.shader_environment_map.environment_map    , 0);
	r_render_object(scene, &cube, &r_render_state.shader_environment_map);
	glEnable(GL_DEPTH_TEST);
	glEnable(GL_CULL_FACE);

	// Render objects
	glUseProgram(r_render_state.shader_pbr.id);

	// Set shadow map
	// @Cleanup: this should be in r_render_objects
	glActiveTexture(GL_TEXTURE5);
	glBindTexture(GL_TEXTURE_2D, scene->shadow_map.gl_depth_texture);
	glUniform1i       (r_render_state.shader_pbr.has_shadow_map, 1);
	glUniform1i       (r_render_state.shader_pbr.shadow_map    , 5);
	glUniformMatrix4fv(r_render_state.shader_pbr.shadow_matrix , 1, GL_TRUE, (const f32 *)(scene->shadow_map.view_matrix.E));

	// Environment map
	glActiveTexture(GL_TEXTURE6);
	glBindTexture(GL_TEXTURE_CUBE_MAP, scene->environment_map.gl_id);
	glUniform1i(r_render_state.shader_pbr.has_environment_map, 1);
	glUniform1i(r_render_state.shader_pbr.environment_map    , 6);

	// Render objects
	for(u32 i = 0; i < scene->objects_count; i++)
	{
		r_render_object(scene, &scene->objects[i]);
	}

	glUniform1i(r_render_state.shader_pbr.has_shadow_map, 0);
	glUniform1i(r_render_state.shader_pbr.has_environment_map, 0);
	r_free_shadow_map(&scene->shadow_map);
}





// Single object rendering functions (batched funcs are below)
void r_render_object(r_scene *scene, r_object *object, r_shader *shader_override)
{
	// @Performance: Minimize draw calls by batching mesh rendering
	for(u32 i = 0; i < object->count; i++)
	{
		r_mesh     *mesh     = object->meshes[i];
		r_material *material = object->mesh_material[i];
		r_shader   *shader   = material->shader;
		if(shader_override)
			shader = shader_override;

		glUseProgram(shader->id);

		glUniform1f(shader->time  , (f32)r_render_state.time  );
		glUniform1f(shader->width , (f32)r_render_state.width );
		glUniform1f(shader->height, (f32)r_render_state.height);

		// Transform matrices
		m4 projview_matrix = scene->projection_matrix * scene->view_matrix;
		m4 projview_matrix_inverse = inverse(projview_matrix);
		glUniformMatrix4fv(shader->view_matrix,         1, GL_TRUE, (const f32 *)(&projview_matrix));
		glUniformMatrix4fv(shader->view_matrix_inverse, 1, GL_TRUE, (const f32 *)(&projview_matrix_inverse));

		m4 model_transform = r_object_mesh_transform_matrix(object, i);
		glUniformMatrix4fv(shader->model_matrix, 1, GL_TRUE, (const f32 *)(&model_transform));

		// Textures
		// Albedo
		glUniform1i(shader->has_albedo_texture, material->albedo_texture != NULL);
		if(material->albedo_texture)
		{
			glActiveTexture(GL_TEXTURE0);
			glBindTexture(GL_TEXTURE_2D, material->albedo_texture->gl_id);
			glUniform1i(shader->albedo_texture, 0);
		}
		glUniform4f(shader->albedo_factor, material->albedo_factor.r, material->albedo_factor.g, material->albedo_factor.b, material->albedo_factor.a);

		// Metallic
		glUniform1i(shader->has_metallic_texture, material->metallic_texture != NULL);
		if(material->metallic_texture)
		{
			glActiveTexture(GL_TEXTURE1);
			glBindTexture(GL_TEXTURE_2D, material->metallic_texture->gl_id);
			glUniform1i(shader->metallic_texture, 1);
		}
		glUniform1f(shader->metallic_factor, material->metallic_factor);

		// Roughness
		glUniform1i(shader->has_roughness_texture, material->roughness_texture != NULL);
		if(material->roughness_texture)
		{
			glActiveTexture(GL_TEXTURE2);
			glBindTexture(GL_TEXTURE_2D, material->roughness_texture->gl_id);
			glUniform1i(shader->roughness_texture, 2);
		}
		glUniform1f(shader->roughness_factor, material->roughness_factor);

		// Normal
		glUniform1i(shader->has_normal_texture, material->normal_texture != NULL);
		if(material->normal_texture)
		{
			glActiveTexture(GL_TEXTURE3);
			glBindTexture(GL_TEXTURE_2D, material->normal_texture->gl_id);
			glUniform1i(shader->normal_texture, 3);
		}

		// Emissive
		glUniform1i(shader->has_emissive_texture, material->emissive_texture != NULL);
		if(material->emissive_texture)
		{
			glActiveTexture(GL_TEXTURE4);
			glBindTexture(GL_TEXTURE_2D, material->emissive_texture->gl_id);
			glUniform1i(shader->emissive_texture, 4);
		}
		glUniform4f(shader->emissive_factor, material->emissive_factor.r, material->emissive_factor.g, material->emissive_factor.b, material->emissive_factor.a);



		// Draw call
		glBindVertexArray(mesh->gl_VAO);
		glDrawElements(GL_TRIANGLES, mesh->indices_count, GL_UNSIGNED_INT, (void*)(0));
	}
}

void r_render_object_wireframe(r_scene *scene, r_object *object)
{
	glPolygonMode(GL_FRONT_AND_BACK, GL_LINE);
	r_render_object(scene, object);
	glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);
}

void r_render_line(r_scene *scene, v3 a, v3 b, v4 a_color, v4 b_color, f32 thickness)
{
	m4 projview_matrix = scene->projection_matrix * scene->view_matrix;
	v4 a_trans = projview_matrix * V4(a, 1);
	v4 b_trans = projview_matrix * V4(b, 1);
	a_trans /= a_trans.w;
	b_trans /= b_trans.w;
	if(a_trans.z > 1 || b_trans.z > 1)
		return;
	v2 a_2d = a_trans.xy;
	v2 b_2d = b_trans.xy;
	a_2d = (v2{.5,.5} + v2{.5,-0.5} * a_2d) * v2{(f32)r_render_state.width, (f32)r_render_state.height};
	b_2d = (v2{.5,.5} + v2{.5,-0.5} * b_2d) * v2{(f32)r_render_state.width, (f32)r_render_state.height};

	r_2d_immediate_segment(a_2d, b_2d, a_color, b_color, thickness); // @Cleanup: This should be independent from the 2D drawing functions, probably.
}




// Batch render functions
void r_render_objects(r_scene *scene, u64 count, r_object **objects)
{
	// @Feature: implement
}




r_shadow_map r_build_shadow_map_sun(r_scene *scene, r_sun_light *sun)
{
	glCullFace(GL_FRONT);

	r_shadow_map sm;
	u32 width = 2*1024;
	u32 height = 2*1024;

	// Init depth texture/framebuffer
	glGenTextures(1, &sm.gl_depth_texture);
	glBindTexture(GL_TEXTURE_2D, sm.gl_depth_texture);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_DEPTH_COMPONENT, width, height, 0, GL_DEPTH_COMPONENT, GL_FLOAT, NULL);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_MODE, GL_COMPARE_REF_TO_TEXTURE);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_COMPARE_FUNC, GL_LEQUAL);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
	glBindTexture(GL_TEXTURE_2D, 0);

	glGenFramebuffers(1, &sm.gl_FBO);
	glBindFramebuffer(GL_FRAMEBUFFER, sm.gl_FBO);
	glFramebufferTexture2D(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_TEXTURE_2D, sm.gl_depth_texture, 0);
	glDrawBuffer(GL_NONE);

	glBindTexture(GL_TEXTURE_2D, 0);

	// Compute light frustum
	m4 projview_matrix = scene->projection_matrix * scene->view_matrix;
	m4 projview_matrix_inverse = inverse(projview_matrix);
	v4 center_point_4 = projview_matrix_inverse * v4{0, 0, 0, 1};
	v3 center_point = center_point_4.xyz / center_point_4.w;
	f32 size = 1;
	v4 bottom = projview_matrix_inverse * v4{ 0,-1, 0, 1};
	v4 top    = projview_matrix_inverse * v4{ 0, 1, 0, 1};
	v4 left   = projview_matrix_inverse * v4{-1, 0, 0, 1};
	v4 right  = projview_matrix_inverse * v4{ 1, 0, 0, 1};
	v4 front  = projview_matrix_inverse * v4{ 0, 0,-1, 1};
	v4 back   = projview_matrix_inverse * v4{ 0, 0, 1, 1};
	size = maximum(size, length(top.xyz/top.w - bottom.xyz/bottom.w));
	size = maximum(size, length(right.xyz/right.w - left.xyz/left.w));
	size = maximum(size, length(back.xyz/back.w - front.xyz/front.w));
	size /= 2; // @Cleanup: hack to have better precision. There are multiple ways to do it right: smallest box around all objects, cascaded shadow maps


	v3 light_position = center_point - size * sun->direction;
	m4 light_camera_matrix = r_view_matrix(light_position, sun->direction, v3{0, 0, 1});
	m4 projection_matrix = r_orthographic_matrix(-size/2, size/2, -size/2, size/2, 0.1, 100.0);

	sm.view_matrix = projection_matrix * light_camera_matrix;

	// Render shadow map
	glViewport(0, 0, width, height);
	glClear(GL_DEPTH_BUFFER_BIT);

	// @Feature @Cleanup: Some objects might not need to be drawn. Add a "hidden" flag. Maybe take the list of objects to render as an argument
	m4 old_view = scene->view_matrix;
	m4 old_proj = scene->projection_matrix;
	r_scene_set_view      (scene, sm.view_matrix);
	r_scene_set_projection(scene, m4_identity);
	for(u32 i = 0; i < scene->objects_count; i++)
	{
		if(scene->objects[i].has_shadow)
		{
			r_render_object(scene, &scene->objects[i], &r_render_state.shader_shadow_map);
		}
	}
	r_scene_set_view      (scene, old_view);
	r_scene_set_projection(scene, old_proj);

	// Restore framebuffer
	r_framebuffer_select(r_render_state.current_framebuffer);

	glCullFace(GL_BACK);
	return sm;
}

void r_free_shadow_map(r_shadow_map *sm)
{
	glDeleteTextures(1, &sm->gl_depth_texture);
	glDeleteFramebuffers(1, &sm->gl_FBO);
}




void r_merge_and_postprocess()
{
	glDisable(GL_DEPTH_TEST);
	glEnable(GL_FRAMEBUFFER_SRGB);

	glUseProgram(r_render_state.shader_postprocessing.id);
	glActiveTexture(GL_TEXTURE0);
	glBindTexture(GL_TEXTURE_2D, r_render_state.framebuffer_3D.color_texture.gl_id);
	glUniform1i(r_render_state.shader_postprocessing.texture[0], 0);

	glActiveTexture(GL_TEXTURE1);
	glBindTexture(GL_TEXTURE_2D, r_render_state.framebuffer_HUD.color_texture.gl_id);
	glUniform1i(r_render_state.shader_postprocessing.texture[1], 1);

	glBindVertexArray(r_render_state.gl_screen_quad_VAO);
	glDrawArrays(GL_TRIANGLES, 0, 6);

	glEnable(GL_DEPTH_TEST);
	glDisable(GL_FRAMEBUFFER_SRGB);
}



void r_framebuffer_select(r_framebuffer *fb)
{

	glBindFramebuffer(GL_FRAMEBUFFER, fb->gl_id);
	glViewport(0, 0, fb->size.x, fb->size.y);

	glUseProgram(r_render_state.shader_2d.id);
	glUniform1i(r_render_state.shader_2d.width , fb->size.x);
	glUniform1i(r_render_state.shader_2d.height, fb->size.y);

	r_render_state.current_framebuffer = fb;
}