#include "FaceApp.h"
#include "hardcode_data.h"
#include <sstream>

inline VkDescriptorSetLayoutBinding descriptor_set_layout_binding(
	VkDescriptorType   type,
	VkShaderStageFlags flags,
	uint32_t           binding,
	uint32_t           count = 1)
{
	VkDescriptorSetLayoutBinding set_layout_binding{};
	set_layout_binding.descriptorType = type;
	set_layout_binding.stageFlags = flags;
	set_layout_binding.binding = binding;
	set_layout_binding.descriptorCount = count;
	return set_layout_binding;
}

inline VkDescriptorSetLayoutCreateInfo descriptor_set_layout_create_info(
	const VkDescriptorSetLayoutBinding* bindings,
	uint32_t                            binding_count)
{
	VkDescriptorSetLayoutCreateInfo descriptor_set_layout_create_info{};
	descriptor_set_layout_create_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
	descriptor_set_layout_create_info.pBindings = bindings;
	descriptor_set_layout_create_info.bindingCount = binding_count;
	return descriptor_set_layout_create_info;
}

inline VkPipelineLayoutCreateInfo fun_pipeline_layout_create_info(
	const VkDescriptorSetLayout* set_layouts,
	uint32_t                     set_layout_count = 1)
{
	VkPipelineLayoutCreateInfo pipeline_layout_create_info{};
	pipeline_layout_create_info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
	pipeline_layout_create_info.setLayoutCount = set_layout_count;
	pipeline_layout_create_info.pSetLayouts = set_layouts;
	return pipeline_layout_create_info;
}

inline VkDescriptorSetAllocateInfo descriptor_set_allocate_info(
	VkDescriptorPool             descriptor_pool,
	const VkDescriptorSetLayout* set_layouts,
	uint32_t                     descriptor_set_count)
{
	VkDescriptorSetAllocateInfo descriptor_set_allocate_info{};
	descriptor_set_allocate_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
	descriptor_set_allocate_info.descriptorPool = descriptor_pool;
	descriptor_set_allocate_info.pSetLayouts = set_layouts;
	descriptor_set_allocate_info.descriptorSetCount = descriptor_set_count;
	return descriptor_set_allocate_info;
}

VkDescriptorBufferInfo create_descriptor(VkBuffer buffer, VkDeviceSize size = VK_WHOLE_SIZE, VkDeviceSize offset = 0)
{
	VkDescriptorBufferInfo descriptor{};
	descriptor.buffer = buffer;
	descriptor.range = size;
	descriptor.offset = offset;
	return descriptor;
}

inline VkWriteDescriptorSet write_descriptor_set(
	VkDescriptorSet         dst_set,
	VkDescriptorType        type,
	uint32_t                binding,
	VkDescriptorBufferInfo* buffer_info,
	uint32_t                descriptor_count = 1)
{
	VkWriteDescriptorSet write_descriptor_set{};
	write_descriptor_set.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
	write_descriptor_set.dstSet = dst_set;
	write_descriptor_set.descriptorType = type;
	write_descriptor_set.dstBinding = binding;
	write_descriptor_set.pBufferInfo = buffer_info;
	write_descriptor_set.descriptorCount = descriptor_count;
	return write_descriptor_set;
}


inline VkWriteDescriptorSet write_descriptor_set(
	VkDescriptorSet        dst_set,
	VkDescriptorType       type,
	uint32_t               binding,
	VkDescriptorImageInfo* image_info,
	uint32_t               descriptor_count = 1)
{
	VkWriteDescriptorSet write_descriptor_set{};
	write_descriptor_set.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
	write_descriptor_set.dstSet = dst_set;
	write_descriptor_set.descriptorType = type;
	write_descriptor_set.dstBinding = binding;
	write_descriptor_set.pImageInfo = image_info;
	write_descriptor_set.descriptorCount = descriptor_count;
	return write_descriptor_set;
}

FaceApp* FaceApp::faceIns = nullptr;

FaceApp::FaceApp(/* args */)
{
	faceIns = this;
}

FaceApp::~FaceApp()
{
}

