/*
 * ╔══════════════════════════════════════════════════════════════════╗
 * ║                     VKZip v1.0                                  ║
 * ║         GPU-Accelerated File Compressor/Decompressor            ║
 * ║              Powered by Vulkan Compute Shaders                  ║
 * ╚══════════════════════════════════════════════════════════════════╝
 *
 * Usage:
 *   vkzip compress   <input_file> [output.vkz]
 *   vkzip decompress <archive.vkz> [output_file]
 *   vkzip info       <archive.vkz>
 *   vkzip benchmark  <input_file>
 *   vkzip --help
 *
 * Cross-platform: Linux & Windows
 * GPU Support: Any Vulkan-capable GPU (NVIDIA, AMD, Intel)
 * Fallback: CPU compression when no GPU is available
 */

#include "cpu_fallback.h"
#include "gpu_compress.h"
#include "gpu_context.h"
#include "gpu_decompress.h"
#include "utils.h"
#include "vkz_format.h"
#include <errno.h>
#include <sys/stat.h>

#ifdef _WIN32
#include <io.h>
#include <process.h>
#define F_OK 0
#define W_OK 2
#define R_OK 4
#define access _access
#define unlink _unlink
#define getpid _getpid
#else
#include <unistd.h>
#endif

// ── Banner ─────────────────────────────────────────────────────────
static void print_banner(void) {
  printf("\n");
  printf(COL_BOLD COL_GREEN
         "  ██╗   ██╗██╗  ██╗███████╗██╗██████╗ \n"
         "  ██║   ██║██║ ██╔╝╚══███╔╝██║██╔══██╗\n"
         "  ██║   ██║█████╔╝   ███╔╝ ██║██████╔╝\n"
         "  ╚██╗ ██╔╝██╔═██╗  ███╔╝  ██║██╔═══╝ \n"
         "   ╚████╔╝ ██║  ██╗███████╗██║██║     \n"
         "    ╚═══╝  ╚═╝  ╚═╝╚══════╝╚═╝╚═╝     \n" COL_RESET);
  printf(COL_CYAN "  GPU-Accelerated File Compressor v1.0\n" COL_RESET);
  printf(COL_CYAN "  Powered by Vulkan Compute Shaders\n" COL_RESET);
  printf("\n");
}

// ── Help ───────────────────────────────────────────────────────────
static void print_help(void) {
  print_banner();
  printf("Usage:\n");
  printf("  vkzip " COL_GREEN "compress" COL_RESET
         "   <input_file> [output.vkz]   Compress a file\n");
  printf("  vkzip " COL_GREEN "decompress" COL_RESET
         " <archive.vkz> [output_file] Decompress archive\n");
  printf("  vkzip " COL_GREEN "info" COL_RESET
         "       <archive.vkz>               Show archive info\n");
  printf("  vkzip " COL_GREEN "benchmark" COL_RESET
         "  <input_file>                Benchmark GPU vs CPU\n");
  printf("  vkzip " COL_GREEN "--gpu-info" COL_RESET
         "                              Show GPU information\n");
  printf("  vkzip " COL_GREEN "--help" COL_RESET
         "                                  Show this help\n");
  printf("\n");
  printf("Options:\n");
  printf("  --cpu-only    Force CPU mode (no GPU)\n");
  printf("  --block-size  Block size in KB (default: 64)\n");
  printf("\n");
  printf("Examples:\n");
  printf("  vkzip compress myfile.bin\n");
  printf("  vkzip compress myfile.bin output.vkz\n");
  printf("  vkzip decompress output.vkz\n");
  printf("  vkzip decompress output.vkz restored.bin\n");
  printf("  vkzip benchmark largefile.dat\n");
  printf("\n");
}

// ── Find shader path ───────────────────────────────────────────────
static int find_shader_path(char *out_path, size_t max_len,
                            const char *shader_name) {
  // Try build directory first (defined by CMake)
#ifdef SHADER_DIR
  snprintf(out_path, max_len, "%s/%s", SHADER_DIR, shader_name);
  if (access(out_path, R_OK) == 0)
    return 0;
#endif

  // Try relative to executable
  snprintf(out_path, max_len, "shaders/%s", shader_name);
  if (access(out_path, R_OK) == 0)
    return 0;

  // Try installed location
  snprintf(out_path, max_len, "/usr/share/vkzip/shaders/%s", shader_name);
  if (access(out_path, R_OK) == 0)
    return 0;

  // Try local share
  snprintf(out_path, max_len, "/usr/local/share/vkzip/shaders/%s", shader_name);
  if (access(out_path, R_OK) == 0)
    return 0;

  return -1;
}

