#include "render_elements.h"

#include "font/font_system.h"
#include "shaders/mirage_rounded_rect.hlsl.h"
#include "shaders/mirage_image.hlsl.h"
#include "shaders/mirage_text.hlsl.h"
#include "shaders/mirage_wireframe.hlsl.h"
#include "shaders/mirage_line.hlsl.h"

template<typename Derived>
void compute_rect_vertices(const Eigen::MatrixBase<Derived>& in_pos,
                           const Eigen::MatrixBase<Derived>& in_size,
                           Eigen::Matrix<float, 2, 4>&       positions,
                           float                             rotation_radians = 0.0f,
                           const Eigen::Vector2f&            scale            = Eigen::Vector2f(1.0f, 1.0f),
                           const Eigen::Vector2f&            pivot            = Eigen::Vector2f(0.5f, 0.5f)) {
	// 计算缩放后的尺寸
	Eigen::Array2f scaled_size = in_size.array() * scale.array();

	// 计算旋转矩阵
	const Eigen::Rotation2D<float> rotation(rotation_radians);
	const Eigen::Matrix2f          rot_matrix = rotation.toRotationMatrix();

	// 计算基于原点的矩形顶点（左上角在原点）
	Eigen::Matrix<float, 2, 4> rect_points;
	rect_points.col(0) = Eigen::Vector2f(0, 0);                             // 左上
	rect_points.col(1) = Eigen::Vector2f(scaled_size.x(), 0);               // 右上
	rect_points.col(2) = Eigen::Vector2f(0, scaled_size.y());               // 左下
	rect_points.col(3) = Eigen::Vector2f(scaled_size.x(), scaled_size.y()); // 右下

	// 计算锚点偏移（相对于左上角）
	const auto& pivot_offset = Eigen::Vector2f(scaled_size.x() * pivot.x(),
	                                           scaled_size.y() * pivot.y());

	// 应用旋转（绕锚点旋转）
	for (int i = 0; i < 4; ++i) {
		// 将点移到锚点为中心
		Eigen::Vector2f centered = rect_points.col(i) - pivot_offset;
		// 应用旋转
		rect_points.col(i) = rot_matrix * centered + pivot_offset;
		// 添加位置偏移
		positions.col(i) = rect_points.col(i) + in_pos;
	}
}

// 开始新一帧
void render_elements::begin_frame() {
	vertices_.clear();
	indices_.clear();
	batches_.clear();
	current_batch_index_ = -1;
	current_key_         = batch_key{};
	draw_call_count_     = 0;
	total_triangles_     = 0;
}

// 设置渲染管线
void render_elements::set_pipeline(sg_pipeline pipeline) {
	batch_key new_key = current_key_;
	new_key.pipeline  = pipeline;
	ensure_batch_compatibility(new_key);
}

// 设置纹理
void render_elements::set_texture(sg_image image) {
	batch_key new_key = current_key_;
	new_key.image     = image;
	ensure_batch_compatibility(new_key);
}

void render_elements::set_sampler(sg_sampler sampler) {
	batch_key new_key = current_key_;
	new_key.sampler   = sampler;
	ensure_batch_compatibility(new_key);
}

Eigen::Matrix4f render_elements::create_projection_matrix(const Eigen::Vector2i& in_size) {
	Eigen::Matrix4f projection_matrix = Eigen::Matrix4f::Identity();

	// 缩放因子
	const float scale_x = 2.0f / static_cast<float>(in_size.x());
	const float scale_y = -2.0f / static_cast<float>(in_size.y());  // Y轴翻转，因为窗口坐标系Y轴向下

	// 平移因子
	constexpr float translate_x = -1.0f;
	constexpr float translate_y = 1.0f;

	// 设置缩放
	projection_matrix(0, 0) = scale_x;
	projection_matrix(1, 1) = scale_y;

	// 设置平移
	projection_matrix(0, 3) = translate_x;
	projection_matrix(1, 3) = translate_y;

	return projection_matrix;
}

void render_elements::init_window_size(const Eigen::Vector2i& in_size) {
	window_size_ = in_size;
	projection_matrix_ = create_projection_matrix(in_size);
}