bool FaceApp::LoadOBJ(const std::string& filename,
	std::vector<TextureLoadingVertexStructure>& vertices,
	std::vector<uint32_t>& indices) {

	// 临时存储从OBJ文件读取的原始数据
	std::vector<float> temp_positions;
	std::vector<float> temp_texcoords;
	std::vector<float> temp_normals;

	// 用于处理顶点索引
	std::vector<int> vertexIndices, uvIndices, normalIndices;

	std::vector<char> data = readFile(filename);

	// 将 vector<char> 转换为以 null 结尾的字符串（安全做法）
	std::string content(data.begin(), data.end());
	std::istringstream iss(content); // 用字符串创建字符串流

	std::string line;
	while (std::getline(iss, line)) {
		// 跳过空行和注释行
		if (line.empty() || line[0] == '#') {
			continue;
		}

		std::istringstream iss(line);
		std::string type;
		iss >> type;

		if (type == "v") { // 顶点位置
			float x, y, z;
			iss >> x >> y >> z;
			temp_positions.push_back(x);
			temp_positions.push_back(y);
			temp_positions.push_back(z);
		}
		else if (type == "vt") { // 纹理坐标
			float u, v;
			iss >> u >> v;
			temp_texcoords.push_back(u);
			temp_texcoords.push_back(1 - v);
		}
		else if (type == "vn") { // 法线
			float nx, ny, nz;
			iss >> nx >> ny >> nz;
			temp_normals.push_back(nx);
			temp_normals.push_back(ny);
			temp_normals.push_back(nz);
		}
		else if (type == "f") { // 面（三角形）
			std::string vertex1, vertex2, vertex3;
			iss >> vertex1 >> vertex2 >> vertex3;

			// 处理每个顶点的索引
			for (const std::string& vertex : { vertex1, vertex2, vertex3 }) {
				std::istringstream viss(vertex);
				std::string v, vt, vn;

				// 解析顶点索引格式：v/vt/vn 或 v//vn 或 v
				std::getline(viss, v, '/');
				std::getline(viss, vt, '/');
				std::getline(viss, vn, '/');

				int posIndex = std::stoi(v) - 1; // OBJ索引从1开始
				int texIndex = -1, normIndex = -1;

				if (!vt.empty()) texIndex = std::stoi(vt) - 1;
				if (!vn.empty()) normIndex = std::stoi(vn) - 1;

				vertexIndices.push_back(posIndex);
				uvIndices.push_back(texIndex);
				normalIndices.push_back(normIndex);
			}
		}
	}


	// 创建顶点数据
	vertices.clear();
	indices.clear();

	// 用于去重的哈希映射
	std::map<std::string, uint32_t> vertexMap;

	for (size_t i = 0; i < vertexIndices.size(); i++) {
		int posIndex = vertexIndices[i];
		int texIndex = uvIndices[i];
		int normIndex = normalIndices[i];

		// 创建唯一标识符
		std::string vertexKey = std::to_string(posIndex) + "/" +
			std::to_string(texIndex) + "/" +
			std::to_string(normIndex);

		// 检查是否已经存在相同的顶点
		if (vertexMap.find(vertexKey) != vertexMap.end()) {
			// 使用现有顶点的索引
			indices.push_back(vertexMap[vertexKey]);
		}
		else {
			// 创建新顶点
			TextureLoadingVertexStructure vertex;

			// 设置位置
			if (posIndex >= 0 && posIndex * 3 + 2 < temp_positions.size()) {
				vertex.pos[0] = temp_positions[posIndex * 3];
				vertex.pos[1] = temp_positions[posIndex * 3 + 1];
				vertex.pos[2] = temp_positions[posIndex * 3 + 2];
			}
			else {
				vertex.pos[0] = vertex.pos[1] = vertex.pos[2] = 0.0f;
			}

			// 设置纹理坐标
			if (texIndex >= 0 && texIndex * 2 + 1 < temp_texcoords.size()) {
				vertex.uv[0] = temp_texcoords[texIndex * 2];
				vertex.uv[1] = temp_texcoords[texIndex * 2 + 1];
			}
			else {
				vertex.uv[0] = vertex.uv[1] = 0.0f;
			}

			// 设置法线
			if (normIndex >= 0 && normIndex * 3 + 2 < temp_normals.size()) {
				vertex.normal[0] = temp_normals[normIndex * 3];
				vertex.normal[1] = temp_normals[normIndex * 3 + 1];
				vertex.normal[2] = temp_normals[normIndex * 3 + 2];
			}
			else {
				vertex.normal[0] = vertex.normal[1] = 0.0f;
				vertex.normal[2] = 1.0f; // 默认法线
			}

			// 添加新顶点并记录索引
			uint32_t newIndex = static_cast<uint32_t>(vertices.size());
			vertices.push_back(vertex);
			indices.push_back(newIndex);
			obj_vertices_map[newIndex] = posIndex;
			vertexMap[vertexKey] = newIndex;
		}
	}

	return true;
}

