#include <psemek/gfx/painter2.hpp>
#include <psemek/gfx/program.hpp>
#include <psemek/gfx/buffer.hpp>
#include <psemek/gfx/array.hpp>
#include <psemek/gfx/gl.hpp>
#include <psemek/math/translation.hpp>
#include <psemek/math/rotation.hpp>
#include <psemek/math/scale.hpp>

namespace psemek::gfx
{

	namespace
	{

		// MVP (Model-View-Projection) matrix is assumed to be always a 2D affine transform
		// Thus, it can be passed as 2x 3D vectors (the rows of the corresponding 2x3 affine matrix)

		// MVPW (Model-View-Projection-Window) matrix composes MVP with the viewport transform ([-1..1] -> [0..width] in pixels)
		// Its inverse is used in gradient computations for analytical SDF rendering

		char const sdf_vertex_common[] = R"(
uniform vec2 u_viewport_size;

layout (location = 0) in vec3 in_mvp_0;
layout (location = 1) in vec3 in_mvp_1;
layout (location = 2) in vec2 in_mvpw_inv_0;
layout (location = 3) in vec2 in_mvpw_inv_1;
layout (location = 4) in vec4 in_color;
layout (location = 5) in float in_half_width;

mat2 adj(mat2 m)
{
	return mat2(
		m[1][1], -m[1][0],
		-m[0][1], m[0][0]
	);
}

out mat2 mvpw_inv_t;
out mat2 mvpw_t;
out vec2 vertex;
out vec4 color;
out float half_width;

mat2x3 mvp_t = mat2x3(in_mvp_0, in_mvp_1);
mat2 w = mat2(u_viewport_size.x / 2.0, 0.0, 0.0, - u_viewport_size.y / 2.0);
mat2 mvpw_adj_t = adj(mat2(mvp_t) * w);

vec2 extend(vec2 n1, vec2 n2)
{
	// Compute the screen-space normals
	n1 = normalize(n1 * mvpw_adj_t);
	n2 = normalize(n2 * mvpw_adj_t);

	// Find the screen-space bisector for n1 and n2
	vec2 n = (n1 + n2) / (1.0 + dot(n1, n2));

	// Transform to object space
	return n * mvpw_inv_t;
}

vec2 get_vertex();

void main()
{
	mvpw_t = mat2(mvp_t) * w;
	mvpw_inv_t = mat2(in_mvpw_inv_0, in_mvpw_inv_1);

	vertex = get_vertex();
	color = in_color;
	half_width = in_half_width;
	gl_Position = vec4(vec3(vertex, 1.0) * mvp_t, 0.0, 1.0);
}
)";

		char const sdf_fill_fragment_common[] = R"(
in vec2 vertex;
in mat2 mvpw_inv_t;
in vec4 color;

layout (location = 0) out vec4 out_color;

// (grad_x, grad_y, distance)
vec3 eval_sdf(vec2 p);

void main()
{
	vec3 sdf = eval_sdf(vertex);

	vec2 d_grad_screen = mvpw_inv_t * sdf.xy;
	float d_grad_screen_length = length(d_grad_screen);
	float aa = d_grad_screen_length * sqrt(0.5);

	out_color = color * smoothstep(aa, -aa, sdf.z);
}
)";

		char const quad_vs[] = R"(
const vec2 QUAD_VERTICES[6] = vec2[6](
	vec2(-1.0, -1.0),
	vec2( 1.0, -1.0),
	vec2(-1.0,  1.0),
	vec2(-1.0,  1.0),
	vec2( 1.0, -1.0),
	vec2( 1.0,  1.0)
);

vec2 get_vertex()
{
	vec2 v = QUAD_VERTICES[gl_VertexID];

	// Extend the quad in screen-space by half-width and an extra 1px for anti-aliasing
	v += (in_half_width + 1.0) * extend(vec2(v.x, 0.0), vec2(0.0, v.y));

	return v;
}
)";

		char const circle_sdf[] = R"(
vec3 eval_sdf(vec2 p)
{
	float l = length(p);
	return vec3(p / l, l - 1.0);
}
)";

		char const square_sdf[] = R"(
vec3 eval_sdf(vec2 p)
{
	vec2 delta = abs(p) - vec2(1.0);

	if (delta.x > 0.0 || delta.y > 0.0) {
		vec2 g = max(delta, vec2(0.0));
		float l = length(g);
		return vec3(g * (sign(p) / l), l);
	} else {
		if (delta.x > delta.y) {
			return vec3(sign(p.x), 0.0, delta.x);
		} else {
			return vec3(0.0, sign(p.y), delta.y);
		}
	}
}
)";

		char const sdf_stroke_fragment_common[] = R"(