// ── Check if path is a directory ───────────────────────────────────
static bool is_directory(const char *path) {
  struct stat st;
  if (stat(path, &st) != 0)
    return false;
  return S_ISDIR(st.st_mode);
}

// ── Strip trailing slashes ─────────────────────────────────────────
static void strip_trailing_slash(char *path) {
  size_t len = strlen(path);
  while (len > 1 && (path[len - 1] == '/' || path[len - 1] == '\\')) {
    path[len - 1] = '\0';
    len--;
  }
}

// ── Check file extension ───────────────────────────────────────────
static bool ends_with(const char *str, const char *suffix) {
  if (!str || !suffix)
    return false;
  size_t lenstr = strlen(str);
  size_t lensuffix = strlen(suffix);
  if (lensuffix > lenstr)
    return false;
  return strcasecmp(str + lenstr - lensuffix, suffix) == 0;
}

// ── Extract external archives via system ───────────────────────────
static int extract_external_archive(const char *input_path,
                                    const char *output_dir, const char *type) {
  char cmd[VKZ_MAX_FILENAME * 3];

  // Create output dir if specified
  if (output_dir) {
#ifdef _WIN32
    snprintf(cmd, sizeof(cmd), "if not exist \"%s\" mkdir \"%s\"", output_dir,
             output_dir);
#else
    snprintf(cmd, sizeof(cmd), "mkdir -p '%s'", output_dir);
#endif
    system(cmd);
  }

  const char *out_target = output_dir ? output_dir : ".";

  if (strcmp(type, "zip") == 0) {
    LOG_INFO("Extracting ZIP archive using system tools...");
#ifdef _WIN32
    snprintf(cmd, sizeof(cmd), "tar -xf \"%s\" -C \"%s\"", input_path,
             out_target);
#else
    snprintf(cmd, sizeof(cmd), "unzip -q -o '%s' -d '%s'", input_path,
             out_target);
#endif
  } else if (strcmp(type, "rar") == 0) {
    LOG_INFO("Extracting RAR archive using system tools...");
#ifdef _WIN32
    snprintf(cmd, sizeof(cmd), "unrar x -y \"%s\" \"%s\\\"", input_path,
             out_target);
#else
    snprintf(cmd, sizeof(cmd), "unrar x -y '%s' '%s/'", input_path, out_target);
#endif
  } else {
    return -1;
  }

  Timer timer;
  timer_start(&timer);
  if (system(cmd) == 0) {
    double time = timer_stop(&timer);
    LOG_OK("Extraction complete! (%.3f seconds)", time);
    return 0;
  } else {
    LOG_ERR("Extraction failed. Do you have %s installed?", type);
    return 1;
  }
}

// ── Create tar from directory ──────────────────────────────────────
static char *tar_directory(const char *dir_path) {
  static char tar_path[VKZ_MAX_FILENAME];
  snprintf(tar_path, sizeof(tar_path), "/tmp/vkzip_%u.tar", (uint32_t)getpid());

  // Get parent dir and basename
  char dir_copy[VKZ_MAX_FILENAME];
  strncpy(dir_copy, dir_path, sizeof(dir_copy) - 1);
  strip_trailing_slash(dir_copy);

  char *last_slash = strrchr(dir_copy, '/');
  char parent[VKZ_MAX_FILENAME] = ".";
  const char *basename = dir_copy;

  if (last_slash) {
    *last_slash = '\0';
    strncpy(parent, dir_copy, sizeof(parent) - 1);
    basename = last_slash + 1;
  }

  char cmd[VKZ_MAX_FILENAME * 3];
  snprintf(cmd, sizeof(cmd), "tar cf '%s' -C '%s' '%s' 2>/dev/null", tar_path,
           parent, basename);

  LOG_INFO("Packing directory with tar...");
  int ret = system(cmd);
  if (ret != 0) {
    LOG_ERR("Failed to create tar archive from directory");
    return NULL;
  }

  char size_buf[32];
  format_size(get_file_size(tar_path), size_buf, sizeof(size_buf));
  LOG_OK("Directory packed: %s", size_buf);

  return tar_path;
}

