#include <psemek/ui/painter_impl.hpp>

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

#include <psemek/util/overload.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 bitmap_text_vs[] =
R"(#version 330

uniform mat4 u_transform;

uniform vec2 u_atlas_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_atlas_size;
}
)";

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

uniform sampler2D u_atlas;

in vec2 texcoord;
in vec4 color;

out vec4 out_color;

void main()
{
	out_color = vec4(color.rgb, texture(u_atlas, texcoord) * color.a);
}
)";

		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_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);
	out_color = vec4(mix(tex_color.rgb, color.rgb, color.a), tex_color.a);
}
)";

		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 bitmap_text_vertex
		{
			geom::point<float, 2> position;
			std::uint32_t depth;
			gfx::color_rgba color;
			geom::vector<std::uint16_t, 2> texcoords;
		};
		static_assert(sizeof(bitmap_text_vertex) == 20);

		struct bitmap_text_batch
		{
			gfx::texture_2d const * atlas = nullptr;
			std::vector<bitmap_text_vertex> vertices{};
			std::vector<std::uint32_t> indices{};
		};

		bool operator == (bitmap_text_batch const & b1, bitmap_text_batch const & b2)
		{
			return b1.atlas == b2.atlas;
		}

		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;
			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;
		}

		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 bitmap_text_program;
		gfx::mesh bitmap_text_mesh;

		gfx::program textured_program;
		gfx::mesh textured_mesh;

		int stencil_level = 0;
		static constexpr int max_stencil_level = 8;

		geom::box<float, 2> draw_bbox;

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

		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}
		, bitmap_text_program{bitmap_text_vs, bitmap_text_fs}
		, textured_program{textured_vs, textured_fs}
	{
		bitmap_text_program.bind();
		bitmap_text_program["u_atlas"] = 0;

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

		colored_mesh.setup<geom::point<float, 2>, std::uint32_t, gfx::normalized<gfx::color_rgba>>();
		bitmap_text_mesh.setup<geom::point<float, 2>, std::uint32_t, gfx::normalized<gfx::color_rgba>, geom::vector<std::uint16_t, 2>>();
		textured_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() = default;

	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 & 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({tri[0], depth, color});
		batch.vertices.push_back({tri[1], depth, color});
		batch.vertices.push_back({tri[2], depth, color});

		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_glyph(font const & f, char32_t c, geom::box<float, 2> const & rect, gfx::color_rgba const & color)
	{
		auto tbox = f.texcoords(c);
		if (!tbox) return;

		auto size = f.atlas().size();

		auto const tc = [&](float x, float y)
		{
			auto t = tbox->corner(x, y);
			geom::vector<std::uint16_t, 2> r;
			r[0] = geom::clamp(t[0], {0, size[0]});
			r[1] = geom::clamp(t[1], {0, size[1]});
			return r;
		};

		auto & batch = impl().batch<bitmap_text_batch>(bitmap_text_batch{&f.atlas()});

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

		auto round = [](geom::point<float, 2> p)
		{
			p[0] = std::round(p[0]);
			p[1] = std::round(p[1]);
			return p;
		};

		batch.vertices.push_back({round(rect.corner(0.f, 0.f)), depth, color, tc(0.f, 0.f)});
		batch.vertices.push_back({round(rect.corner(1.f, 0.f)), depth, color, tc(1.f, 0.f)});
		batch.vertices.push_back({round(rect.corner(0.f, 1.f)), depth, color, tc(0.f, 1.f)});
		batch.vertices.push_back({round(rect.corner(1.f, 1.f)), depth, color, tc(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_image(geom::box<float, 2> const & rect, gfx::texture_2d const & tex, gfx::color_rgba const & color)
	{
		geom::box<float, 2> part;
		part[0] = {0.f, tex.width()};
		part[1] = {0.f, tex.height()};
		draw_subimage(rect, tex, part, color);
	}

	void painter_impl::draw_subimage(geom::box<float, 2> const & rect, gfx::texture_2d const & tex, geom::box<float, 2> const & part, gfx::color_rgba const & color)
	{
		auto & batch = impl().batch<textured_batch>(textured_batch{&tex});

		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, part.corner(0.f, 0.f)});
		batch.vertices.push_back({rect.corner(1.f, 0.f), depth, color, part.corner(1.f, 0.f)});
		batch.vertices.push_back({rect.corner(0.f, 1.f), depth, color, part.corner(0.f, 1.f)});
		batch.vertices.push_back({rect.corner(1.f, 1.f), depth, color, 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("Only 8 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 = 2 * (1 << impl().stencil_level) - 1;
		batch.mask = (1 << impl().stencil_level) - 1;

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

		++impl().stencil_level;
	}

	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 = (1 << impl().stencil_level) - 1;
		batch.mask = (1 << impl().stencil_level) - 1;

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

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

			clear_batch.func = gl::EQUAL;
			clear_batch.ref = (1 << impl().stencil_level) / 2 - 1;
			clear_batch.mask = (1 << impl().stencil_level) / 2 - 1;

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

		draw_rect(impl().draw_bbox, {0, 0, 0, 255});

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

			batch.func = gl::NEVER;
			batch.ref = 0;
			batch.mask = 0;

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

		--impl().stencil_level;
	}

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

		gl::ActiveTexture(gl::TEXTURE0);

		auto batch_visitor = util::overload(
		[&](colored_batch const & b){
			impl().colored_program.bind();
			impl().colored_mesh.load(b.vertices, b.indices, gl::TRIANGLES);
			impl().colored_mesh.draw();
		},
		[&](bitmap_text_batch const & b){
			if (!b.atlas) return;
			impl().bitmap_text_program.bind();
			impl().bitmap_text_program["u_atlas_size"] = geom::cast<float>(b.atlas->size());
			b.atlas->bind();
			impl().bitmap_text_mesh.load(b.vertices, b.indices, gl::TRIANGLES);
			impl().bitmap_text_mesh.draw();
		},
		[&](textured_batch const & b){
			if (!b.texture) return;
			impl().textured_program.bind();
			impl().textured_program["u_texture_size"] = geom::cast<float>(b.texture->size());
			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().depth = 0;
		impl().batches.clear();
		impl().stencil_level = 0;
		impl().draw_bbox = {};
	}

}