void render_elements::update_projection_matrix(const Eigen::Vector2i& in_size) {
	window_size_ = in_size;
	projection_matrix_ = create_projection_matrix(in_size);
}

// 确保批次兼容性
void render_elements::ensure_batch_compatibility(const batch_key& key) {
	// 如果当前没有批次或者渲染状态改变了，创建新批次
	if (current_batch_index_ < 0 || current_key_ != key) {
		current_key_ = key;

		draw_batch new_batch;
		new_batch.key          = key;
		new_batch.vertex_start = vertices_.size();
		new_batch.index_start  = indices_.size();

		batches_.push_back(new_batch);
		current_batch_index_ = batches_.size() - 1;
	}
}

// 添加矩形到当前批次
void render_elements::add_rect_to_batch(
		const Eigen::Vector2f&       in_pos,
		const Eigen::Vector2f&       in_size,
		const rect_color&            in_color,
		const mirage_vertex_param_t& in_param_a,
		const mirage_vertex_param_t& in_param_b,
		const mirage_vertex_param_t& in_param_c,
		const rect_uv&               in_uv,
		float                        in_rotation_radians,
		const Eigen::Vector2f&       in_scale,
		const Eigen::Vector2f&       in_pivot) {
	// 确保有活跃的批次
	if (current_batch_index_ < 0) {
		set_pipeline(sg_pipeline{}); // 使用默认管线
		set_texture(sg_image{});     // 使用默认纹理
	}

	// 计算顶点位置
	Eigen::Matrix<float, 2, 4> positions;
	compute_rect_vertices(in_pos, in_size, positions, in_rotation_radians, in_scale, in_pivot);

	// 记录起始顶点索引
	const auto base_index = vertices_.size();

	// 添加顶点
	for (int i = 0; i < 4; ++i) {
		auto& vertex    = vertices_.emplace_back();
		vertex.position = positions.col(i);
		vertex.color    = in_color[i];
		vertex.uv       = in_uv[i];
		vertex.param_a  = in_param_a;
		vertex.param_b  = in_param_b;
		vertex.param_c  = in_param_c;
	}

	// 添加索引（两个三角形）
	indices_.push_back(base_index);
	indices_.push_back(base_index + 1);
	indices_.push_back(base_index + 2);

	indices_.push_back(base_index + 1);
	indices_.push_back(base_index + 3);
	indices_.push_back(base_index + 2);

	// 更新当前批次的计数
	draw_batch& batch  = batches_[current_batch_index_];
	batch.vertex_count = vertices_.size() - batch.vertex_start;
	batch.index_count  = indices_.size() - batch.index_start;

	// 更新统计信息
	total_triangles_ += 2;
}

// 创建矩形（公开接口）
void render_elements::make_rect(
	const Eigen::Vector2f&       in_pos,
	const Eigen::Vector2f&       in_size,
	const rect_color&            in_color,
	const geometry_t&            in_geometry,
	const mirage_vertex_param_t& in_param_a,
	const mirage_vertex_param_t& in_param_b,
	const mirage_vertex_param_t& in_param_c,
	const rect_uv&               in_uv,
	float                        in_rotation_radians,
	const Eigen::Vector2f&       in_pivot,
	const Eigen::Vector2f&       in_scale) {
	// 从几何体转换到窗口坐标
	const auto& pos = in_geometry.local_to_window(in_pos);
	add_rect_to_batch(pos,
	                  in_size,
	                  in_color,
	                  in_param_a,
	                  in_param_b,
	                  in_param_c,
	                  in_uv,
	                  in_rotation_radians,
	                  in_scale,
	                  in_pivot);
}

void render_elements::make_rounded_rect(const Eigen::Vector2f& in_pos, const Eigen::Vector2f& in_size,
	const geometry_t& in_geometry, const rect_color& in_color, const rect_round& in_round, float in_rotation_radians,
	const Eigen::Vector2f& in_pivot, const rect_uv& in_uv, const Eigen::Vector2f& in_scale) {
	set_pipeline(rounded_rect_pipeline_);
	const mirage_vertex_param_t param_a{ in_size };
	const mirage_vertex_param_t param_b{ in_round };
	make_rect(in_pos, in_size, in_color, in_geometry, param_a, param_b, {}, in_uv, in_rotation_radians, in_pivot, in_scale);
}