// ── Extract tar to directory ───────────────────────────────────────
static int untar_file(const char *tar_path, const char *output_dir) {
  // Create output directory if needed
  char cmd[VKZ_MAX_FILENAME * 3];

#ifdef _WIN32
  if (output_dir) {
    snprintf(cmd, sizeof(cmd),
             "if not exist \"%s\" mkdir \"%s\" && tar xf \"%s\" -C \"%s\"",
             output_dir, output_dir, tar_path, output_dir);
  } else {
    snprintf(cmd, sizeof(cmd), "tar xf \"%s\"", tar_path);
  }
#else
  if (output_dir) {
    snprintf(cmd, sizeof(cmd),
             "mkdir -p '%s' && tar xf '%s' -C '%s' 2>/dev/null", output_dir,
             tar_path, output_dir);
  } else {
    snprintf(cmd, sizeof(cmd), "tar xf '%s' 2>/dev/null", tar_path);
  }
#endif

  int ret = system(cmd);
  return (ret == 0) ? 0 : -1;
}

// ── Generate output filename ───────────────────────────────────────
static void make_compress_output(const char *input, char *output,
                                 size_t max_len) {
  // Strip trailing slash for directories
  char clean[VKZ_MAX_FILENAME];
  strncpy(clean, input, sizeof(clean) - 1);
  strip_trailing_slash(clean);
  snprintf(output, max_len, "%s.vkz", clean);
}

static void make_decompress_output(const VkzArchive *archive, char *output,
                                   size_t max_len) {
  if (archive->original_filename[0] != '\0') {
    snprintf(output, max_len, "%s", archive->original_filename);
  } else {
    snprintf(output, max_len, "output.bin");
  }
}

