#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
{

	static 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;
}
)";

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

in vec4 color;

out vec4 out_color;

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

	static 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;
}
)";

	static 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);
}
)";

	static 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;
}
)";

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

uniform sampler2D u_texture;

in vec2 texcoord;
in vec4 color;

out vec4 out_color;

void main()
{
	out_color = mix(texture(u_texture, texcoord), color, 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 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;

		std::vector<std::variant<colored_batch, bitmap_text_batch, textured_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});
	}

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

	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();

		batch.vertices.push_back({rect.corner(0.f, 0.f), depth, color, tc(0.f, 0.f)});
		batch.vertices.push_back({rect.corner(1.f, 0.f), depth, color, tc(1.f, 0.f)});
		batch.vertices.push_back({rect.corner(0.f, 1.f), depth, color, tc(0.f, 1.f)});
		batch.vertices.push_back({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});
	}

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

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

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

		impl().depth = 0;
		impl().batches.clear();
	}

}