void FaceApp::setup_descriptor_pool()
{
	VkDescriptorPoolSize descriptor_pool_size{};
	descriptor_pool_size.type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
	descriptor_pool_size.descriptorCount = 6;

	VkDescriptorPoolSize descriptor_pool_image_size{};
	descriptor_pool_size.type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
	descriptor_pool_size.descriptorCount = 6;


	std::vector<VkDescriptorPoolSize> pool_sizes;
	pool_sizes.push_back(descriptor_pool_size);
	pool_sizes.push_back(descriptor_pool_image_size);


	VkDescriptorPoolCreateInfo descriptor_pool_info{};
	descriptor_pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
	descriptor_pool_info.poolSizeCount = pool_sizes.size();
	descriptor_pool_info.pPoolSizes = pool_sizes.data();
	descriptor_pool_info.maxSets = 6;

	VK_CHECK(vkCreateDescriptorPool(device, &descriptor_pool_info, nullptr, &descriptor_pool));
}



void FaceApp::setup_descriptor_set_layout()
{
	std::vector<VkDescriptorSetLayoutBinding> set_layout_bindings =
	{
		// Binding 0 : Vertex shader uniform buffer
		descriptor_set_layout_binding(
			VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
			VK_SHADER_STAGE_VERTEX_BIT,
			0),
		// Binding 1 : Fragment shader image sampler
		descriptor_set_layout_binding(
			VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
			VK_SHADER_STAGE_FRAGMENT_BIT,
			1) };

	VkDescriptorSetLayoutCreateInfo descriptor_layout =
		descriptor_set_layout_create_info(
			set_layout_bindings.data(),
			static_cast<uint32_t>(set_layout_bindings.size()));

	VK_CHECK(vkCreateDescriptorSetLayout(device, &descriptor_layout, nullptr, &m_descriptorSetLayout));


	VkPipelineLayoutCreateInfo pipeline_layout_create_info =
		fun_pipeline_layout_create_info(
			&m_descriptorSetLayout,
			1);

	VkPushConstantRange pushConstantRange{};
	pushConstantRange.stageFlags = VK_SHADER_STAGE_VERTEX_BIT; // 只在片段着色器中使用
	pushConstantRange.offset = 0;
	pushConstantRange.size = sizeof(float); // 或者 sizeof(PushConstants)

	pipeline_layout_create_info.pushConstantRangeCount = 1;
	pipeline_layout_create_info.pPushConstantRanges = &pushConstantRange;

	VK_CHECK(vkCreatePipelineLayout(device, &pipeline_layout_create_info, nullptr, &m_pipelineLayout));
}

void FaceApp::setup_descriptor_set()
{
	VkDescriptorSetAllocateInfo alloc_info =
		descriptor_set_allocate_info(
			descriptor_pool,
			&m_descriptorSetLayout,
			1);

	VK_CHECK(vkAllocateDescriptorSets(device, &alloc_info, &m_descriptor_set));

	VkDescriptorBufferInfo buffer_descriptor = create_descriptor(uniform_buffer_vs);


	VkDescriptorImageInfo image_descriptor;
	image_descriptor.imageView = tex_demo0.view;                // The image's view (images are never directly accessed by the shader, but rather through views defining subresources)
	image_descriptor.sampler = tex_demo0.sampler;             // The sampler (Telling the pipeline how to sample the texture, including repeat, border, etc.)
	image_descriptor.imageLayout = tex_demo0.image_layout;        // The current layout of the image (Note: Should always fit the actual use, e.g. shader read)

	std::vector<VkWriteDescriptorSet> write_descriptor_sets =
	{
		// Binding 0 : Vertex shader uniform buffer
		write_descriptor_set(
			m_descriptor_set,
			VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
			0,
			&buffer_descriptor),
		// Binding 1 : Fragment shader texture sampler
		//	Fragment shader: layout (binding = 1) uniform sampler2D samplerColor;
		write_descriptor_set(
			m_descriptor_set,
			VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,        // The descriptor set will use a combined image sampler (sampler and image could be split)
			1,                                                // Shader binding point 1
			&image_descriptor)                                // Pointer to the descriptor image for our texture
	};

	vkUpdateDescriptorSets(device, static_cast<uint32_t>(write_descriptor_sets.size()), write_descriptor_sets.data(), 0, NULL);
}