void render_elements::make_image(const Eigen::Vector2f& in_pos, const Eigen::Vector2f& in_size,
		const geometry_t& in_geometry, const sg_image& in_image, sampler_type in_sampler_type,
		const rect_color& in_color, const rect_uv& in_uv, float in_rotation_radians, const Eigen::Vector2f& in_pivot,
		const Eigen::Vector2f& in_scale) {
	const auto& sampler = texture_sampler_builder::get_sampler(in_sampler_type);
	make_image(in_pos,
	           in_size,
	           in_geometry,
	           in_image,
	           sampler,
	           in_color,
	           in_uv,
	           in_rotation_radians,
	           in_pivot,
	           in_scale);
}

void render_elements::make_image(const Eigen::Vector2f& in_pos, const Eigen::Vector2f& in_size,
		const geometry_t& in_geometry, const sg_image& in_image, const std::shared_ptr<texture_sampler>& in_sampler,
		const rect_color& in_color, const rect_uv& in_uv, float in_rotation_radians, const Eigen::Vector2f& in_pivot,
		const Eigen::Vector2f& in_scale) {
	if (!in_sampler) {
		return;
	}

	// 构建完整的批次键
	batch_key new_key;
	new_key.pipeline = image_pipeline_;
	new_key.image = in_image;
	new_key.sampler = *in_sampler;

	ensure_batch_compatibility(new_key);

	make_rect(in_pos, in_size, { in_color }, in_geometry, {}, {}, {}, in_uv, in_rotation_radians, in_pivot, in_scale);
}

void render_elements::make_text(const text_layout_t& in_layout, const Eigen::Vector2f& in_pos,
		const Eigen::Vector2f& in_size, const geometry_t& in_geometry, const rect_color& in_color,
		float in_rotation_radians, const Eigen::Vector2f& in_pivot, const Eigen::Vector2f& in_scale) {
	const auto& sampler = texture_sampler_builder::get_sampler(sampler_type::linear_clamp);

	for (const auto& position : in_layout.glyphs) {
		const auto& p      = position;
		const auto& region = p.region;
		const auto& size = p.size;
		const auto& uv     = region.uv_rect;

		if (auto texture = region.texture.lock()) {
			// 构建完整的批次键
			batch_key new_key;
			new_key.pipeline  = position.is_emoji ? image_pipeline_ : text_pipeline_;
			new_key.image     = texture->get_image();
			new_key.sampler   = *sampler;

			Eigen::Vector2f real_pos = p.position.array() + in_pos.array();
			// snap to pixel grid
			real_pos.x() = std::round(real_pos.x());
			real_pos.y() = std::round(real_pos.y());

			ensure_batch_compatibility(new_key);
			make_rect(real_pos, size, in_color, in_geometry, {}, {}, {}, uv, in_rotation_radians, in_pivot, in_scale);
			// make_wireframe(real_pos, size, in_geometry, 1.0f, in_color, in_rotation_radians, in_pivot, in_scale);
		}
	}
	Eigen::Vector2f font_ascent_start = in_pos;
	font_ascent_start.y() += in_layout.ascent;
	Eigen::Vector2f font_ascent_end = font_ascent_start;
	font_ascent_end.x() += in_layout.total_size.x();

	// 绘制文本的基线
	make_line(font_ascent_start, font_ascent_end, in_geometry, 1.0f, in_color, in_rotation_radians, in_pivot, in_scale);
	// 绘制文本的线框
	make_wireframe(in_pos, in_size, in_geometry, 1.0f, in_color, in_rotation_radians, in_pivot, in_scale);
}