// ── Compress command ───────────────────────────────────────────────
static int cmd_compress(const char *input_path, const char *output_path,
                        bool cpu_only, uint32_t block_size_kb) {
  Timer total_timer;
  timer_start(&total_timer);

  // Check if input is a directory
  bool is_dir = is_directory(input_path);
  const char *actual_input = input_path;
  char *tar_temp_path = NULL;

  if (is_dir) {
    LOG_INFO("Input is a directory: %s", input_path);
    tar_temp_path = tar_directory(input_path);
    if (!tar_temp_path)
      return 1;
    actual_input = tar_temp_path;
  }

  // Read input file
  uint64_t input_size = get_file_size(actual_input);
  if (input_size == 0) {
    LOG_ERR("Cannot read file or file is empty: %s", actual_input);
    if (tar_temp_path)
      unlink(tar_temp_path);
    return 1;
  }

  char size_buf[32];
  format_size(input_size, size_buf, sizeof(size_buf));
  LOG_INFO("Input: %s (%s)", input_path, size_buf);
  if (is_dir)
    LOG_INFO("Mode: directory (tar + compress)");

  FILE *f = fopen(actual_input, "rb");
  if (!f) {
    LOG_ERR("Cannot open input file: %s", input_path);
    return 1;
  }

  uint8_t *input_data = (uint8_t *)malloc((size_t)input_size);
  if (!input_data) {
    LOG_ERR("Out of memory (need %s)", size_buf);
    fclose(f);
    if (tar_temp_path)
      unlink(tar_temp_path);
    return 1;
  }
  fread(input_data, 1, (size_t)input_size, f);
  fclose(f);

  uint32_t block_size = block_size_kb * 1024;

  // Auto-generate output filename if needed
  char auto_output[VKZ_MAX_FILENAME];
  if (!output_path) {
    make_compress_output(input_path, auto_output, sizeof(auto_output));
    output_path = auto_output;
  }

  LOG_INFO("Output: %s", output_path);
  LOG_INFO("Block size: %u KB", block_size_kb);

  // ── Try GPU compression ────────────────────────────────────────
  uint8_t **compressed_blocks = NULL;
  uint32_t *block_sizes = NULL;
  uint32_t block_count = 0;
  bool used_gpu = false;

  if (!cpu_only) {
    GpuContext gpu;
    if (gpu_init(&gpu) == 0) {
      gpu_print_info(&gpu);

      // Load compression shader
      char shader_path[1024];
      if (find_shader_path(shader_path, sizeof(shader_path),
                           "compress.comp.spv") == 0) {
        gpu.compress_shader = gpu_load_shader(&gpu, shader_path);

        if (gpu.compress_shader != VK_NULL_HANDLE) {
          GpuCompressPipeline pipe;
          if (gpu_compress_init(&pipe, &gpu) == 0) {
            LOG_INFO("Compressing with GPU...");
            Timer gpu_timer;
            timer_start(&gpu_timer);

            int result = gpu_compress_data(&pipe, input_data, input_size,
                                           block_size, &compressed_blocks,
                                           &block_sizes, &block_count);

            double gpu_time = timer_stop(&gpu_timer);

            if (result == 0) {
              used_gpu = true;
              LOG_OK("GPU compression done in %.3f seconds", gpu_time);

              double throughput =
                  (double)input_size / (1024.0 * 1024.0) / gpu_time;
              LOG_OK("Throughput: %.1f MB/s", throughput);
            } else {
              LOG_WARN("GPU compression failed, falling back to CPU");
            }

            gpu_compress_cleanup(&pipe);
          }
        } else {
          LOG_WARN("Could not load compression shader");
        }
      } else {
        LOG_WARN("Shader file not found, falling back to CPU");
      }

      gpu_cleanup(&gpu);
    } else {
      LOG_WARN("No Vulkan GPU available, using CPU fallback");
    }
  }

  // ── CPU fallback ───────────────────────────────────────────────
  if (!used_gpu) {
    LOG_INFO("Compressing with CPU...");
    Timer cpu_timer;
    timer_start(&cpu_timer);

    int result =
        cpu_compress_file(input_data, input_size, block_size,
                          &compressed_blocks, &block_sizes, &block_count);

    double cpu_time = timer_stop(&cpu_timer);

    if (result != 0) {
      LOG_ERR("CPU compression failed!");
      free(input_data);
      return 1;
    }

    LOG_OK("CPU compression done in %.3f seconds", cpu_time);
    double throughput = (double)input_size / (1024.0 * 1024.0) / cpu_time;
    LOG_OK("Throughput: %.1f MB/s", throughput);
  }

  // ── Build and write archive ────────────────────────────────────
  VkzArchive *archive = vkz_create(input_size, block_size, input_path);
  if (!archive) {
    LOG_ERR("Failed to create archive structure");
    free(input_data);
    return 1;
  }

  // Set block compressed sizes and checksums
  for (uint32_t i = 0; i < block_count; i++) {
    archive->blocks[i].compressed_size = block_sizes[i];
    archive->blocks[i].checksum =
        vkz_crc32(compressed_blocks[i], block_sizes[i]);
  }

  // Set directory flag if input was a directory
  if (is_dir) {
    archive->header.flags |= VKZ_FLAG_DIRECTORY;
  }

  // Set overall checksum
  archive->header.checksum = vkz_crc32(input_data, (size_t)input_size);

  int write_result =
      vkz_write(archive, (const uint8_t **)compressed_blocks, output_path);

  // Print results
  if (write_result == 0) {
    vkz_print_info(archive);

    double total_time = timer_stop(&total_timer);
    uint64_t output_file_size = get_file_size(output_path);
    char out_buf[32];
    format_size(output_file_size, out_buf, sizeof(out_buf));

    printf(COL_BOLD COL_GREEN "✓ Compression complete!" COL_RESET "\n");
    printf("  Method: %s\n", used_gpu ? "GPU (Vulkan)" : "CPU");
    printf("  Time:   %.3f seconds\n", total_time);
    printf("  Output: %s (%s)\n", output_path, out_buf);
    printf("\n");
  }

  // Cleanup
  for (uint32_t i = 0; i < block_count; i++) {
    free(compressed_blocks[i]);
  }
  free(compressed_blocks);
  free(block_sizes);
  free(input_data);
  vkz_free(archive);

  // Remove temp tar
  if (tar_temp_path)
    unlink(tar_temp_path);

  return write_result;
}