void FaceApp::render(VkCommandBuffer commandBuffer)
{
	Application::render(commandBuffer);
	vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_pipelineLayout, 0, 1, &m_descriptor_set, 0, NULL);
	vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_graphicsPipeline);

	VkDeviceSize offsets[1] = { 0 };
	VkBuffer vertexBuffers[] = { m_vertexBuffer };
	vkCmdBindVertexBuffers(commandBuffer, 0, 1, vertexBuffers, offsets);
	vkCmdBindIndexBuffer(commandBuffer, m_indexBuffer, 0, VK_INDEX_TYPE_UINT32);

#ifdef _WIN32
	vkCmdDrawIndexed(commandBuffer, obj_indices.size(), 1, 0, 0, 0);
#else
	if (getCurrentTimeMillis() - last_update_time < 1000)
	{
		vkCmdDrawIndexed(draw_cmd_buffers[i], index_count, 1, 0, 0, 0);
	}
#endif
}

void FaceApp::initVulkan()
{
	Application::initVulkan();
    VmaAllocatorCreateInfo allocatorInfo = {};
    allocatorInfo.physicalDevice = physicalDevice;
    allocatorInfo.device = device;
    allocatorInfo.instance = instance;
    vmaCreateAllocator(&allocatorInfo, &allocator);
    LoadOBJ("face_picture_3dmax.obj", obj_vertices, obj_indices);
	loadTexture("demo0.png", tex_demo0);
	createVertexBuffer();
	createUniformBuffer();
	setup_descriptor_pool();
	setup_descriptor_set_layout();
	setup_descriptor_set();
	uploadVertexData();
}

void FaceApp::update_uniform_buffers()
{
	uint32_t width = 480;
	uint32_t height = 480;
	float zoom = 2;
	// Vertex shader
	ubo_vs.projection = glm::perspective(glm::radians(60.0f), static_cast<float>(width) / static_cast<float>(height), 0.001f, 256.0f);
	glm::mat4 view_matrix = glm::translate(glm::mat4(1.0f), glm::vec3(0.0f, 0.0f, zoom));

	ubo_vs.model = view_matrix * glm::translate(glm::mat4(1.0f), camera_pos);
	ubo_vs.model = glm::rotate(ubo_vs.model, glm::radians(rotation.x), glm::vec3(1.0f, 0.0f, 0.0f));
	ubo_vs.model = glm::rotate(ubo_vs.model, glm::radians(rotation.y), glm::vec3(0.0f, 1.0f, 0.0f));
	ubo_vs.model = glm::rotate(ubo_vs.model, glm::radians(rotation.z), glm::vec3(0.0f, 0.0f, 1.0f));

	ubo_vs.view_pos = glm::vec4(0.0f, 0.0f, -zoom, 0.0f);

	memcpy(uniform_buffer_mapped, &ubo_vs, sizeof(ubo_vs));
}

void FaceApp::createVertexBuffer()
{
    VkDeviceSize vertexBufferSize = sizeof(TextureLoadingVertexStructure) * obj_vertices.capacity();
    VkDeviceSize indexBufferSize = sizeof(uint32_t) * obj_indices.capacity();

    // 创建顶点缓冲区（设备本地，用于渲染）
    VkBufferCreateInfo bufferInfo = { VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO };
    bufferInfo.size = vertexBufferSize;
    bufferInfo.usage = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT;
    bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

    VmaAllocationCreateInfo allocInfo = {};
    allocInfo.usage = VMA_MEMORY_USAGE_GPU_ONLY;

    vmaCreateBuffer(allocator, &bufferInfo, &allocInfo, &m_vertexBuffer, &m_vertexBufferAllocation, nullptr);

    // 创建暂存缓冲区（CPU可见，用于上传数据）
    bufferInfo.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
    allocInfo.usage = VMA_MEMORY_USAGE_CPU_TO_GPU;

    vmaCreateBuffer(allocator, &bufferInfo, &allocInfo, &m_stagingBuffer, &m_stagingBufferAllocation, nullptr);

    // 创建索引缓冲区
    bufferInfo.size = indexBufferSize;
    bufferInfo.usage = VK_BUFFER_USAGE_INDEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT;
    allocInfo.usage = VMA_MEMORY_USAGE_GPU_ONLY;

    vmaCreateBuffer(allocator, &bufferInfo, &allocInfo, &m_indexBuffer, &m_indexBufferAllocation, nullptr);
}

