#include <psemek/app/application_base.hpp>
#include <psemek/app/default_application_factory.hpp>

#include <psemek/gfx/program.hpp>
#include <psemek/gfx/mesh.hpp>
#include <psemek/gfx/color.hpp>
#include <psemek/gfx/texture.hpp>
#include <psemek/gfx/renderbuffer.hpp>
#include <psemek/gfx/framebuffer.hpp>
#include <psemek/gfx/error.hpp>

#include <psemek/math/camera.hpp>
#include <psemek/math/math.hpp>
#include <psemek/math/mesh.hpp>
#include <psemek/math/homogeneous.hpp>
#include <psemek/math/gram_schmidt.hpp>

#include <psemek/util/clock.hpp>

#include <fstream>

using namespace psemek;

static char const phong_vertex_source[] =
R"(#version 330

uniform mat4 u_transform;

layout (location = 0) in vec4 in_position;
layout (location = 1) in vec3 in_normal;
layout (location = 2) in vec4 in_color;

out vec3 position;
out vec4 color;
out vec3 normal;

void main()
{
	position = in_position.xyz;
	gl_Position = u_transform * in_position;
	color = in_color;
	normal = in_normal;
}
)";

static char const phong_fragment_source[] =
R"(#version 330

uniform vec4 u_light_position;
uniform vec4 u_camera_position;

uniform vec4 u_material;

in vec3 position;
in vec4 color;
in vec3 normal;

out vec4 out_color;

void main()
{
	vec3 n = normalize(normal);

	vec3 r = u_light_position.xyz - u_light_position.w * position;
	r = normalize(r);

	vec3 v = u_camera_position.xyz - u_camera_position.w * position;
	v = normalize(v);

	vec3 d = 2.0 * dot(r, n) * n - r;

	vec3 light_color = vec3(1.0, 1.0, 1.0);
	vec3 specular_color = vec3(1.0, 1.0, 1.0);

	vec3 col = u_material.x * color.rgb + u_material.y * max(0.0, dot(n, r)) * color.rgb * light_color + u_material.z * (pow(max(0.0, dot(d, v)), u_material.w)) * specular_color;

	out_color = vec4(col, color.a);
}
)";

struct phong_renderer
{
	struct material
	{
		float ambient;
		float diffuse;
		float specular;
		float shininess;
	};

	struct render_state
	{
		struct material material;
		gfx::mesh const * mesh;
	};

	struct light
	{
		math::vector<float, 4> position;
	};

	struct render_options
	{
		gfx::framebuffer const * framebuffer;
		GLenum draw_buffer;
		math::box<int, 2> viewport;
		math::matrix<float, 4, 4> transform;
		math::point<float, 3> camera_position;
		struct light light;
	};

	void push(render_state const & state);

	void render(render_options const & options);

private:
	std::vector<render_state> render_states_;
	gfx::program program_{phong_vertex_source, phong_fragment_source};
};

void phong_renderer::push(render_state const & state)
{
	render_states_.push_back(state);
}

void phong_renderer::render(render_options const & options)
{
	if (!options.framebuffer) return;

	options.framebuffer->bind();
	gl::DrawBuffer(options.draw_buffer);
	gl::Viewport(options.viewport[0].min, options.viewport[1].min, options.viewport[0].length(), options.viewport[1].length());

	program_.bind();
	program_["u_transform"] = options.transform;
	program_["u_light_position"] = options.light.position;
	program_["u_camera_position"] = math::homogeneous(options.camera_position);

	for (auto const & state : render_states_)
	{
		if (!state.mesh) continue;

		program_["u_material"] = math::vector{state.material.ambient, state.material.diffuse, state.material.specular, state.material.shininess};
		state.mesh->draw();
	}

	render_states_.clear();
}

static char const shadow_builder_vertex_source[] =
R"(#version 330

uniform mat4 u_transform;

layout (location = 0) in vec4 in_position;

void main()
{
	gl_Position = u_transform * in_position;
}
)";