// ── Decompress command ─────────────────────────────────────────────
static int cmd_decompress(const char *archive_path, const char *output_path,
                          bool cpu_only) {
  // Check for standard external archives (zip, rar)
  if (ends_with(archive_path, ".zip")) {
    return extract_external_archive(archive_path, output_path, "zip");
  } else if (ends_with(archive_path, ".rar")) {
    return extract_external_archive(archive_path, output_path, "rar");
  }

  // Native VKZ decompression
  Timer total_timer;
  timer_start(&total_timer);

  VkzArchive *archive = vkz_read(archive_path);
  if (!archive) {
    return 1;
  }
  vkz_print_info(archive);

  // Generate output path if not provided
  char auto_output[VKZ_MAX_FILENAME];
  if (!output_path) {
    make_decompress_output(archive, auto_output, sizeof(auto_output));
    output_path = auto_output;
  }

  LOG_INFO("Output: %s", output_path);

  // Allocate output buffer
  uint8_t *output_data =
      (uint8_t *)calloc(1, (size_t)archive->header.original_size);
  if (!output_data) {
    char buf[32];
    format_size(archive->header.original_size, buf, sizeof(buf));
    LOG_ERR("Out of memory (need %s)", buf);
    vkz_free(archive);
    return 1;
  }

  bool used_gpu = false;

  // ── Try GPU decompression ──────────────────────────────────────
  if (!cpu_only) {
    GpuContext gpu;
    if (gpu_init(&gpu) == 0) {
      char shader_path[1024];
      if (find_shader_path(shader_path, sizeof(shader_path),
                           "decompress.comp.spv") == 0) {
        gpu.decompress_shader = gpu_load_shader(&gpu, shader_path);

        if (gpu.decompress_shader != VK_NULL_HANDLE) {
          GpuDecompressPipeline pipe;
          if (gpu_decompress_init(&pipe, &gpu) == 0) {
            LOG_INFO("Decompressing with GPU...");
            Timer gpu_timer;
            timer_start(&gpu_timer);

            int result =
                gpu_decompress_data(&pipe, archive, archive_path, output_data,
                                    archive->header.original_size);

            double gpu_time = timer_stop(&gpu_timer);

            if (result == 0) {
              used_gpu = true;
              LOG_OK("GPU decompression done in %.3f seconds", gpu_time);
            } else {
              LOG_WARN("GPU decompression failed, falling back to CPU");
            }

            gpu_decompress_cleanup(&pipe);
          }
        }
      }

      gpu_cleanup(&gpu);
    }
  }

  // ── CPU fallback ───────────────────────────────────────────────
  if (!used_gpu) {
    LOG_INFO("Decompressing with CPU...");
    Timer cpu_timer;
    timer_start(&cpu_timer);

    int result = cpu_decompress_file(archive, archive_path, output_data,
                                     archive->header.original_size);

    double cpu_time = timer_stop(&cpu_timer);

    if (result != 0) {
      LOG_ERR("Decompression failed!");
      free(output_data);
      vkz_free(archive);
      return 1;
    }

    LOG_OK("CPU decompression done in %.3f seconds", cpu_time);
  }

  // ── Verify checksum ────────────────────────────────────────────
  uint32_t checksum =
      vkz_crc32(output_data, (size_t)archive->header.original_size);
  if (archive->header.checksum != 0 && checksum != archive->header.checksum) {
    LOG_WARN("CRC32 mismatch! File may be corrupted.");
    LOG_WARN("Expected: 0x%08X  Got: 0x%08X", archive->header.checksum,
             checksum);
  } else if (archive->header.checksum != 0) {
    LOG_OK("CRC32 checksum verified: 0x%08X", checksum);
  }

  // ── Handle output ──────────────────────────────────────────────
  bool is_dir_archive = (archive->header.flags & VKZ_FLAG_DIRECTORY) != 0;

  if (is_dir_archive) {
    // Write tar to temp file, then extract
    char tar_path[VKZ_MAX_FILENAME];
    snprintf(tar_path, sizeof(tar_path), "/tmp/vkzip_decomp_%u.tar",
             (uint32_t)getpid());

    FILE *f = fopen(tar_path, "wb");
    if (!f) {
      LOG_ERR("Cannot create temp tar: %s", tar_path);
      free(output_data);
      vkz_free(archive);
      return 1;
    }
    fwrite(output_data, 1, (size_t)archive->header.original_size, f);
    fclose(f);

    LOG_INFO("Extracting directory...");
    // If output_path is specified, extract there; otherwise extract to current
    // dir
    int ret = untar_file(tar_path, output_path);
    unlink(tar_path);

    if (ret != 0) {
      LOG_ERR("Failed to extract directory from tar");
      free(output_data);
      vkz_free(archive);
      return 1;
    }

    double total_time = timer_stop(&total_timer);
    char size_buf[32];
    format_size(archive->header.original_size, size_buf, sizeof(size_buf));

    printf(COL_BOLD COL_GREEN "✓ Directory extracted!" COL_RESET "\n");
    printf("  Method: %s\n", used_gpu ? "GPU (Vulkan)" : "CPU");
    printf("  Time:   %.3f seconds\n", total_time);
    printf("  Output: %s\n", output_path ? output_path : "(current directory)");
    printf("  Tar:    %s\n", size_buf);
    printf("\n");
  } else {
    // Regular file output
    FILE *f = fopen(output_path, "wb");
    if (!f) {
      LOG_ERR("Cannot create output file: %s", output_path);
      free(output_data);
      vkz_free(archive);
      return 1;
    }

    fwrite(output_data, 1, (size_t)archive->header.original_size, f);
    fclose(f);

    double total_time = timer_stop(&total_timer);
    char size_buf[32];
    format_size(archive->header.original_size, size_buf, sizeof(size_buf));

    printf(COL_BOLD COL_GREEN "✓ Decompression complete!" COL_RESET "\n");
    printf("  Method: %s\n", used_gpu ? "GPU (Vulkan)" : "CPU");
    printf("  Time:   %.3f seconds\n", total_time);
    printf("  Output: %s (%s)\n", output_path, size_buf);
    printf("\n");
  }

  free(output_data);
  vkz_free(archive);
  return 0;
}

