#ifndef COMPRESS_C
#define COMPRESS_C

#include "types.c"
#include <stdlib.h>

// Compress data in-place and return new length
byte_t *compress(byte_t *data, size_t length, size_t *out_compressed_len) {
  byte_t *buffer = malloc(length * 2 * sizeof(byte_t));
  if (buffer == NULL) {
    return NULL;
  }

  size_t buffer_index = 0;
  byte_t holder = 0;
  size_t count = 0;
  for (size_t i = 0; i < length; i++) {
    if (i == 0) {
      holder = data[i];
      count = 1;
      continue;
    }

    if (data[i] == holder) {
      count++;
    } else {
      buffer[buffer_index++] = holder;
      buffer[buffer_index++] = count;
      holder = data[i];
      count = 1;
    }
  }
  if (count > 0) {
    buffer[buffer_index++] = holder;
    buffer[buffer_index++] = count;
  }

  *out_compressed_len = buffer_index;

  byte_t *resized_buffer = malloc(length * sizeof(byte_t));
  if (resized_buffer == NULL) {
    free(buffer);
    return NULL;
  }

  memcpy(resized_buffer, buffer, length * sizeof(byte_t));
  free(buffer);
  return resized_buffer;
}

byte_t *decompress(byte_t *data, size_t length,
                   size_t *out_decompressed_length) {
  size_t output_length = 0;
  for (size_t i = 1; i < length; i += 2) {
    output_length += data[i];
  }

  byte_t *buffer = malloc(output_length * sizeof(byte_t));
  if (buffer == NULL) {
    return NULL;
  }

  size_t buffer_index = 0;
  for (size_t i = 1; i < length; i += 2) {
    byte_t value = data[i - 1];
    byte_t count = data[i];
    for (size_t j = 0; j < count; j++) {
      buffer[buffer_index++] = value;
    }
  }

  *out_decompressed_length = buffer_index;
  return buffer;
}

int test_compress() {
  printf("--- START TEST:\t compress ---\n");
  allocator_storage_t a = allocator_storage_create(2);

  size_t length = 20;
  byte_t *data = malloc(length * sizeof(byte_t));
  if (data == NULL) {
    allocator_storage_free_inner_pointers(&a);
    return 1;
  }
  allocator_storage_add(&a, (generic_pointer_t)data);
  memset(data, 0x22, 10);
  memset(data + 10, 0x33, 10);

  printf("Original data:\t\t");
  print_array(data, length);

  size_t new_data_len;
  byte_t *new_data = compress((byte_t *)data, length, &new_data_len);
  if (new_data == NULL) {
    allocator_storage_free_inner_pointers(&a);
    return 1;
  }
  allocator_storage_add(&a, (generic_pointer_t)new_data);

  printf("Compressed data:\t");
  print_array(new_data, new_data_len);

  size_t decompressed_data_len;
  byte_t *decompressed_data =
      decompress(new_data, length, &decompressed_data_len);
  if (decompressed_data == NULL) {
    allocator_storage_free_inner_pointers(&a);
    return 1;
  }
  allocator_storage_add(&a, (generic_pointer_t)decompressed_data);

  printf("Decompressed data:\t");
  print_array(decompressed_data, length);

  if (memcmp(decompressed_data, data, length * sizeof(byte_t)) != 0) {
    printf("Decompressed data does not match original data!\n");
    allocator_storage_free_inner_pointers(&a);
    return 1;
  }

  allocator_storage_free_inner_pointers(&a);
  printf("--- END TEST:\t compress ---\n");
  return 0;
}

#endif