in vec2 vertex;
in mat2 mvpw_t;
in mat2 mvpw_inv_t;
in vec4 color;
in float half_width;

layout (location = 0) out vec4 out_color;

vec3 eval_sdf(vec2 p);

void main()
{
	// Newton-Raphson iteration to find the closest point on the boundary
	vec2 vertex_screen = vertex * mvpw_t;
	vec2 p_screen = vertex_screen;
	for (int i = 0; i < 4; ++i)
	{
		vec2 p_local = p_screen * mvpw_inv_t;
		vec3 sdf_value = eval_sdf(p_local);
		p_screen = p_local * mvpw_t;
		vec2 d_grad_screen = mvpw_inv_t * sdf_value.xy;
		p_screen -= (sdf_value.z / dot(d_grad_screen, d_grad_screen)) * d_grad_screen;
	}

	float d = length(vertex_screen - p_screen);
	float aa = sqrt(0.5);

	out_color = color * smoothstep(aa, -aa, d - half_width);
}
)";

		gfx::shader_source const quad_vs_source{gl::sys::shader_prelude(), sdf_vertex_common, quad_vs};

		using transform_type = painter2::transform_type;
		using fill_style = painter2::fill_style;
		using stroke_style = painter2::stroke_style;

		struct context
		{
			gfx::buffer instance_buffer;

			gfx::program fill_circle_program{quad_vs_source, {gl::sys::shader_prelude(), sdf_fill_fragment_common, circle_sdf}};
			gfx::program fill_square_program{quad_vs_source, {gl::sys::shader_prelude(), sdf_fill_fragment_common, square_sdf}};
			gfx::program stroke_circle_program{quad_vs_source, {gl::sys::shader_prelude(), sdf_stroke_fragment_common, circle_sdf}};
			gfx::program stroke_square_program{quad_vs_source, {gl::sys::shader_prelude(), sdf_stroke_fragment_common, square_sdf}};

			gfx::array array;

			context();
		};

		struct primitive_instance
		{
			transform_type mvp;
			math::matrix<float, 2, 2> mvpw_inv;
			gfx::color_rgba color;
			float half_width; // unused when filling
		};

		struct batch
		{
			gfx::program * program;
			std::vector<primitive_instance> instances;
		};

		struct frame_builder_impl
			: painter2::frame_builder
		{
			frame_builder_impl(context & context, math::vector<int, 2> viewport_size);

			void push_matrix(transform_type const & matrix) override;
			void pop_matrix() override;
			void reset_matrix() override;

			void fill_circle(math::point<float, 2> const & center, float radius, fill_style const & style) override;
			void fill_ellipse(math::point<float, 2> const & center, math::vector<float, 2> const & radii, float rotation, fill_style const & style) override;
			void fill_square(math::point<float, 2> const & center, float extent, float rotation, fill_style const & style) override;
			void fill_rectangle(math::point<float, 2> const & center, math::vector<float, 2> const & extent, float rotation, fill_style const & style) override;

			void stroke_circle(math::point<float, 2> const & center, float radius, stroke_style const & style) override;
			void stroke_ellipse(math::point<float, 2> const & center, math::vector<float, 2> const & radii, float rotation, stroke_style const & style) override;
			void stroke_square(math::point<float, 2> const & center, float extent, float rotation, stroke_style const & style) override;
			void stroke_rectangle(math::point<float, 2> const & center, math::vector<float, 2> const & extent, float rotation, stroke_style const & style) override;

			void flush() override;

			~frame_builder_impl() override
			{
				flush();
			}

		private:
			context & context_;
			math::vector<int, 2> viewport_size_;
			transform_type viewport_transform_;

			std::vector<transform_type> matrix_stack_;

			std::vector<batch> batches_;

			batch * get_batch(gfx::program * program);

			void add_ellipse(math::point<float, 2> const & center, math::vector<float, 2> const & radii, float rotation, fill_style const & style);
			void add_rectangle(math::point<float, 2> const & center, math::vector<float, 2> const & extent, float rotation, fill_style const & style);

			void add_ellipse(math::point<float, 2> const & center, math::vector<float, 2> const & radii, float rotation, stroke_style const & style);
			void add_rectangle(math::point<float, 2> const & center, math::vector<float, 2> const & extent, float rotation, stroke_style const & style);

			primitive_instance make_primitive(transform_type const & model, fill_style const & style);
			primitive_instance make_primitive(transform_type const & model, stroke_style const & style);

			transform_type make_transform(math::point<float, 2> const & center, math::vector<float, 2> const & scale, float rotation);
			std::pair<transform_type, math::matrix<float, 2, 2>> make_mvp(transform_type const & model);
		};

		context::context()
		{
			array.bind();
			instance_buffer.bind();
			gl::EnableVertexAttribArray(0);
			gl::EnableVertexAttribArray(1);
			gl::EnableVertexAttribArray(2);
			gl::EnableVertexAttribArray(3);
			gl::EnableVertexAttribArray(4);
			gl::EnableVertexAttribArray(5);
			gl::VertexAttribDivisor(0, 1);
			gl::VertexAttribDivisor(1, 1);
			gl::VertexAttribDivisor(2, 1);
			gl::VertexAttribDivisor(3, 1);
			gl::VertexAttribDivisor(4, 1);
			gl::VertexAttribDivisor(5, 1);
			gl::VertexAttribPointer(0, 3, gl::FLOAT, gl::FALSE, sizeof(primitive_instance), (char *)offsetof(primitive_instance, mvp));
			gl::VertexAttribPointer(1, 3, gl::FLOAT, gl::FALSE, sizeof(primitive_instance), (char *)offsetof(primitive_instance, mvp) + 3 * sizeof(float));
			gl::VertexAttribPointer(2, 2, gl::FLOAT, gl::FALSE, sizeof(primitive_instance), (char *)offsetof(primitive_instance, mvpw_inv));
			gl::VertexAttribPointer(3, 2, gl::FLOAT, gl::FALSE, sizeof(primitive_instance), (char *)offsetof(primitive_instance, mvpw_inv) + 2 * sizeof(float));
			gl::VertexAttribPointer(4, 4, gl::UNSIGNED_BYTE, gl::TRUE, sizeof(primitive_instance), (char *)offsetof(primitive_instance, color));
			gl::VertexAttribPointer(5, 1, gl::FLOAT, gl::FALSE, sizeof(primitive_instance), (char *)offsetof(primitive_instance, half_width));

			gfx::array::null().bind();
		}

		frame_builder_impl::frame_builder_impl(context & context, math::vector<int, 2> viewport_size)
			: context_(context)
			, viewport_size_(viewport_size)
			, viewport_transform_(
				math::translation{math::cast<float>(viewport_size) / 2.f}.transform() *
				math::scale{math::pointwise_mult(
					math::cast<float>(viewport_size),
					math::vector{0.5f, -0.5f}
				)}.transform()
			)
		{
			matrix_stack_.push_back(transform_type::identity());
		}

		void frame_builder_impl::push_matrix(painter2::transform_type const & matrix)
		{
			matrix_stack_.push_back(matrix_stack_.back() * matrix);
		}

		void frame_builder_impl::pop_matrix()
		{
			if (matrix_stack_.size() > 1)
				matrix_stack_.pop_back();
		}

		void frame_builder_impl::reset_matrix()
		{
			matrix_stack_.resize(1);
		}

		void frame_builder_impl::fill_circle(math::point<float, 2> const & center, float radius, fill_style const & style)
		{
			add_ellipse(center, {radius, radius}, 0.f, style);
		}

		void frame_builder_impl::fill_ellipse(math::point<float, 2> const & center, math::vector<float, 2> const & radii, float rotation, fill_style const & style)
		{
			add_ellipse(center, radii, rotation, style);
		}

		void frame_builder_impl::fill_square(math::point<float, 2> const & center, float extent, float rotation, fill_style const & style)
		{
			add_rectangle(center, {extent, extent}, rotation, style);
		}

		void frame_builder_impl::fill_rectangle(math::point<float, 2> const & center, math::vector<float, 2> const & extent, float rotation, fill_style const & style)
		{
			add_rectangle(center, extent, rotation, style);
		}

		void frame_builder_impl::stroke_circle(math::point<float, 2> const & center, float radius, stroke_style const & style)
		{
			add_ellipse(center, {radius, radius}, 0.f, style);
		}

		void frame_builder_impl::stroke_ellipse(math::point<float, 2> const & center, math::vector<float, 2> const & radii, float rotation, stroke_style const & style)
		{
			add_ellipse(center, radii, rotation, style);
		}

		void frame_builder_impl::stroke_square(math::point<float, 2> const & center, float extent, float rotation, stroke_style const & style)
		{
			add_rectangle(center, {extent, extent}, rotation, style);
		}

		void frame_builder_impl::stroke_rectangle(math::point<float, 2> const & center, math::vector<float, 2> const & extent, float rotation, stroke_style const & style)
		{
			add_rectangle(center, extent, rotation, style);
		}

		void frame_builder_impl::flush()
		{
			if (batches_.empty())
				return;

			gl::Enable(gl::BLEND);
			gl::BlendEquation(gl::FUNC_ADD);
			gl::BlendFunc(gl::ONE, gl::ONE_MINUS_SRC_ALPHA);

			auto const viewport_size_f = math::cast<float>(viewport_size_);

			context_.array.bind();

			for (auto const & batch : batches_)
			{
				if (batch.instances.empty()) continue;

				auto & program = *batch.program;
				program.bind();
				program["u_viewport_size"] = viewport_size_f;
				context_.instance_buffer.load(batch.instances);
				gl::DrawArraysInstanced(gl::TRIANGLES, 0, 6, batch.instances.size());
			}

			gfx::array::null().bind();
			gfx::program::null().bind();

			batches_.clear();
		}

		batch * frame_builder_impl::get_batch(gfx::program * program)
		{
			if (batches_.empty() || batches_.back().program != program)
				batches_.push_back({.program = program, .instances = {}});
			return &batches_.back();
		}

		void frame_builder_impl::add_ellipse(math::point<float, 2> const & center, math::vector<float, 2> const & radii, float rotation, fill_style const & style)
		{
			get_batch(&context_.fill_circle_program)->instances.push_back(make_primitive(make_transform(center, radii, rotation), style));
		}

		void frame_builder_impl::add_rectangle(math::point<float, 2> const & center, math::vector<float, 2> const & extent, float rotation, fill_style const & style)
		{
			get_batch(&context_.fill_square_program)->instances.push_back(make_primitive(make_transform(center, extent, rotation), style));
		}

		void frame_builder_impl::add_ellipse(math::point<float, 2> const & center, math::vector<float, 2> const & radii, float rotation, stroke_style const & style)
		{
			get_batch(&context_.stroke_circle_program)->instances.push_back(make_primitive(make_transform(center, radii, rotation), style));
		}

		void frame_builder_impl::add_rectangle(math::point<float, 2> const & center, math::vector<float, 2> const & extent, float rotation, stroke_style const & style)
		{
			get_batch(&context_.stroke_square_program)->instances.push_back(make_primitive(make_transform(center, extent, rotation), style));
		}

		primitive_instance frame_builder_impl::make_primitive(transform_type const & model, fill_style const & style)
		{
			auto [mvp, mvpw_inv] = make_mvp(model);

			return {
				.mvp = mvp,
				.mvpw_inv = mvpw_inv,
				.color = style.color,
				.half_width = 0.f,
			};
		}

		primitive_instance frame_builder_impl::make_primitive(transform_type const & model, stroke_style const & style)
		{
			auto [mvp, mvpw_inv] = make_mvp(model);

			return {
				.mvp = mvp,
				.mvpw_inv = mvpw_inv,
				.color = style.color,
				.half_width = style.width * 0.5f,
			};
		}

		transform_type frame_builder_impl::make_transform(math::point<float, 2> const & center, math::vector<float, 2> const & scale, float rotation)
		{
			return math::translation<float, 2>(center - center.zero()).transform() *
				math::plane_rotation<float, 2>(0, 1, rotation).transform() *
				math::scale<float, 2>(scale).transform();
		}

		std::pair<transform_type, math::matrix<float, 2, 2>> frame_builder_impl::make_mvp(transform_type const & model)
		{
			auto mvp = matrix_stack_.back() * model;
			auto mvpw_inv = math::inverse(viewport_transform_.linear_matrix() * mvp.linear_matrix()).value();
			return {mvp, mvpw_inv};
		}

	}

	struct painter2::impl
	{
		struct context context;
	};

	painter2::painter2()
		: pimpl_(make_impl())
	{}

	painter2::~painter2()
	{}

	std::unique_ptr<painter2::frame_builder> painter2::begin_frame(math::vector<int, 2> viewport_size)
	{
		return std::make_unique<frame_builder_impl>(impl().context, viewport_size);
	}

}