void render_elements::make_wireframe(const Eigen::Vector2f& in_pos, const Eigen::Vector2f& in_size,
		const geometry_t& in_geometry, float in_line_width, const rect_color& in_color, float in_rotation_radians,
		const Eigen::Vector2f& in_pivot, const Eigen::Vector2f& in_scale) {
	set_pipeline(wireframe_pipeline_);
	mirage_vertex_param_t param_a{};
	param_a.x = in_line_width / in_size.x();
	param_a.y = in_line_width / in_size.y();
	make_rect(in_pos,
	          in_size,
	          { in_color },
	          in_geometry,
	          param_a,
	          {},
	          {},
	          {},
	          in_rotation_radians,
	          in_pivot,
	          in_scale);
}

void render_elements::make_line(const Eigen::Vector2f& in_start, const Eigen::Vector2f& in_end,
		const geometry_t& in_geometry, float in_thickness, const rect_color& in_color, float in_rotation_radians,
		const Eigen::Vector2f& in_pivot, const Eigen::Vector2f& in_scale) {
	set_pipeline(line_pipeline_);
	auto rect = rect_t<>::bounding_box<float>({in_start, in_end});

	if (rect.size().x() == 0 || rect.size().y() == 0) {
		// 根据线段方向添加厚度
		Eigen::Vector2f direction = (in_end - in_start).normalized();
		Eigen::Vector2f normal   = Eigen::Vector2f(-direction.y(), direction.x());
		rect.position() -= normal * in_thickness * 0.5f;
		rect.size() += normal * in_thickness;
	}

	mirage_vertex_param_t param_a{};
	mirage_vertex_param_t param_b{};
	param_a.x = in_thickness;

	// 需要将线段的起始和结束点转换为窗口坐标, 因为着色器中使用的是窗口坐标
	const auto& start_window = in_geometry.local_to_window(in_start);
	const auto& end_window   = in_geometry.local_to_window(in_end);

	param_b.x = start_window.x();
	param_b.y = start_window.y();
	param_b.z = end_window.x();
	param_b.w = end_window.y();

	make_rect(
		rect.position(),
		rect.size(),
		{ in_color },
		in_geometry,
		param_a,
		param_b,
		{},
		{},
		in_rotation_radians,
		in_pivot,
		in_scale
	);
}

// 确保缓冲区容量
void render_elements::ensure_buffer_capacity(uint32_t vertex_count, uint32_t index_count) {
	// 如果现有缓冲区不够大，重新创建
	if (vertex_count > vertex_buffer_capacity_) {
		// 销毁旧缓冲区
		if (vertex_buffer_.id != SG_INVALID_ID) { sg_destroy_buffer(vertex_buffer_); }

		// 创建新缓冲区，容量翻倍以减少重新分配
		vertex_buffer_capacity_   = vertex_count * 2;
		sg_buffer_desc vbuf_desc = {};
		vbuf_desc.size           = vertex_buffer_capacity_ * sizeof(mirage_vertex_t);
		vbuf_desc.usage          = SG_USAGE_STREAM;
		vertex_buffer_            = sg_make_buffer(&vbuf_desc);
	}

	if (index_count > index_buffer_capacity_) {
		if (index_buffer_.id != SG_INVALID_ID) { sg_destroy_buffer(index_buffer_); }

		index_buffer_capacity_    = index_count * 2;
		sg_buffer_desc ibuf_desc = {};
		ibuf_desc.size           = index_buffer_capacity_ * sizeof(uint32_t);
		ibuf_desc.usage          = SG_USAGE_STREAM;
		ibuf_desc.type           = SG_BUFFERTYPE_INDEXBUFFER;
		index_buffer_             = sg_make_buffer(&ibuf_desc);
	}
}

// 上传并绘制所有批次
void render_elements::flush_batches() {
	// 如果没有数据，直接返回
	if (vertices_.empty() || indices_.empty() || batches_.empty()) { return; }

	// 确保缓冲区足够大
	ensure_buffer_capacity(vertices_.size(), indices_.size());

	// 上传顶点和索引数据
	sg_update_buffer(vertex_buffer_, sg_range{ vertices_.data(), vertices_.size() * sizeof(mirage_vertex_t) });
	sg_update_buffer(index_buffer_, sg_range{ indices_.data(), indices_.size() * sizeof(uint32_t) });

	// 绑定顶点和索引缓冲区
	sg_bindings bind       = {};
	bind.vertex_buffers[0] = vertex_buffer_;
	bind.index_buffer      = index_buffer_;

	// 渲染每个批次
	for (const auto& batch: batches_) {
		// 设置纹理
		bind.images[0] = batch.key.image;
		bind.samplers[0] = batch.key.sampler;

		// 绑定管线
		sg_apply_pipeline(batch.key.pipeline);
		sg_apply_bindings(&bind);

		sg_apply_uniforms(0, SG_RANGE(projection_matrix_));

		// 绘制批次
		sg_draw(batch.index_start, batch.index_count, 1);
		draw_call_count_++;
	}
}