static char const shadow_builder_fragment_source[] =
R"(#version 330

layout (location = 0) out float depth;

void main()
{
	depth = gl_FragCoord.z;
}

)";

struct shadow_map_builder
{
	struct render_state
	{
		gfx::mesh const * mesh;
		math::box<float, 3> bbox;
	};

	struct render_options
	{
		math::vector<float, 3> light;
	};

	shadow_map_builder(std::size_t width, std::size_t height);

	void push(render_state const & state);

	void build(render_options const & options);

	math::matrix<float, 4, 4> const & transform() const;

	gfx::texture_2d const & texture() const;

private:
	std::vector<render_state> render_states_;
	gfx::program program_{shadow_builder_vertex_source, shadow_builder_fragment_source};
	gfx::framebuffer framebuffer_;
	gfx::texture_2d depth_texture_;

	math::matrix<float, 4, 4> transform_;
};

shadow_map_builder::shadow_map_builder(std::size_t width, std::size_t height)
{
	depth_texture_.load<gfx::depth24_pixel>({width, height});
	gl::TexParameteri(gl::TEXTURE_2D, gl::TEXTURE_MIN_FILTER, gl::LINEAR);
	gl::TexParameteri(gl::TEXTURE_2D, gl::TEXTURE_MAG_FILTER, gl::LINEAR);
	gl::TexParameteri(gl::TEXTURE_2D, gl::TEXTURE_WRAP_S, gl::CLAMP_TO_EDGE);
	gl::TexParameteri(gl::TEXTURE_2D, gl::TEXTURE_WRAP_T, gl::CLAMP_TO_EDGE);
	gl::TexParameteri(gl::TEXTURE_2D, gl::TEXTURE_COMPARE_MODE, gl::COMPARE_REF_TO_TEXTURE);
	gl::TexParameteri(gl::TEXTURE_2D, gl::TEXTURE_COMPARE_FUNC, gl::LEQUAL);
	framebuffer_.depth(depth_texture_);
	framebuffer_.assert_complete();
}

void shadow_map_builder::push(render_state const & state)
{
	render_states_.push_back(state);
}

void shadow_map_builder::build(render_options const & options)
{
	math::vector<float, 3> light_axes[3];

	light_axes[2] = -options.light;
	light_axes[1] = {0.f, 0.f, 1.f};

	if (light_axes[2][2] > 0.5f)
		light_axes[1] = {0.f, 1.f, 0.f};

	math::gram_schmidt(light_axes[2], light_axes[1]);

	light_axes[0] = math::cross(light_axes[2], light_axes[1]);

	math::box<float, 3> light_bbox;

	math::point<float, 3> origin = math::point<float, 3>::zero();

	for (auto const & state : render_states_)
	{
		for (auto const & v : math::vertices(state.bbox))
		{
			for (std::size_t i = 0; i < 3; ++i)
				light_bbox[i] |= math::dot(light_axes[i], v - origin);
		}
	}

	transform_ = math::orthographic_camera{light_bbox}.projection() * math::homogeneous(math::by_rows(light_axes[0], light_axes[1], light_axes[2]));

	framebuffer_.bind();
	gl::DrawBuffer(gl::NONE);

	gl::Viewport(0, 0, depth_texture_.width(), depth_texture_.height());

	gl::Clear(gl::DEPTH_BUFFER_BIT);

	gl::Enable(gl::DEPTH_TEST);
	gl::DepthFunc(gl::LEQUAL);

	gl::Enable(gl::CULL_FACE);
	gl::CullFace(gl::FRONT);

	program_.bind();
	program_["u_transform"] = transform_;
	gfx::check_error();

	for (auto const & state: render_states_)
	{
		state.mesh->draw();
		gfx::check_error();
	}

	render_states_.clear();

/*
	framebuffer_.null().bind();

	gfx::pixmap_float pixmap({depth_texture_.width(), depth_texture_.height()});
	depth_texture_.pixels(gl::DEPTH_COMPONENT, gl::FLOAT, pixmap.data());

	auto pixels = util::map([](float x){ return std::uint8_t(x * 255); }, pixmap);

	std::ofstream out{"depth.pgm"};
	gfx::write_pgm(pixels, out);
*/
}

