#include "gpu_context.h"

// ── Initialize Vulkan ──────────────────────────────────────────────
int gpu_init(GpuContext *ctx) {
  memset(ctx, 0, sizeof(GpuContext));

#ifdef _WIN32
  if (volkInitialize() != VK_SUCCESS) {
    LOG_ERR("Failed to initialize Volk (Vulkan loader). Is Vulkan driver "
            "installed?");
    return -1;
  }
#endif

  // ── Create Vulkan instance ─────────────────────────────────────
  VkApplicationInfo app_info = {
      .sType = VK_STRUCTURE_TYPE_APPLICATION_INFO,
      .pApplicationName = "VKZip",
      .applicationVersion = VK_MAKE_VERSION(1, 0, 0),
      .pEngineName = "VKZip Engine",
      .engineVersion = VK_MAKE_VERSION(1, 0, 0),
      .apiVersion = VK_API_VERSION_1_2,
  };

  VkInstanceCreateInfo instance_info = {
      .sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
      .pApplicationInfo = &app_info,
  };

  VkResult res = vkCreateInstance(&instance_info, NULL, &ctx->instance);
  if (res != VK_SUCCESS) {
    LOG_ERR("Failed to create Vulkan instance (VkResult: %d)", res);
    LOG_ERR("Make sure Vulkan drivers are installed!");
    return -1;
  }

#ifdef _WIN32
  volkLoadInstance(ctx->instance);
#endif

  // ── Select physical device (prefer discrete GPU) ───────────────
  uint32_t device_count = 0;
  vkEnumeratePhysicalDevices(ctx->instance, &device_count, NULL);
  if (device_count == 0) {
    LOG_ERR("No Vulkan-capable GPU found!");
    return -1;
  }

  VkPhysicalDevice *devices =
      (VkPhysicalDevice *)malloc(device_count * sizeof(VkPhysicalDevice));
  vkEnumeratePhysicalDevices(ctx->instance, &device_count, devices);

  // Score and pick best device
  int best_score = -1;
  uint32_t best_idx = 0;

  for (uint32_t i = 0; i < device_count; i++) {
    VkPhysicalDeviceProperties props;
    vkGetPhysicalDeviceProperties(devices[i], &props);

    int score = 0;
    if (props.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU)
      score += 1000;
    else if (props.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU)
      score += 100;

    // Prefer more compute units / higher limits
    score += props.limits.maxComputeWorkGroupInvocations / 32;

    if (score > best_score) {
      best_score = score;
      best_idx = i;
    }
  }

  ctx->physical_device = devices[best_idx];
  free(devices);

  vkGetPhysicalDeviceProperties(ctx->physical_device, &ctx->device_props);
  vkGetPhysicalDeviceMemoryProperties(ctx->physical_device, &ctx->mem_props);

  strncpy(ctx->device_name, ctx->device_props.deviceName,
          sizeof(ctx->device_name) - 1);
  ctx->max_workgroup_size =
      ctx->device_props.limits.maxComputeWorkGroupInvocations;

  // Calculate VRAM
  ctx->vram_size = 0;
  for (uint32_t i = 0; i < ctx->mem_props.memoryHeapCount; i++) {
    if (ctx->mem_props.memoryHeaps[i].flags & VK_MEMORY_HEAP_DEVICE_LOCAL_BIT) {
      ctx->vram_size += ctx->mem_props.memoryHeaps[i].size;
    }
  }

  // ── Find compute queue family ──────────────────────────────────
  uint32_t queue_family_count = 0;
  vkGetPhysicalDeviceQueueFamilyProperties(ctx->physical_device,
                                           &queue_family_count, NULL);
  VkQueueFamilyProperties *queue_families = (VkQueueFamilyProperties *)malloc(
      queue_family_count * sizeof(VkQueueFamilyProperties));
  vkGetPhysicalDeviceQueueFamilyProperties(ctx->physical_device,
                                           &queue_family_count, queue_families);

  ctx->compute_queue_family = UINT32_MAX;

  // Prefer a dedicated compute queue (no graphics)
  for (uint32_t i = 0; i < queue_family_count; i++) {
    if ((queue_families[i].queueFlags & VK_QUEUE_COMPUTE_BIT) &&
        !(queue_families[i].queueFlags & VK_QUEUE_GRAPHICS_BIT)) {
      ctx->compute_queue_family = i;
      break;
    }
  }

  // Fallback to any compute queue
  if (ctx->compute_queue_family == UINT32_MAX) {
    for (uint32_t i = 0; i < queue_family_count; i++) {
      if (queue_families[i].queueFlags & VK_QUEUE_COMPUTE_BIT) {
        ctx->compute_queue_family = i;
        break;
      }
    }
  }

  free(queue_families);

  if (ctx->compute_queue_family == UINT32_MAX) {
    LOG_ERR("No compute queue family found on GPU!");
    return -1;
  }

  // ── Create logical device ──────────────────────────────────────
  float queue_priority = 1.0f;
  VkDeviceQueueCreateInfo queue_info = {
      .sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
      .queueFamilyIndex = ctx->compute_queue_family,
      .queueCount = 1,
      .pQueuePriorities = &queue_priority,
  };

  VkDeviceCreateInfo device_info = {
      .sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
      .queueCreateInfoCount = 1,
      .pQueueCreateInfos = &queue_info,
  };

  res = vkCreateDevice(ctx->physical_device, &device_info, NULL, &ctx->device);
  if (res != VK_SUCCESS) {
    LOG_ERR("Failed to create Vulkan device (VkResult: %d)", res);
    return -1;
  }

#ifdef _WIN32
  volkLoadDevice(ctx->device);
#endif

  vkGetDeviceQueue(ctx->device, ctx->compute_queue_family, 0,
                   &ctx->compute_queue);

  // ── Create command pool ────────────────────────────────────────
  VkCommandPoolCreateInfo pool_info = {
      .sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
      .queueFamilyIndex = ctx->compute_queue_family,
      .flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
  };

  res = vkCreateCommandPool(ctx->device, &pool_info, NULL, &ctx->command_pool);
  if (res != VK_SUCCESS) {
    LOG_ERR("Failed to create command pool (VkResult: %d)", res);
    return -1;
  }

  ctx->is_initialized = true;
  return 0;
}