// ── Info command ───────────────────────────────────────────────────
static int cmd_info(const char *archive_path) {
  print_banner();

  VkzArchive *archive = vkz_read(archive_path);
  if (!archive)
    return 1;

  vkz_print_info(archive);

  // Print block details
  printf(COL_BOLD "Block Details:\n" COL_RESET);
  printf("  %-6s  %-12s  %-12s  %-8s  %-10s\n", "Block", "Original",
         "Compressed", "Ratio", "CRC32");
  printf("  ──────  ────────────  ────────────  ────────  ──────────\n");

  for (uint32_t i = 0; i < archive->header.block_count && i < 20; i++) {
    char orig_buf[32], comp_buf[32];
    format_size(archive->blocks[i].original_size, orig_buf, sizeof(orig_buf));
    format_size(archive->blocks[i].compressed_size, comp_buf, sizeof(comp_buf));

    double ratio = archive->blocks[i].original_size > 0
                       ? (double)archive->blocks[i].compressed_size /
                             (double)archive->blocks[i].original_size * 100.0
                       : 0.0;

    printf("  %-6u  %-12s  %-12s  %5.1f%%   0x%08X\n", i, orig_buf, comp_buf,
           ratio, archive->blocks[i].checksum);
  }

  if (archive->header.block_count > 20) {
    printf("  ... and %u more blocks\n", archive->header.block_count - 20);
  }

  printf("\n");
  vkz_free(archive);
  return 0;
}