math::matrix<float, 4, 4> const & shadow_map_builder::transform() const
{
	return transform_;
}

gfx::texture_2d const & shadow_map_builder::texture() const
{
	return depth_texture_;
}

static char const shadow_vertex_source[] =
R"(#version 330

uniform mat4 u_transform;

layout (location = 0) in vec4 in_position;
layout (location = 1) in vec3 in_normal;
layout (location = 2) in vec4 in_color;

out vec3 position;
out vec4 color;
out vec3 normal;

void main()
{
	position = in_position.xyz;
	gl_Position = u_transform * in_position;
	color = in_color;
	normal = in_normal;
}
)";

static char const shadow_fragment_source[] =
R"(#version 330

uniform vec4 u_light_position;
uniform vec4 u_camera_position;

uniform mat4 u_shadow_transform;
uniform sampler2DShadow u_shadow_map;

uniform vec4 u_material;

in vec3 position;
in vec4 color;
in vec3 normal;

out vec4 out_color;

void main()
{
	vec3 n = normalize(normal);

	vec3 r = u_light_position.xyz - u_light_position.w * position;
	r = normalize(r);

	vec3 v = u_camera_position.xyz - u_camera_position.w * position;
	v = normalize(v);

	vec3 d = 2.0 * dot(r, n) * n - r;

	vec3 light_color = vec3(1.0, 1.0, 1.0);
	vec3 specular_color = vec3(1.0, 1.0, 1.0);

	float light_dot = dot(n, r);

	vec3 ambient = u_material.x * color.rgb;
	vec3 diffuse = u_material.y * max(0.0, light_dot) * color.rgb * light_color;
	vec3 specular = u_material.z * (pow(max(0.0, dot(d, v)), u_material.w)) * specular_color;

	vec3 result = ambient;

	vec4 shadow_space = u_shadow_transform * vec4(position, 1.0);

	vec3 tc = shadow_space.xyz / shadow_space.w;
	tc = tc * 0.5 + vec3(0.5);

	float fragment_depth = tc.z;

	float shadow_value = 1.0;

	if (tc.x >= 0.0 && tc.x <= 1.0 && tc.y >= 0.0 && tc.y <= 1.0 && tc.z >= 0.0)
	{
		if (light_dot > 0.0)
		{
			shadow_value = texture(u_shadow_map, tc);
		}
	}

	result += shadow_value * (diffuse + specular);

	out_color = vec4(result, color.a);
}
)";

struct shadow_renderer
{
	struct material
	{
		float ambient;
		float diffuse;
		float specular;
		float shininess;
	};

	struct render_state
	{
		struct material material;
		gfx::mesh const * mesh;
		bool casts_shadow;
		std::optional<math::box<float, 3>> bbox;
	};

	struct light
	{
		math::vector<float, 4> position;
	};

	struct render_options
	{
		gfx::framebuffer const * framebuffer;
		GLenum draw_buffer;
		math::box<int, 2> viewport;
		math::matrix<float, 4, 4> transform;
		math::point<float, 3> camera_position;
		struct light light;
	};

	void push(render_state const & state);

	void render(render_options const & options);

	gfx::texture_2d const & depth_texture() const { return builder_.texture(); }

private:
	std::vector<render_state> render_states_;
	gfx::program program_{shadow_vertex_source, shadow_fragment_source};

	shadow_map_builder builder_{1024, 1024};
};

void shadow_renderer::push(render_state const & state)
{
	render_states_.push_back(state);
}

