#include "gpu_decompress.h"
#include "gpu_compress.h" // for BlockMeta
#include <string.h>

// ── Initialize decompression pipeline ──────────────────────────────
int gpu_decompress_init(GpuDecompressPipeline *pipe, GpuContext *ctx) {
  memset(pipe, 0, sizeof(GpuDecompressPipeline));
  pipe->ctx = ctx;

  // ── Descriptor set layout (same as compression) ────────────────
  VkDescriptorSetLayoutBinding bindings[3] = {
      {.binding = 0,
       .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
       .descriptorCount = 1,
       .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT},
      {.binding = 1,
       .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
       .descriptorCount = 1,
       .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT},
      {.binding = 2,
       .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
       .descriptorCount = 1,
       .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT},
  };

  VkDescriptorSetLayoutCreateInfo layout_info = {
      .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
      .bindingCount = 3,
      .pBindings = bindings,
  };

  VkResult res = vkCreateDescriptorSetLayout(ctx->device, &layout_info, NULL,
                                             &pipe->desc_layout);
  if (res != VK_SUCCESS)
    return -1;

  VkPushConstantRange push_range = {
      .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
      .offset = 0,
      .size = sizeof(DecompressPushConstants),
  };

  VkPipelineLayoutCreateInfo pipe_layout_info = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
      .setLayoutCount = 1,
      .pSetLayouts = &pipe->desc_layout,
      .pushConstantRangeCount = 1,
      .pPushConstantRanges = &push_range,
  };

  res = vkCreatePipelineLayout(ctx->device, &pipe_layout_info, NULL,
                               &pipe->pipeline_layout);
  if (res != VK_SUCCESS)
    return -1;

  VkComputePipelineCreateInfo pipeline_info = {
      .sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
      .stage =
          {
              .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
              .stage = VK_SHADER_STAGE_COMPUTE_BIT,
              .module = ctx->decompress_shader,
              .pName = "main",
          },
      .layout = pipe->pipeline_layout,
  };

  res = vkCreateComputePipelines(ctx->device, VK_NULL_HANDLE, 1, &pipeline_info,
                                 NULL, &pipe->pipeline);
  if (res != VK_SUCCESS) {
    LOG_ERR("Failed to create decompression compute pipeline (VkResult: %d)",
            res);
    return -1;
  }

  VkDescriptorPoolSize pool_size = {
      .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
      .descriptorCount = 3,
  };

  VkDescriptorPoolCreateInfo pool_info = {
      .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
      .maxSets = 1,
      .poolSizeCount = 1,
      .pPoolSizes = &pool_size,
  };

  res = vkCreateDescriptorPool(ctx->device, &pool_info, NULL, &pipe->desc_pool);
  if (res != VK_SUCCESS)
    return -1;

  VkDescriptorSetAllocateInfo alloc_info = {
      .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
      .descriptorPool = pipe->desc_pool,
      .descriptorSetCount = 1,
      .pSetLayouts = &pipe->desc_layout,
  };

  res = vkAllocateDescriptorSets(ctx->device, &alloc_info, &pipe->desc_set);
  if (res != VK_SUCCESS)
    return -1;

  LOG_OK("GPU decompression pipeline initialized");
  return 0;
}

