#include "render_elements.h"
#include "shaders/mirage_rounded_rect.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);
}

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 / in_size.x();
	const float scale_y = -2.0f / 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_pivot,
	const Eigen::Vector2f&       in_scale) {
	// 确保有活跃的批次
	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_pivot, in_scale);

	// 记录起始顶点索引
	uint32_t 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_pivot,
	                  in_scale);
}

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_scale, in_pivot);
}

// 确保缓冲区容量
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;

		// 绑定管线
		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
	auto rounded_rect_shader        = sg_make_shader(get_mirage_rounded_rect_shader_desc());
	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);
}