// ── Load SPIR-V shader ─────────────────────────────────────────────
VkShaderModule gpu_load_shader(GpuContext *ctx, const char *spv_path) {
  FILE *f = fopen(spv_path, "rb");
  if (!f) {
    LOG_ERR("Cannot open shader: %s", spv_path);
    return VK_NULL_HANDLE;
  }

  fseek(f, 0, SEEK_END);
  size_t size = (size_t)ftell(f);
  fseek(f, 0, SEEK_SET);

  uint32_t *code = (uint32_t *)malloc(size);
  if (!code) {
    fclose(f);
    return VK_NULL_HANDLE;
  }

  fread(code, 1, size, f);
  fclose(f);

  VkShaderModuleCreateInfo create_info = {
      .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
      .codeSize = size,
      .pCode = code,
  };

  VkShaderModule shader;
  VkResult res = vkCreateShaderModule(ctx->device, &create_info, NULL, &shader);
  free(code);

  if (res != VK_SUCCESS) {
    LOG_ERR("Failed to create shader module from %s (VkResult: %d)", spv_path,
            res);
    return VK_NULL_HANDLE;
  }

  return shader;
}

// ── Find memory type ───────────────────────────────────────────────
uint32_t gpu_find_memory_type(GpuContext *ctx, uint32_t type_bits,
                              VkMemoryPropertyFlags properties) {
  for (uint32_t i = 0; i < ctx->mem_props.memoryTypeCount; i++) {
    if ((type_bits & (1 << i)) && (ctx->mem_props.memoryTypes[i].propertyFlags &
                                   properties) == properties) {
      return i;
    }
  }
  return UINT32_MAX;
}

// ── Create buffer ──────────────────────────────────────────────────
int gpu_create_buffer(GpuContext *ctx, VkDeviceSize size,
                      VkBufferUsageFlags usage,
                      VkMemoryPropertyFlags mem_props_flags, VkBuffer *buffer,
                      VkDeviceMemory *memory) {
  VkBufferCreateInfo buffer_info = {
      .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
      .size = size,
      .usage = usage,
      .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };

  VkResult res = vkCreateBuffer(ctx->device, &buffer_info, NULL, buffer);
  if (res != VK_SUCCESS)
    return -1;

  VkMemoryRequirements mem_req;
  vkGetBufferMemoryRequirements(ctx->device, *buffer, &mem_req);

  uint32_t mem_type =
      gpu_find_memory_type(ctx, mem_req.memoryTypeBits, mem_props_flags);
  if (mem_type == UINT32_MAX) {
    vkDestroyBuffer(ctx->device, *buffer, NULL);
    return -1;
  }

  VkMemoryAllocateInfo alloc_info = {
      .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
      .allocationSize = mem_req.size,
      .memoryTypeIndex = mem_type,
  };

  res = vkAllocateMemory(ctx->device, &alloc_info, NULL, memory);
  if (res != VK_SUCCESS) {
    vkDestroyBuffer(ctx->device, *buffer, NULL);
    return -1;
  }

  vkBindBufferMemory(ctx->device, *buffer, *memory, 0);
  return 0;
}