void ReceiveFacePoint(float* pos, int pointCount, int width, int height)
{
	FaceApp* self = FaceApp::Get();
	if (self != nullptr)
	{
		FaceApp::Get()->update_face_vertex_buffer(pos, pointCount);
	}
}

void FaceApp::update_face_vertex_buffer(float* pos, int pointCount)
{
	std::lock_guard<std::mutex> lock(mtx_point);
	last_update_time = getCurrentTimeMillis();


	for (int i = 0; i < obj_vertices.size(); ++i)
	{

		int face_index = obj_vertices_map[HardCodeData::Get().indexMap[i]];
		float x = pos[face_index * 3 + 0];
		float y = pos[face_index * 3 + 1];
		float z = pos[face_index * 3 + 2];

		obj_vertices[i].pos[0] = x;
		obj_vertices[i].pos[1] = y;
		obj_vertices[i].pos[2] = z;
	}

	uploadVertexData();
}

void FaceApp::uploadVertexData() {
    // 上传顶点数据
    void* data;
    vmaMapMemory(allocator, m_stagingBufferAllocation, &data);
    memcpy(data, obj_vertices.data(), sizeof(TextureLoadingVertexStructure) * obj_vertices.size());
    vmaUnmapMemory(allocator, m_stagingBufferAllocation);

    // 复制到设备内存
    copyBuffer(m_stagingBuffer, m_vertexBuffer, sizeof(TextureLoadingVertexStructure) * obj_vertices.size());

    // 上传索引数据（如果需要暂存缓冲区，可以创建另一个）
    vmaMapMemory(allocator, m_stagingBufferAllocation, &data);
    memcpy(data, obj_indices.data(), sizeof(uint32_t) * obj_indices.size());
    vmaUnmapMemory(allocator, m_stagingBufferAllocation);

    copyBuffer(m_stagingBuffer, m_indexBuffer, sizeof(uint32_t) * obj_indices.size());
}

void FaceApp::copyBuffer(VkBuffer srcBuffer, VkBuffer dstBuffer, VkDeviceSize size)
{
    VkCommandBuffer commandBuffer = beginSingleTimeCommands();

    VkBufferCopy copyRegion = {};
    copyRegion.size = size;
    vkCmdCopyBuffer(commandBuffer, srcBuffer, dstBuffer, 1, &copyRegion);

    endSingleTimeCommands(commandBuffer);
}

VkCommandBuffer FaceApp::beginSingleTimeCommands() {
    VkCommandBufferAllocateInfo allocInfo = { VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO };
    allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandPool = commandPool;
    allocInfo.commandBufferCount = 1;

    VkCommandBuffer commandBuffer;
    vkAllocateCommandBuffers(device, &allocInfo, &commandBuffer);

    VkCommandBufferBeginInfo beginInfo = { VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO };
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;

    vkBeginCommandBuffer(commandBuffer, &beginInfo);
    return commandBuffer;
}

void FaceApp::endSingleTimeCommands(VkCommandBuffer commandBuffer) {
    vkEndCommandBuffer(commandBuffer);

    VkSubmitInfo submitInfo = { VK_STRUCTURE_TYPE_SUBMIT_INFO };
    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &commandBuffer;

    vkQueueSubmit(graphicsQueue, 1, &submitInfo, VK_NULL_HANDLE);
    vkQueueWaitIdle(graphicsQueue);

    vkFreeCommandBuffers(device, commandPool, 1, &commandBuffer);
}

void FaceApp::createUniformBuffer()
{
	VkDeviceSize bufferSize = sizeof(ubo_vs);

	VkBufferCreateInfo bufferInfo = { VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO };
	bufferInfo.size = bufferSize;
	bufferInfo.usage = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT;
	bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

	VmaAllocationCreateInfo allocInfo = {};
	allocInfo.usage = VMA_MEMORY_USAGE_CPU_TO_GPU;
	allocInfo.requiredFlags = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;

	// 创建缓冲区和内存分配
	vmaCreateBuffer(allocator, &bufferInfo, &allocInfo,
		&uniform_buffer_vs,
		&uniform_buffer_allocation,
		nullptr);

	// 映射内存以便直接写入
	vmaMapMemory(allocator, uniform_buffer_allocation, &uniform_buffer_mapped);

	update_uniform_buffers();
}