// 结束帧并提交所有批次
void render_elements::end_frame() { flush_batches(); }

// 析构函数
render_elements::~render_elements() {
	if (!sg_isvalid())
		return;
	if (vertex_buffer_.id != SG_INVALID_ID) { sg_destroy_buffer(vertex_buffer_); }
	if (index_buffer_.id != SG_INVALID_ID) { sg_destroy_buffer(index_buffer_); }
}

void render_elements::create_resources() {
	constexpr uint32_t default_vertex_buffer_size = 512; // 512kb
	constexpr uint32_t default_index_buffer_size  = 256; // 256kb
	// 计算默认顶点, 使用default_vertex_buffer_size对齐到sizeof(mirage_vertex_t)
	vertex_buffer_capacity_ = default_vertex_buffer_size * 1024 / sizeof(mirage_vertex_t);
	index_buffer_capacity_  = default_index_buffer_size * 1024 / sizeof(uint32_t);

	// 创建顶点缓冲区
	sg_buffer_desc vbuf_desc = {};
	vbuf_desc.size           = vertex_buffer_capacity_ * sizeof(mirage_vertex_t);
	vbuf_desc.usage          = SG_USAGE_STREAM;
	vbuf_desc.type           = SG_BUFFERTYPE_VERTEXBUFFER;
	vertex_buffer_            = sg_make_buffer(&vbuf_desc);

	// 创建索引缓冲区
	sg_buffer_desc ibuf_desc = {};
	ibuf_desc.size           = index_buffer_capacity_ * sizeof(mirage_triangle_t);
	ibuf_desc.usage          = SG_USAGE_STREAM;
	ibuf_desc.type           = SG_BUFFERTYPE_INDEXBUFFER;
	index_buffer_             = sg_make_buffer(&ibuf_desc);
}

void render_elements::load_mirage_pipelines() {
#if MIRAGE_USE_HDR
	auto format = MIRAGE_HDR_FORMAT;
#else
	auto format = MIRAGE_PIXEL_FORMAT;
#endif
	const auto rounded_rect_shader        = sg_make_shader(get_mirage_rounded_rect_shader_desc());
	const auto rounded_rect_pipeline_desc = get_mirage_rounded_rect_pipeline_desc(rounded_rect_shader, format, 1);
	rounded_rect_pipeline_ = sg_make_pipeline(rounded_rect_pipeline_desc);

	const auto image_shader        = sg_make_shader(get_mirage_image_shader_desc());
	const auto image_pipeline_desc = get_mirage_image_pipeline_desc(image_shader, format, 1);
	image_pipeline_          = sg_make_pipeline(image_pipeline_desc);

	const auto font_shader        = sg_make_shader(get_mirage_text_shader_desc());
	const auto font_pipeline_desc = get_mirage_text_pipeline_desc(font_shader, format, 1);
	text_pipeline_          = sg_make_pipeline(font_pipeline_desc);

	const auto wireframe_shader        = sg_make_shader(get_mirage_wireframe_shader_desc());
	const auto wireframe_pipeline_desc = get_mirage_wireframe_pipeline_desc(wireframe_shader, format, 1);
	wireframe_pipeline_ = sg_make_pipeline(wireframe_pipeline_desc);

	const auto line_shader        = sg_make_shader(get_mirage_line_shader_desc());
	const auto line_pipeline_desc = get_mirage_line_pipeline_desc(line_shader, format, 1);
	line_pipeline_ = sg_make_pipeline(line_pipeline_desc);
}