// ── Upload data to GPU ─────────────────────────────────────────────
int gpu_upload_data(GpuContext *ctx, VkDeviceMemory memory, const void *data,
                    VkDeviceSize size) {
  void *mapped;
  VkResult res = vkMapMemory(ctx->device, memory, 0, size, 0, &mapped);
  if (res != VK_SUCCESS)
    return -1;
  memcpy(mapped, data, (size_t)size);
  vkUnmapMemory(ctx->device, memory);
  return 0;
}

// ── Download data from GPU ─────────────────────────────────────────
int gpu_download_data(GpuContext *ctx, VkDeviceMemory memory, void *data,
                      VkDeviceSize size) {
  void *mapped;
  VkResult res = vkMapMemory(ctx->device, memory, 0, size, 0, &mapped);
  if (res != VK_SUCCESS)
    return -1;
  memcpy(data, mapped, (size_t)size);
  vkUnmapMemory(ctx->device, memory);
  return 0;
}

// ── Print GPU info ─────────────────────────────────────────────────
void gpu_print_info(const GpuContext *ctx) {
  char vram_buf[32];
  format_size(ctx->vram_size, vram_buf, sizeof(vram_buf));

  const char *type_str = "Unknown";
  switch (ctx->device_props.deviceType) {
  case VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU:
    type_str = "Discrete GPU";
    break;
  case VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU:
    type_str = "Integrated GPU";
    break;
  case VK_PHYSICAL_DEVICE_TYPE_VIRTUAL_GPU:
    type_str = "Virtual GPU";
    break;
  case VK_PHYSICAL_DEVICE_TYPE_CPU:
    type_str = "CPU";
    break;
  default:
    break;
  }

  printf("\n");
  printf(COL_BOLD COL_MAGENTA
         "╔══════════════════════════════════════════╗\n" COL_RESET);
  printf(COL_BOLD COL_MAGENTA
         "║            GPU Information               ║\n" COL_RESET);
  printf(COL_BOLD COL_MAGENTA
         "╠══════════════════════════════════════════╣\n" COL_RESET);
  printf(COL_MAGENTA "║" COL_RESET " Device:     %-28s" COL_MAGENTA
                     "║\n" COL_RESET,
         ctx->device_name);
  printf(COL_MAGENTA "║" COL_RESET " Type:       %-28s" COL_MAGENTA
                     "║\n" COL_RESET,
         type_str);
  printf(COL_MAGENTA "║" COL_RESET " VRAM:       %-28s" COL_MAGENTA
                     "║\n" COL_RESET,
         vram_buf);
  printf(COL_MAGENTA "║" COL_RESET " Max WG:     %-28u" COL_MAGENTA
                     "║\n" COL_RESET,
         ctx->max_workgroup_size);
  printf(COL_MAGENTA "║" COL_RESET " Vulkan API: %u.%u.%u%*s" COL_MAGENTA
                     "║\n" COL_RESET,
         VK_API_VERSION_MAJOR(ctx->device_props.apiVersion),
         VK_API_VERSION_MINOR(ctx->device_props.apiVersion),
         VK_API_VERSION_PATCH(ctx->device_props.apiVersion), 21, "");
  printf(COL_BOLD COL_MAGENTA
         "╚══════════════════════════════════════════╝\n" COL_RESET);
  printf("\n");
}

// ── Cleanup ────────────────────────────────────────────────────────
void gpu_cleanup(GpuContext *ctx) {
  if (!ctx || !ctx->is_initialized)
    return;

  if (ctx->compress_shader != VK_NULL_HANDLE)
    vkDestroyShaderModule(ctx->device, ctx->compress_shader, NULL);
  if (ctx->decompress_shader != VK_NULL_HANDLE)
    vkDestroyShaderModule(ctx->device, ctx->decompress_shader, NULL);

  vkDestroyCommandPool(ctx->device, ctx->command_pool, NULL);
  vkDestroyDevice(ctx->device, NULL);
  vkDestroyInstance(ctx->instance, NULL);

  ctx->is_initialized = false;
}
