#include <psemek/ui/painter_impl.hpp>

#include <psemek/gfx/program.hpp>
#include <psemek/gfx/mesh.hpp>

#include <psemek/util/overload.hpp>
#include <psemek/util/to_string.hpp>
#include <psemek/log/log.hpp>

#include <variant>

namespace psemek::ui
{

	namespace
	{

		char const colored_vs[] =
R"(#version 330

uniform mat4 u_transform;

layout (location = 0) in vec2 in_position;
layout (location = 1) in float in_depth;
layout (location = 2) in vec4 in_color;

out vec4 color;

void main()
{
	gl_Position = u_transform * vec4(in_position, 1.0 - float(in_depth) / 16777216.0 * 2.0, 1.0);
	color = in_color;
}
)";

		char const colored_fs[] =
R"(#version 330

in vec4 color;

out vec4 out_color;

void main()
{
	out_color = color;
}
)";

		char const textured_vs[] =
R"(#version 330

uniform mat4 u_transform;

uniform vec2 u_texture_size;

layout (location = 0) in vec2 in_position;
layout (location = 1) in float in_depth;
layout (location = 2) in vec4 in_color;
layout (location = 3) in vec2 in_texcoord;

out vec4 color;
out vec2 texcoord;

void main()
{
	gl_Position = u_transform * vec4(in_position, 1.0 - float(in_depth) / 16777216.0 * 2.0, 1.0);
	color = in_color;
	texcoord = in_texcoord / u_texture_size;
}
)";

		char const textured_mix_fs[] =
R"(#version 330

uniform sampler2D u_texture;

in vec2 texcoord;
in vec4 color;

out vec4 out_color;

void main()
{
	vec4 tex_color = texture(u_texture, texcoord);
	if (tex_color.a == 0.0) discard;
	out_color = vec4(mix(tex_color.rgb, color.rgb, color.a), tex_color.a);
}
)";

		char const textured_multiply_fs[] =
R"(#version 330

uniform sampler2D u_texture;

in vec2 texcoord;
in vec4 color;

out vec4 out_color;

void main()
{
	vec4 tex_color = texture(u_texture, texcoord);
	if (tex_color.a == 0.0) discard;
	out_color = tex_color * color;
}
)";

		char const msdf_vs[] =
R"(#version 330

uniform mat4 u_transform;

uniform vec2 u_texture_size;

layout (location = 0) in vec2 in_position;
layout (location = 1) in float in_depth;
layout (location = 2) in vec4 in_color;
layout (location = 3) in vec2 in_texcoord;

out vec4 color;
out vec2 texcoord;

void main()
{
	gl_Position = u_transform * vec4(in_position, 1.0 - float(in_depth) / 16777216.0 * 2.0, 1.0);
	color = in_color;
	texcoord = in_texcoord / u_texture_size;
}
)";

		char const msdf_fs[] =
R"(#version 330

uniform sampler2D u_texture;
uniform float u_sdf_scale;

in vec2 texcoord;
in vec4 color;

out vec4 out_color;

float median(vec3 v) {
	return max(min(v.r, v.g), min(max(v.r, v.g), v.b));
}

