#include "cpu_fallback.h"

// ── CPU LZ77 Compress Block ────────────────────────────────────────
uint8_t *cpu_compress_block(const uint8_t *input, uint32_t input_size,
                            uint32_t *out_size) {
  // Worst case: 2 bytes per input byte (literal tokens)
  uint32_t max_output = input_size * 2 + 16;
  uint8_t *output = (uint8_t *)malloc(max_output);
  if (!output)
    return NULL;

  // Simple hash table for match finding
  uint32_t hash_table[VKZ_HASH_SIZE];
  memset(hash_table, 0xFF, sizeof(hash_table));

  uint32_t pos = 0;
  uint32_t out_pos = 0;

  while (pos < input_size) {
    uint32_t best_len = 0;
    uint32_t best_dist = 0;

    // Try to find a match (need at least 3 bytes)
    if (pos + 2 < input_size) {
      // Hash 3 bytes
      uint32_t h = ((uint32_t)input[pos] << 16 | (uint32_t)input[pos + 1] << 8 |
                    (uint32_t)input[pos + 2]) %
                   VKZ_HASH_SIZE;

      uint32_t match_pos = hash_table[h];
      hash_table[h] = pos;

      if (match_pos != 0xFFFFFFFF && pos > match_pos) {
        uint32_t dist = pos - match_pos;
        if (dist <= VKZ_WINDOW_SIZE && dist > 0) {
          uint32_t len = 0;
          uint32_t max_len = input_size - pos;
          if (max_len > VKZ_MAX_MATCH_LEN)
            max_len = VKZ_MAX_MATCH_LEN;

          while (len < max_len && input[match_pos + len] == input[pos + len]) {
            len++;
          }

          if (len >= VKZ_MIN_MATCH_LEN) {
            best_len = len;
            best_dist = dist;
          }
        }
      }
    }

    if (best_len >= VKZ_MIN_MATCH_LEN) {
      // Match token: [0x01] [len_lo] [len_hi] [dist_lo] [dist_hi]
      output[out_pos++] = 0x01;
      output[out_pos++] = (uint8_t)(best_len & 0xFF);
      output[out_pos++] = (uint8_t)((best_len >> 8) & 0xFF);
      output[out_pos++] = (uint8_t)(best_dist & 0xFF);
      output[out_pos++] = (uint8_t)((best_dist >> 8) & 0xFF);
      pos += best_len;
    } else {
      // Literal token: [0x00] [byte]
      output[out_pos++] = 0x00;
      output[out_pos++] = input[pos];
      pos++;
    }
  }

  *out_size = out_pos;

  // Shrink to actual size
  output = (uint8_t *)realloc(output, out_pos);
  return output;
}

// ── CPU LZ77 Decompress Block ──────────────────────────────────────
uint8_t *cpu_decompress_block(const uint8_t *compressed,
                              uint32_t compressed_size, uint32_t original_size,
                              uint32_t *out_size) {
  uint8_t *output = (uint8_t *)malloc(original_size);
  if (!output)
    return NULL;

  uint32_t in_pos = 0;
  uint32_t out_pos = 0;

  while (in_pos < compressed_size && out_pos < original_size) {
    uint8_t token = compressed[in_pos++];

    if (token == 0x00) {
      // Literal
      if (in_pos >= compressed_size)
        break;
      output[out_pos++] = compressed[in_pos++];
    } else if (token == 0x01) {
      // Match
      if (in_pos + 4 > compressed_size)
        break;

      uint32_t match_len = (uint32_t)compressed[in_pos] |
                           ((uint32_t)compressed[in_pos + 1] << 8);
      in_pos += 2;
      uint32_t match_dist = (uint32_t)compressed[in_pos] |
                            ((uint32_t)compressed[in_pos + 1] << 8);
      in_pos += 2;

      uint32_t copy_src = out_pos - match_dist;
      for (uint32_t i = 0; i < match_len && out_pos < original_size; i++) {
        output[out_pos++] = output[copy_src + i];
      }
    }
  }

  *out_size = out_pos;
  return output;
}

// ── Compress file (multi-block, CPU) ───────────────────────────────
int cpu_compress_file(const uint8_t *input_data, uint64_t input_size,
                      uint32_t block_size, uint8_t ***out_blocks,
                      uint32_t **out_block_sizes, uint32_t *out_block_count) {

  uint32_t n_blocks = (uint32_t)((input_size + block_size - 1) / block_size);

  *out_block_count = n_blocks;
  *out_blocks = (uint8_t **)calloc(n_blocks, sizeof(uint8_t *));
  *out_block_sizes = (uint32_t *)calloc(n_blocks, sizeof(uint32_t));

  for (uint32_t i = 0; i < n_blocks; i++) {
    uint64_t offset = (uint64_t)i * block_size;
    uint64_t remaining = input_size - offset;
    uint32_t this_size =
        (uint32_t)(remaining < block_size ? remaining : block_size);

    uint32_t comp_size = 0;
    uint8_t *compressed =
        cpu_compress_block(input_data + offset, this_size, &comp_size);

    if (!compressed || comp_size >= this_size) {
      // Store raw if compression doesn't help
      if (compressed)
        free(compressed);
      (*out_blocks)[i] = (uint8_t *)malloc(this_size);
      memcpy((*out_blocks)[i], input_data + offset, this_size);
      (*out_block_sizes)[i] = this_size;
    } else {
      (*out_blocks)[i] = compressed;
      (*out_block_sizes)[i] = comp_size;
    }

    print_progress("Compressing (CPU)", i + 1, n_blocks);
  }

  return 0;
}

// ── Decompress file (multi-block, CPU) ─────────────────────────────
int cpu_decompress_file(const VkzArchive *archive, const char *archive_path,
                        uint8_t *output_data, uint64_t output_size) {

  uint32_t n_blocks = archive->header.block_count;
  (void)output_size;

  for (uint32_t i = 0; i < n_blocks; i++) {
    uint32_t comp_size;
    uint8_t *compressed = vkz_read_block(archive_path, archive, i, &comp_size);
    if (!compressed)
      return -1;

    uint32_t orig_size = archive->blocks[i].original_size;
    uint64_t offset = (uint64_t)i * archive->header.block_size;

    // Check if block was stored raw (compressed_size == original_size)
    if (comp_size == orig_size) {
      memcpy(output_data + offset, compressed, orig_size);
    } else {
      uint32_t dec_size;
      uint8_t *decompressed =
          cpu_decompress_block(compressed, comp_size, orig_size, &dec_size);
      if (!decompressed) {
        free(compressed);
        return -1;
      }
      memcpy(output_data + offset, decompressed, dec_size);
      free(decompressed);
    }

    free(compressed);
    print_progress("Decompressing (CPU)", i + 1, n_blocks);
  }

  return 0;
}