void shadow_renderer::render(render_options const & options)
{
	for (auto const & state : render_states_)
	{
		if (state.casts_shadow && state.bbox)
			builder_.push({state.mesh, *state.bbox});
	}

	shadow_map_builder::render_options build_options;
	build_options.light = {options.light.position[0], options.light.position[1], options.light.position[2]};
	builder_.build(build_options);

	options.framebuffer->bind();
	gl::DrawBuffer(options.draw_buffer);
	gl::Viewport(options.viewport[0].min, options.viewport[1].min, options.viewport[0].length(), options.viewport[1].length());

	gl::Enable(gl::DEPTH_TEST);
	gl::DepthFunc(gl::LEQUAL);

	gl::Enable(gl::CULL_FACE);
	gl::CullFace(gl::BACK);

	program_.bind();
	program_["u_transform"] = options.transform;
	program_["u_light_position"] = options.light.position;
	program_["u_camera_position"] = math::homogeneous(options.camera_position);
	program_["u_shadow_transform"] = builder_.transform();
	program_["u_shadow_map"] = 0;

	gl::ActiveTexture(gl::TEXTURE0);
	builder_.texture().bind();

	for (auto const & state : render_states_)
	{
		program_["u_material"] = math::vector{state.material.ambient, state.material.diffuse, state.material.specular, state.material.shininess};
		state.mesh->draw();
	}

	render_states_.clear();
}

struct vertex
{
	math::point<float, 3> position;
	math::vector<float, 3> normal;
	gfx::color_rgba color;
};

struct shadow_app
	: app::application_base
{
	math::spherical_camera camera;

	shadow_renderer renderer;

	gfx::mesh plane_mesh;
	shadow_renderer::material plane_material;

	gfx::mesh cube_mesh;
	shadow_renderer::material cube_material;
	math::box<float, 3> cube_bbox;

	gfx::mesh sphere_mesh;
	shadow_renderer::material sphere_material;
	math::box<float, 3> sphere_bbox;

	gfx::mesh torus_mesh;
	shadow_renderer::material torus_material;
	math::box<float, 3> torus_bbox;

	util::clock<std::chrono::duration<float>> clock;
	float time = 0.f;
	bool paused = false;

	shadow_app(options const &, context const &);

	void on_event(app::resize_event const & event) override;

	void on_event(app::mouse_move_event const & event) override;

	void on_event(app::mouse_wheel_event const & event) override;

	void on_event(app::key_event const & event) override;

	void update() override {}
	void present() override;
};