void main()
{
	float sdf_value = u_sdf_scale * (median(texture(u_texture, texcoord).rgb) - 0.5);

	float aa_step = length(vec2(dFdx(sdf_value), dFdy(sdf_value))) * sqrt(2.0);

	float alpha = smoothstep(-aa_step * 0.5, aa_step * 0.5, sdf_value);

	if (alpha == 0.0) discard;

	out_color = vec4(color.rgb, color.a * alpha);
}
)";

		struct colored_vertex
		{
			geom::point<float, 2> position;
			std::uint32_t depth;
			gfx::color_rgba color;
		};
		static_assert(sizeof(colored_vertex) == 16);

		struct colored_batch
		{
			std::vector<colored_vertex> vertices;
			std::vector<std::uint32_t> indices;
		};

		bool operator == (colored_batch const &, colored_batch const &)
		{
			return true;
		}

		struct textured_vertex
		{
			geom::point<float, 2> position;
			std::uint32_t depth;
			gfx::color_rgba color;
			geom::point<float, 2> texcoords;
		};
		static_assert(sizeof(textured_vertex) == 24);

		struct textured_batch
		{
			gfx::texture_2d const * texture = nullptr;
			painter::color_mode mode = painter::color_mode::mix;
			std::vector<textured_vertex> vertices{};
			std::vector<std::uint32_t> indices{};
		};

		bool operator == (textured_batch const & b1, textured_batch const & b2)
		{
			return b1.texture == b2.texture && b1.mode == b2.mode;
		}

		struct msdf_vertex
		{
			geom::point<float, 2> position;
			std::uint32_t depth;
			gfx::color_rgba color;
			geom::point<float, 2> texcoords;
		};
		static_assert(sizeof(textured_vertex) == 24);

		struct msdf_batch
		{
			gfx::texture_2d const * texture = nullptr;
			painter::color_mode mode = painter::color_mode::mix;
			float sdf_scale{0.f};
			std::vector<textured_vertex> vertices{};
			std::vector<std::uint32_t> indices{};
		};

		bool operator == (msdf_batch const & b1, msdf_batch const & b2)
		{
			return b1.texture == b2.texture && b1.mode == b2.mode && b1.sdf_scale == b2.sdf_scale;
		}

		struct stencil_batch
		{
			std::optional<bool> enable_stencil;
			std::optional<bool> enable_depth;
			std::optional<bool> enable_color;

			GLenum sfail, dfail, dpass;
			GLenum func;
			GLint ref;
			GLuint mask;
		};

		bool operator == (stencil_batch const &, stencil_batch const &)
		{
			return false;
		}

	}

	struct painter_impl::impl
	{
		std::uint32_t depth = 0;

		gfx::program colored_program;
		gfx::mesh colored_mesh;

		gfx::program textured_mix_program;
		gfx::program textured_multiply_program;
		gfx::mesh textured_mesh;

		gfx::program msdf_program;
		gfx::mesh msdf_mesh;

		int stencil_level = 0;
		static constexpr int max_stencil_level = 255;
		std::vector<geom::box<float, 2>> stencil_bbox;

		geom::box<float, 2> draw_bbox;

		std::vector<std::variant<colored_batch, textured_batch, msdf_batch, stencil_batch>> batches;

		std::size_t max_batch_count = 0;
		std::size_t max_primitive_count = 0;

		template <typename Batch>
		Batch & batch(Batch && n)
		{
			if (!batches.empty())
			{
				auto b = std::get_if<Batch>(&batches.back());
				if (b && *b == n) return *b;
			}

			batches.push_back(std::move(n));
			return std::get<Batch>(batches.back());
		}

		impl();
	};

	painter_impl::impl::impl()
		: colored_program{colored_vs, colored_fs}
		, textured_mix_program{textured_vs, textured_mix_fs}
		, textured_multiply_program{textured_vs, textured_multiply_fs}
		, msdf_program{msdf_vs, msdf_fs}
	{
		textured_mix_program.bind();
		textured_mix_program["u_texture"] = 0;

		textured_multiply_program.bind();
		textured_multiply_program["u_texture"] = 0;

		msdf_program.bind();
		msdf_program["u_texture"] = 0;

		colored_mesh.setup<geom::point<float, 2>, std::uint32_t, gfx::normalized<gfx::color_rgba>>();
		textured_mesh.setup<geom::point<float, 2>, std::uint32_t, gfx::normalized<gfx::color_rgba>, geom::vector<float, 2>>();
		msdf_mesh.setup<geom::point<float, 2>, std::uint32_t, gfx::normalized<gfx::color_rgba>, geom::vector<float, 2>>();
	}

	painter_impl::painter_impl()
		: pimpl_{make_impl()}
	{}

	painter_impl::~painter_impl()
	{
		log::debug() << "UI painter max batch count: " << impl().max_batch_count;
		log::debug() << "UI painter max primitive count: " << impl().max_primitive_count;
	}

	void painter_impl::draw_rect(geom::box<float, 2> const & rect, gfx::color_rgba const & color)
	{
		auto & batch = impl().batch<colored_batch>({});

		std::uint32_t const depth = impl().depth++;
		std::uint32_t const base = batch.vertices.size();

		batch.vertices.push_back({rect.corner(0.f, 0.f), depth, color});
		batch.vertices.push_back({rect.corner(1.f, 0.f), depth, color});
		batch.vertices.push_back({rect.corner(0.f, 1.f), depth, color});
		batch.vertices.push_back({rect.corner(1.f, 1.f), depth, color});

		batch.indices.insert(batch.indices.end(), {base + 0, base + 1, base + 2, base + 2, base + 1, base + 3});

		impl().draw_bbox |= rect;
	}

	void painter_impl::draw_triangle(geom::triangle<geom::point<float, 2>> const & tri, gfx::color_rgba const & c0, gfx::color_rgba const & c1, gfx::color_rgba const & c2)
	{
		auto & batch = impl().batch<colored_batch>({});

		std::uint32_t const depth = impl().depth++;
		std::uint32_t const base = batch.vertices.size();

		batch.vertices.push_back({tri[0], depth, c0});
		batch.vertices.push_back({tri[1], depth, c1});
		batch.vertices.push_back({tri[2], depth, c2});

		batch.indices.insert(batch.indices.end(), {base + 0, base + 1, base + 2});

		impl().draw_bbox |= tri[0];
		impl().draw_bbox |= tri[1];
		impl().draw_bbox |= tri[2];
	}

	void painter_impl::draw_image(geom::box<float, 2> const & rect, gfx::texture_view_2d const & tex, image_options const & options)
	{
		auto & batch = impl().batch<textured_batch>(textured_batch{tex.texture, options.mode});

		std::uint32_t const depth = impl().depth++;
		std::uint32_t const base = batch.vertices.size();

		auto p00 = rect.corner(0.f, 0.f);
		auto p10 = rect.corner(1.f, 0.f);
		auto p01 = rect.corner(0.f, 1.f);
		auto p11 = rect.corner(1.f, 1.f);

		if (options.rotation != 0.f)
		{
			auto const c = rect.center();

			p00 = c + geom::rotate(p00 - c, options.rotation);
			p10 = c + geom::rotate(p10 - c, options.rotation);
			p01 = c + geom::rotate(p01 - c, options.rotation);
			p11 = c + geom::rotate(p11 - c, options.rotation);
		}

		batch.vertices.push_back({p00, depth, options.color, tex.part.corner(0.f, 0.f)});
		batch.vertices.push_back({p10, depth, options.color, tex.part.corner(1.f, 0.f)});
		batch.vertices.push_back({p01, depth, options.color, tex.part.corner(0.f, 1.f)});
		batch.vertices.push_back({p11, depth, options.color, tex.part.corner(1.f, 1.f)});

		batch.indices.insert(batch.indices.end(), {base + 0, base + 1, base + 2, base + 2, base + 1, base + 3});

		impl().draw_bbox |= rect;
	}

	void painter_impl::draw_msdf_glyph(geom::box<float, 2> const & rect, gfx::texture_view_2d const & tex, float sdf_scale, image_options const & options)
	{
		auto & batch = impl().batch<msdf_batch>(msdf_batch{tex.texture, options.mode, sdf_scale});

		std::uint32_t const depth = impl().depth++;
		std::uint32_t const base = batch.vertices.size();

		auto p00 = rect.corner(0.f, 0.f);
		auto p10 = rect.corner(1.f, 0.f);
		auto p01 = rect.corner(0.f, 1.f);
		auto p11 = rect.corner(1.f, 1.f);

		if (options.rotation != 0.f)
		{
			auto const c = rect.center();

			p00 = c + geom::rotate(p00 - c, options.rotation);
			p10 = c + geom::rotate(p10 - c, options.rotation);
			p01 = c + geom::rotate(p01 - c, options.rotation);
			p11 = c + geom::rotate(p11 - c, options.rotation);
		}

		batch.vertices.push_back({p00, depth, options.color, tex.part.corner(0.f, 0.f)});
		batch.vertices.push_back({p10, depth, options.color, tex.part.corner(1.f, 0.f)});
		batch.vertices.push_back({p01, depth, options.color, tex.part.corner(0.f, 1.f)});
		batch.vertices.push_back({p11, depth, options.color, tex.part.corner(1.f, 1.f)});

		batch.indices.insert(batch.indices.end(), {base + 0, base + 1, base + 2, base + 2, base + 1, base + 3});

		impl().draw_bbox |= rect;
	}

	void painter_impl::begin_stencil()
	{
		if (impl().stencil_level == impl().max_stencil_level)
			throw std::runtime_error(util::to_string("Only ", impl().max_stencil_level, " recursive stencil levels are supported"));

		auto & batch = impl().batch<stencil_batch>({});
		if (impl().stencil_level == 0)
			batch.enable_stencil = true;
		batch.enable_color = false;
		batch.enable_depth = false;

		batch.func = gl::EQUAL;
		batch.ref = impl().stencil_level;
		batch.mask = -1;

		batch.sfail = gl::KEEP;
		batch.dfail = gl::KEEP;
		batch.dpass = gl::INCR;

		++impl().stencil_level;
		impl().draw_bbox = {};
	}

	void painter_impl::commit_stencil()
	{
		auto & batch = impl().batch<stencil_batch>({});
		batch.enable_color = true;
		batch.enable_depth = true;

		batch.func = gl::EQUAL;
		batch.ref = impl().stencil_level;
		batch.mask = -1;

		batch.sfail = gl::KEEP;
		batch.dfail = gl::KEEP;
		batch.dpass = gl::KEEP;

		impl().stencil_bbox.push_back(impl().draw_bbox);
	}

	void painter_impl::end_stencil()
	{
		--impl().stencil_level;

		{
			auto & clear_batch = impl().batch<stencil_batch>({});
			clear_batch.enable_color = false;
			clear_batch.enable_depth = false;

			clear_batch.func = gl::LEQUAL;
			clear_batch.ref = impl().stencil_level;
			clear_batch.mask = -1;

			clear_batch.sfail = gl::KEEP;
			clear_batch.dfail = gl::KEEP;
			clear_batch.dpass = gl::REPLACE;
		}

		draw_rect(impl().stencil_bbox.back(), {0, 0, 0, 255});

		{
			auto & batch = impl().batch<stencil_batch>({});
			if (impl().stencil_level == 0)
				batch.enable_stencil = false;
			batch.enable_color = true;
			batch.enable_depth = true;

			batch.func = gl::EQUAL;
			batch.ref = impl().stencil_level;
			batch.mask = -1;

			batch.sfail = gl::KEEP;
			batch.dfail = gl::KEEP;
			batch.dpass = gl::KEEP;
		}

		impl().stencil_bbox.pop_back();
	}

	void painter_impl::start_frame(geom::box<float, 2> const & bbox)
	{
		impl().stencil_bbox.push_back(bbox);
	}

	void painter_impl::render(geom::matrix<float, 4, 4> const & transform)
	{
		impl().colored_program.bind();
		impl().colored_program["u_transform"] = transform;
		impl().textured_mix_program.bind();
		impl().textured_mix_program["u_transform"] = transform;
		impl().textured_multiply_program.bind();
		impl().textured_multiply_program["u_transform"] = transform;
		impl().msdf_program.bind();
		impl().msdf_program["u_transform"] = transform;

		gl::ActiveTexture(gl::TEXTURE0);

		std::size_t primitive_count = 0;

		auto batch_visitor = util::overload(
		[&](colored_batch const & b){
			primitive_count += b.indices.size() / 3;

			impl().colored_program.bind();
			impl().colored_mesh.load(b.vertices, b.indices, gl::TRIANGLES);
			impl().colored_mesh.draw();
		},
		[&](textured_batch const & b){
			if (!b.texture) return;

			primitive_count += b.indices.size() / 3;

			gfx::program & program = (b.mode == ui::painter::color_mode::mix) ? impl().textured_mix_program : impl().textured_multiply_program;
			program.bind();
			program["u_texture_size"] = geom::cast<float>(b.texture->size());
			b.texture->bind();
			impl().msdf_mesh.load(b.vertices, b.indices, gl::TRIANGLES);
			impl().msdf_mesh.draw();
		},
		[&](msdf_batch const & b){
			if (!b.texture) return;

			primitive_count += b.indices.size() / 3;

			gfx::program & program = impl().msdf_program;
			program.bind();
			program["u_texture_size"] = geom::cast<float>(b.texture->size());
			program["u_sdf_scale"] = b.sdf_scale;
			b.texture->bind();
			impl().textured_mesh.load(b.vertices, b.indices, gl::TRIANGLES);
			impl().textured_mesh.draw();
		},
		[&](stencil_batch const & b){
			if (b.enable_stencil)
			{
				if (*b.enable_stencil)
					gl::Enable(gl::STENCIL_TEST);
				else
					gl::Disable(gl::STENCIL_TEST);
			}

			if (b.enable_depth)
			{
				if (*b.enable_depth)
					gl::DepthMask(gl::TRUE);
				else
					gl::DepthMask(gl::FALSE);
			}

			if (b.enable_color)
			{
				if (*b.enable_color)
					gl::ColorMask(gl::TRUE, gl::TRUE, gl::TRUE, gl::TRUE);
				else
					gl::ColorMask(gl::FALSE, gl::FALSE, gl::FALSE, gl::FALSE);
			}

			gl::StencilOp(b.sfail, b.dfail, b.dpass);
			gl::StencilFunc(b.func, b.ref, b.mask);
		});

		for (auto const & b : impl().batches)
			std::visit(batch_visitor, b);

		impl().max_batch_count = std::max(impl().max_batch_count, impl().batches.size());

		impl().max_primitive_count = std::max(impl().max_primitive_count, primitive_count);

		impl().depth = 0;
		impl().batches.clear();
		impl().stencil_level = 0;
		impl().stencil_bbox.clear();
		impl().draw_bbox = {};
	}

	geom::box<float, 2> painter_impl::current_bbox() const
	{
		return impl().stencil_bbox.back();
	}

}