// ── Benchmark command ──────────────────────────────────────────────
static int cmd_benchmark(const char *input_path) {
  print_banner();

  uint64_t input_size = get_file_size(input_path);
  if (input_size == 0) {
    LOG_ERR("Cannot read file: %s", input_path);
    return 1;
  }

  char size_buf[32];
  format_size(input_size, size_buf, sizeof(size_buf));
  LOG_INFO("Benchmarking: %s (%s)", input_path, size_buf);

  FILE *f = fopen(input_path, "rb");
  if (!f)
    return 1;

  uint8_t *input_data = (uint8_t *)malloc((size_t)input_size);
  fread(input_data, 1, (size_t)input_size, f);
  fclose(f);

  uint32_t block_size = VKZ_BLOCK_SIZE;
  double cpu_time = 0, gpu_time = 0;
  uint64_t cpu_output_size = 0, gpu_output_size = 0;

  // ── CPU Benchmark ──────────────────────────────────────────────
  {
    LOG_INFO("Running CPU benchmark...");
    uint8_t **blocks = NULL;
    uint32_t *sizes = NULL;
    uint32_t count = 0;

    Timer t;
    timer_start(&t);
    cpu_compress_file(input_data, input_size, block_size, &blocks, &sizes,
                      &count);
    cpu_time = timer_stop(&t);

    for (uint32_t i = 0; i < count; i++) {
      cpu_output_size += sizes[i];
      free(blocks[i]);
    }
    free(blocks);
    free(sizes);
  }

  // ── GPU Benchmark ──────────────────────────────────────────────
  {
    GpuContext gpu;
    if (gpu_init(&gpu) == 0) {
      gpu_print_info(&gpu);

      char shader_path[1024];
      if (find_shader_path(shader_path, sizeof(shader_path),
                           "compress.comp.spv") == 0) {
        gpu.compress_shader = gpu_load_shader(&gpu, shader_path);

        if (gpu.compress_shader != VK_NULL_HANDLE) {
          GpuCompressPipeline pipe;
          if (gpu_compress_init(&pipe, &gpu) == 0) {
            LOG_INFO("Running GPU benchmark...");
            uint8_t **blocks = NULL;
            uint32_t *sizes = NULL;
            uint32_t count = 0;

            Timer t;
            timer_start(&t);
            gpu_compress_data(&pipe, input_data, input_size, block_size,
                              &blocks, &sizes, &count);
            gpu_time = timer_stop(&t);

            for (uint32_t i = 0; i < count; i++) {
              gpu_output_size += sizes[i];
              free(blocks[i]);
            }
            free(blocks);
            free(sizes);

            gpu_compress_cleanup(&pipe);
          }
        }
      }
      gpu_cleanup(&gpu);
    } else {
      LOG_WARN("No GPU available for benchmark");
    }
  }

  // ── Results ────────────────────────────────────────────────────
  printf("\n");
  printf(COL_BOLD COL_YELLOW
         "╔══════════════════════════════════════════════════╗\n" COL_RESET);
  printf(COL_BOLD COL_YELLOW
         "║              Benchmark Results                   ║\n" COL_RESET);
  printf(COL_BOLD COL_YELLOW
         "╠══════════════════════════════════════════════════╣\n" COL_RESET);

  char cpu_buf[32], gpu_buf[32];
  format_size(cpu_output_size, cpu_buf, sizeof(cpu_buf));

  double cpu_ratio =
      input_size > 0 ? (double)cpu_output_size / (double)input_size * 100.0 : 0;
  double cpu_throughput =
      cpu_time > 0 ? (double)input_size / (1024.0 * 1024.0) / cpu_time : 0;

  printf(COL_YELLOW "║" COL_RESET "  CPU:\n");
  printf(COL_YELLOW "║" COL_RESET "    Time:       %.3f sec\n", cpu_time);
  printf(COL_YELLOW "║" COL_RESET "    Output:     %s (%.1f%%)\n", cpu_buf,
         cpu_ratio);
  printf(COL_YELLOW "║" COL_RESET "    Throughput: %.1f MB/s\n",
         cpu_throughput);
  printf(COL_YELLOW "║" COL_RESET "\n");

  if (gpu_time > 0) {
    format_size(gpu_output_size, gpu_buf, sizeof(gpu_buf));
    double gpu_ratio =
        input_size > 0 ? (double)gpu_output_size / (double)input_size * 100.0
                       : 0;
    double gpu_throughput =
        gpu_time > 0 ? (double)input_size / (1024.0 * 1024.0) / gpu_time : 0;

    printf(COL_YELLOW "║" COL_RESET "  GPU:\n");
    printf(COL_YELLOW "║" COL_RESET "    Time:       %.3f sec\n", gpu_time);
    printf(COL_YELLOW "║" COL_RESET "    Output:     %s (%.1f%%)\n", gpu_buf,
           gpu_ratio);
    printf(COL_YELLOW "║" COL_RESET "    Throughput: %.1f MB/s\n",
           gpu_throughput);
    printf(COL_YELLOW "║" COL_RESET "\n");

    double speedup = cpu_time > 0 ? cpu_time / gpu_time : 0;
    printf(COL_YELLOW "║" COL_RESET COL_BOLD "  Speedup: %.2fx %s\n" COL_RESET,
           speedup,
           speedup > 1.0 ? COL_GREEN "GPU wins! 🚀" COL_RESET
                         : COL_RED "CPU wins" COL_RESET);
  } else {
    printf(COL_YELLOW "║" COL_RESET "  GPU: " COL_RED "Not available" COL_RESET
                      "\n");
  }

  printf(COL_BOLD COL_YELLOW
         "╚══════════════════════════════════════════════════╝\n" COL_RESET);
  printf("\n");

  free(input_data);
  return 0;
}