shadow_app::shadow_app(options const &, context const & context)
{
	context.vsync(true);

	camera.near_clip = 0.1f;
	camera.far_clip = 1000.f;
	camera.fov_y = math::rad(45.f);

	camera.azimuthal_angle = 0.f;
	camera.elevation_angle = math::rad(30.f);
	camera.target = {0.f, 0.f, 0.f};
	camera.distance = 10.f;

	{

		std::vector<vertex> vertices;

		vertices.push_back({{-10.f, -10.f, 0.f}, {0.f, 0.f, 1.f}, {127, 127, 127, 255}});
		vertices.push_back({{ 10.f, -10.f, 0.f}, {0.f, 0.f, 1.f}, {127, 127, 127, 255}});
		vertices.push_back({{-10.f,  10.f, 0.f}, {0.f, 0.f, 1.f}, {127, 127, 127, 255}});
		vertices.push_back({{-10.f,  10.f, 0.f}, {0.f, 0.f, 1.f}, {127, 127, 127, 255}});
		vertices.push_back({{ 10.f, -10.f, 0.f}, {0.f, 0.f, 1.f}, {127, 127, 127, 255}});
		vertices.push_back({{ 10.f,  10.f, 0.f}, {0.f, 0.f, 1.f}, {127, 127, 127, 255}});

		plane_mesh.setup<math::point<float, 3>, math::vector<float, 3>, gfx::normalized<gfx::color_rgba>>();
		plane_mesh.load(vertices, gl::TRIANGLES, gl::STATIC_DRAW);

		plane_material.ambient = 0.2f;
		plane_material.diffuse = 1.f;
		plane_material.specular = 0.f;
		plane_material.shininess = 1.f;
	}

	{
		auto cube = math::box<float, 3>{{{-1.f, 1.f}, {-1.f, 1.f}, {0.f, 5.f}}};

		auto vertices = math::vertices(cube);
		auto faces = math::faces(cube);
		auto normals = math::flat_normals(vertices, faces);
		auto flat_vertices = math::deindex(vertices, faces);

		std::vector<vertex> mesh_vertices;

		for (std::size_t i = 0; i < flat_vertices.size(); ++i)
		{
			mesh_vertices.push_back({flat_vertices[i][0], normals[i], {255, 127, 127, 255}});
			mesh_vertices.push_back({flat_vertices[i][1], normals[i], {255, 127, 127, 255}});
			mesh_vertices.push_back({flat_vertices[i][2], normals[i], {255, 127, 127, 255}});
		}

		cube_mesh.setup<math::point<float, 3>, math::vector<float, 3>, gfx::normalized<gfx::color_rgba>>();
		cube_mesh.load(mesh_vertices, gl::TRIANGLES, gl::STATIC_DRAW);

		cube_material.ambient = 0.2f;
		cube_material.diffuse = 1.f;
		cube_material.specular = 0.6f;
		cube_material.shininess = 10000.f;

		cube_bbox = cube;
	}

	{
		std::vector<vertex> vertices;

		math::point<float, 3> const position = {3.f, 2.f, 3.f};

		float const radius = 1.f;

		int const N = 24;

		gfx::color_rgba color { 63, 63, 191, 255 };

		for (int j = - N + 1; j < N; ++j)
		{
			for (int i = 0; i < 4 * N; ++i)
			{
				float a = (math::pi * i) / (2 * N);
				float b = (math::pi * j) / (2 * N);

				math::vector n{std::cos(a) * std::cos(b), std::sin(a) * std::cos(b), std::sin(b)};

				vertices.push_back({position + radius * n, n, color});
			}
		}

		vertices.push_back({position + math::vector{0.f, 0.f, -radius}, {0.f, 0.f, -1.f}, color});
		vertices.push_back({position + math::vector{0.f, 0.f,  radius}, {0.f, 0.f,  1.f}, color});

		std::vector<math::triangle<std::uint32_t>> indices;

		auto idx = [](int i, int j) -> std::uint32_t { return (i % (4 * N)) + 4 * N * (j + N - 1); };

		for (int j = - N + 1; j + 1 < N; ++j)
		{
			for (int i = 0; i < 4 * N; ++i)
			{
				indices.push_back({idx(i, j), idx(i + 1, j), idx(i, j + 1)});
				indices.push_back({idx(i, j + 1), idx(i + 1, j), idx(i + 1, j + 1)});
			}
		}

		for (int i = 0; i < 4 * N; ++i)
		{
			indices.push_back({idx(i, 1 - N), (2 * N - 1) * (4 * N), idx(i + 1, 1 - N)});
		}

		for (int i = 0; i < 4 * N; ++i)
		{
			indices.push_back({idx(i, N - 1), idx(i + 1, N - 1), (2 * N - 1) * (4 * N) + 1});
		}

		sphere_mesh.setup<math::point<float, 3>, math::vector<float, 3>, gfx::normalized<gfx::color_rgba>>();
		sphere_mesh.load(vertices, indices, gl::STATIC_DRAW);

		sphere_material.ambient = 0.2f;
		sphere_material.diffuse = 1.f;
		sphere_material.specular = 1.f;
		sphere_material.shininess = 100.f;

		sphere_bbox[0] = {position[0] - radius, position[0] + radius};
		sphere_bbox[1] = {position[1] - radius, position[1] + radius};
		sphere_bbox[2] = {position[2] - radius, position[2] + radius};
	}

	{
		std::vector<vertex> vertices;

		math::point<float, 3> const position = {-3.f, 2.f, 3.f};

		float const radius1 = 1.f;
		float const radius2 = 0.2f;

		int const N = 72;
		int const M = 24;

		gfx::color_rgba color { 63, 63, 191, 255 };

		for (int j = 0; j < M; ++j)
		{
			for (int i = 0; i < N; ++i)
			{
				float a = (2.f * math::pi * i) / N;
				float b = (2.f * math::pi * j) / M;

				math::vector r{std::cos(a), std::sin(a), 0.f};

				math::vector n{std::cos(a) * std::cos(b), std::sin(a) * std::cos(b), std::sin(b)};

				vertices.push_back({position + radius1 * r + radius2 * n, n, color});
			}
		}

		std::vector<math::triangle<std::uint32_t>> indices;

		auto idx = [](int i, int j) -> std::uint32_t { return (i % N) + N * (j % M); };

		for (int j = 0; j < M; ++j)
		{
			for (int i = 0; i < N; ++i)
			{
				indices.push_back({idx(i, j), idx(i + 1, j), idx(i, j + 1)});
				indices.push_back({idx(i, j + 1), idx(i + 1, j), idx(i + 1, j + 1)});
			}
		}

		torus_mesh.setup<math::point<float, 3>, math::vector<float, 3>, gfx::normalized<gfx::color_rgba>>();
		torus_mesh.load(vertices, indices, gl::STATIC_DRAW);

		torus_material.ambient = 0.2f;
		torus_material.diffuse = 1.f;
		torus_material.specular = 1.f;
		torus_material.shininess = 100.f;

		torus_bbox[0] = {position[0] - radius1 - radius2, position[0] + radius1 + radius2};
		torus_bbox[1] = {position[1] - radius1 - radius2, position[1] + radius1 + radius2};
		torus_bbox[2] = {position[2] - radius2, position[2] + radius2};
	}
}