// ── Decompress data using GPU ──────────────────────────────────────
int gpu_decompress_data(GpuDecompressPipeline *pipe, const VkzArchive *archive,
                        const char *archive_path, uint8_t *output_data,
                        uint64_t output_size) {

  GpuContext *ctx = pipe->ctx;
  uint32_t n_blocks = archive->header.block_count;
  uint32_t block_size = archive->header.block_size;

  // ── Read all compressed blocks into one buffer ─────────────────
  uint64_t total_compressed = 0;
  for (uint32_t i = 0; i < n_blocks; i++) {
    total_compressed += archive->blocks[i].compressed_size;
  }

  uint8_t *compressed_data =
      (uint8_t *)calloc(1, (size_t)total_compressed + 16);
  if (!compressed_data)
    return -1;

  // Read compressed blocks from archive file
  FILE *f = fopen(archive_path, "rb");
  if (!f) {
    free(compressed_data);
    return -1;
  }

  uint64_t comp_offset = 0;
  for (uint32_t i = 0; i < n_blocks; i++) {
    fseek(f, (long)archive->blocks[i].compressed_offset, SEEK_SET);
    fread(compressed_data + comp_offset, 1, archive->blocks[i].compressed_size,
          f);
    comp_offset += archive->blocks[i].compressed_size;
  }
  fclose(f);

  // ── Build block metadata ───────────────────────────────────────
  BlockMeta *block_metas = (BlockMeta *)calloc(n_blocks, sizeof(BlockMeta));
  uint32_t in_off = 0;
  uint32_t out_off = 0;

  for (uint32_t i = 0; i < n_blocks; i++) {
    block_metas[i].input_offset = in_off;
    block_metas[i].input_size = archive->blocks[i].compressed_size;
    block_metas[i].output_offset = out_off;
    block_metas[i].output_size = archive->blocks[i].original_size;

    in_off += archive->blocks[i].compressed_size;
    out_off += archive->blocks[i].original_size;
  }

  // ── Create GPU buffers ─────────────────────────────────────────
  VkDeviceSize in_buf_size = ((total_compressed + 3) / 4) * 4;
  VkDeviceSize out_buf_size = ((output_size + 3) / 4) * 4;
  VkDeviceSize meta_buf_size = n_blocks * sizeof(BlockMeta);

  VkMemoryPropertyFlags host_visible = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
                                       VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;

  if (gpu_create_buffer(ctx, in_buf_size, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
                        host_visible, &pipe->input_buffer,
                        &pipe->input_memory) != 0) {
    free(compressed_data);
    free(block_metas);
    return -1;
  }

  if (gpu_create_buffer(ctx, out_buf_size, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
                        host_visible, &pipe->output_buffer,
                        &pipe->output_memory) != 0) {
    free(compressed_data);
    free(block_metas);
    return -1;
  }

  if (gpu_create_buffer(ctx, meta_buf_size, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
                        host_visible, &pipe->meta_buffer,
                        &pipe->meta_memory) != 0) {
    free(compressed_data);
    free(block_metas);
    return -1;
  }

  // Upload data
  void *mapped;
  vkMapMemory(ctx->device, pipe->input_memory, 0, in_buf_size, 0, &mapped);
  memset(mapped, 0, (size_t)in_buf_size);
  memcpy(mapped, compressed_data, (size_t)total_compressed);
  vkUnmapMemory(ctx->device, pipe->input_memory);

  vkMapMemory(ctx->device, pipe->output_memory, 0, out_buf_size, 0, &mapped);
  memset(mapped, 0, (size_t)out_buf_size);
  vkUnmapMemory(ctx->device, pipe->output_memory);

  gpu_upload_data(ctx, pipe->meta_memory, block_metas, meta_buf_size);

  // ── Update descriptor set ──────────────────────────────────────
  VkDescriptorBufferInfo buffer_infos[3] = {
      {.buffer = pipe->input_buffer, .offset = 0, .range = in_buf_size},
      {.buffer = pipe->output_buffer, .offset = 0, .range = out_buf_size},
      {.buffer = pipe->meta_buffer, .offset = 0, .range = meta_buf_size},
  };

  VkWriteDescriptorSet writes[3];
  for (int i = 0; i < 3; i++) {
    writes[i] = (VkWriteDescriptorSet){
        .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .dstSet = pipe->desc_set,
        .dstBinding = (uint32_t)i,
        .descriptorCount = 1,
        .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .pBufferInfo = &buffer_infos[i],
    };
  }
  vkUpdateDescriptorSets(ctx->device, 3, writes, 0, NULL);

  // ── Dispatch compute ───────────────────────────────────────────
  VkCommandBufferAllocateInfo cmd_alloc = {
      .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
      .commandPool = ctx->command_pool,
      .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
      .commandBufferCount = 1,
  };

  VkCommandBuffer cmd;
  vkAllocateCommandBuffers(ctx->device, &cmd_alloc, &cmd);

  VkCommandBufferBeginInfo begin_info = {
      .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
      .flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
  };

  vkBeginCommandBuffer(cmd, &begin_info);

  vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, pipe->pipeline);
  vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE,
                          pipe->pipeline_layout, 0, 1, &pipe->desc_set, 0,
                          NULL);

  DecompressPushConstants pc = {
      .block_count = n_blocks,
      .block_size = block_size,
      ._pad1 = 0,
      ._pad2 = 0,
  };
  vkCmdPushConstants(cmd, pipe->pipeline_layout, VK_SHADER_STAGE_COMPUTE_BIT, 0,
                     sizeof(DecompressPushConstants), &pc);

  vkCmdDispatch(cmd, n_blocks, 1, 1);

  VkMemoryBarrier barrier = {
      .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER,
      .srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
      .dstAccessMask = VK_ACCESS_HOST_READ_BIT,
  };
  vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
                       VK_PIPELINE_STAGE_HOST_BIT, 0, 1, &barrier, 0, NULL, 0,
                       NULL);

  vkEndCommandBuffer(cmd);

  VkSubmitInfo submit = {
      .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
      .commandBufferCount = 1,
      .pCommandBuffers = &cmd,
  };

  VkFence fence;
  VkFenceCreateInfo fence_info = {.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO};
  vkCreateFence(ctx->device, &fence_info, NULL, &fence);

  vkQueueSubmit(ctx->compute_queue, 1, &submit, fence);
  vkWaitForFences(ctx->device, 1, &fence, VK_TRUE, UINT64_MAX);

  vkDestroyFence(ctx->device, fence, NULL);
  vkFreeCommandBuffers(ctx->device, ctx->command_pool, 1, &cmd);

  // ── Read back decompressed data ────────────────────────────────
  gpu_download_data(ctx, pipe->output_memory, output_data, output_size);

  // ── Cleanup ────────────────────────────────────────────────────
  vkDestroyBuffer(ctx->device, pipe->input_buffer, NULL);
  vkFreeMemory(ctx->device, pipe->input_memory, NULL);
  vkDestroyBuffer(ctx->device, pipe->output_buffer, NULL);
  vkFreeMemory(ctx->device, pipe->output_memory, NULL);
  vkDestroyBuffer(ctx->device, pipe->meta_buffer, NULL);
  vkFreeMemory(ctx->device, pipe->meta_memory, NULL);

  pipe->input_buffer = VK_NULL_HANDLE;
  pipe->input_memory = VK_NULL_HANDLE;
  pipe->output_buffer = VK_NULL_HANDLE;
  pipe->output_memory = VK_NULL_HANDLE;
  pipe->meta_buffer = VK_NULL_HANDLE;
  pipe->meta_memory = VK_NULL_HANDLE;

  free(compressed_data);
  free(block_metas);

  return 0;
}