// ── GPU Info command ───────────────────────────────────────────────
static int cmd_gpu_info(void) {
  print_banner();

  GpuContext gpu;
  if (gpu_init(&gpu) != 0) {
    LOG_ERR("No Vulkan GPU available");
    return 1;
  }

  gpu_print_info(&gpu);
  gpu_cleanup(&gpu);
  return 0;
}

// ── Main ───────────────────────────────────────────────────────────
int main(int argc, char *argv[]) {
  if (argc < 2) {
    print_help();
    return 0;
  }

  // Parse global options
  bool cpu_only = false;
  uint32_t block_size_kb = 64;

  // Check for flags anywhere in args
  for (int i = 1; i < argc; i++) {
    if (strcmp(argv[i], "--cpu-only") == 0) {
      cpu_only = true;
    } else if (strcmp(argv[i], "--block-size") == 0 && i + 1 < argc) {
      block_size_kb = (uint32_t)atoi(argv[++i]);
      if (block_size_kb < 4)
        block_size_kb = 4;
      if (block_size_kb > 1024)
        block_size_kb = 1024;
    }
  }

  const char *command = argv[1];

  // ── Route commands ─────────────────────────────────────────────
  if (strcmp(command, "--help") == 0 || strcmp(command, "-h") == 0) {
    print_help();
    return 0;
  }

  if (strcmp(command, "--gpu-info") == 0) {
    return cmd_gpu_info();
  }

  if (strcmp(command, "compress") == 0 || strcmp(command, "c") == 0) {
    if (argc < 3) {
      LOG_ERR("Usage: vkzip compress <input_file> [output.vkz]");
      return 1;
    }
    print_banner();
    const char *input = argv[2];
    const char *output = (argc > 3 && argv[3][0] != '-') ? argv[3] : NULL;
    return cmd_compress(input, output, cpu_only, block_size_kb);
  }

  if (strcmp(command, "decompress") == 0 || strcmp(command, "d") == 0 ||
      strcmp(command, "extract") == 0 || strcmp(command, "x") == 0) {
    if (argc < 3) {
      LOG_ERR("Usage: vkzip decompress <archive.vkz> [output_file]");
      return 1;
    }
    print_banner();
    const char *input = argv[2];
    const char *output = (argc > 3 && argv[3][0] != '-') ? argv[3] : NULL;
    return cmd_decompress(input, output, cpu_only);
  }

  if (strcmp(command, "info") == 0 || strcmp(command, "i") == 0) {
    if (argc < 3) {
      LOG_ERR("Usage: vkzip info <archive.vkz>");
      return 1;
    }
    return cmd_info(argv[2]);
  }

  if (strcmp(command, "benchmark") == 0 || strcmp(command, "bench") == 0 ||
      strcmp(command, "b") == 0) {
    if (argc < 3) {
      LOG_ERR("Usage: vkzip benchmark <input_file>");
      return 1;
    }
    return cmd_benchmark(argv[2]);
  }

  LOG_ERR("Unknown command: %s", command);
  LOG_ERR("Run 'vkzip --help' for usage information");
  return 1;
}