void shadow_app::on_event(app::resize_event const & event)
{
	app::application_base::on_event(event);

	camera.set_fov(camera.fov_y, (1.f * event.size[0]) / event.size[1]);
}

void shadow_app::on_event(app::mouse_move_event const & event)
{
	auto const old_mouse = state().mouse;

	app::application_base::on_event(event);

	if (state().mouse_button_down.contains(app::mouse_button::middle))
	{
		auto const delta = event.position - old_mouse;
		camera.azimuthal_angle -= delta[0] * 0.01f;
		camera.elevation_angle += delta[1] * 0.01f;
	}
}

void shadow_app::on_event(app::mouse_wheel_event const & event)
{
	app::application_base::on_event(event);

	camera.distance *= std::pow(0.8f, event.delta);
}

void shadow_app::on_event(app::key_event const & event)
{
	app::application_base::on_event(event);

	if (event.down && event.key == app::keycode::SPACE)
		paused = !paused;
}

void shadow_app::present()
{
	if (!paused)
		time += clock.restart().count() / 4.f;
	else
		clock.restart();

	math::vector<float, 3> light_dir = {std::cos(time), std::sin(time), 0.5f};

	gfx::framebuffer::null().bind();
	gl::DrawBuffer(gl::BACK);

	gl::Viewport(0, 0, state().size[0], state().size[1]);

	gl::ClearColor(0.7f, 0.7f, 1.f, 0.f);
	gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);

	gl::Enable(gl::DEPTH_TEST);
	gl::DepthFunc(gl::LEQUAL);

	renderer.push({plane_material, &plane_mesh, false, std::nullopt});
	renderer.push({sphere_material, &sphere_mesh, true, sphere_bbox});
	renderer.push({cube_material, &cube_mesh, true, cube_bbox});
	renderer.push({torus_material, &torus_mesh, true, torus_bbox});

	shadow_renderer::render_options options;
	options.framebuffer = &gfx::framebuffer::null();
	options.viewport = {{{0, state().size[0]}, {0, state().size[1]}}};
	options.draw_buffer = gl::BACK;
	options.transform = camera.transform();
	options.light.position = math::homogeneous(light_dir);
	options.camera_position = camera.position();

	renderer.render(options);

	gfx::check_error();
}

namespace psemek::app
{

	std::unique_ptr<application::factory> make_application_factory()
	{
		return default_application_factory<shadow_app>({.name = "Shadow example"});
	}

}