// ── Cleanup ────────────────────────────────────────────────────────
void gpu_decompress_cleanup(GpuDecompressPipeline *pipe) {
  if (!pipe || !pipe->ctx)
    return;
  VkDevice dev = pipe->ctx->device;

  if (pipe->pipeline != VK_NULL_HANDLE)
    vkDestroyPipeline(dev, pipe->pipeline, NULL);
  if (pipe->pipeline_layout != VK_NULL_HANDLE)
    vkDestroyPipelineLayout(dev, pipe->pipeline_layout, NULL);
  if (pipe->desc_pool != VK_NULL_HANDLE)
    vkDestroyDescriptorPool(dev, pipe->desc_pool, NULL);
  if (pipe->desc_layout != VK_NULL_HANDLE)
    vkDestroyDescriptorSetLayout(dev, pipe->desc_layout, NULL);

  if (pipe->input_buffer != VK_NULL_HANDLE)
    vkDestroyBuffer(dev, pipe->input_buffer, NULL);
  if (pipe->input_memory != VK_NULL_HANDLE)
    vkFreeMemory(dev, pipe->input_memory, NULL);
  if (pipe->output_buffer != VK_NULL_HANDLE)
    vkDestroyBuffer(dev, pipe->output_buffer, NULL);
  if (pipe->output_memory != VK_NULL_HANDLE)
    vkFreeMemory(dev, pipe->output_memory, NULL);
  if (pipe->meta_buffer != VK_NULL_HANDLE)
    vkDestroyBuffer(dev, pipe->meta_buffer, NULL);
  if (pipe->meta_memory != VK_NULL_HANDLE)
    vkFreeMemory(dev, pipe->meta_memory, NULL);
}
