#include "./lib/lib.c"
#include <stdio.h>

void _test_defer_cleanup(int *x) { printf("Deferred: x = %d\n", *x); }
int test_defer() {
  printf("--- START TEST:\t defer ---\n");
  {
    defer(_test_defer_cleanup) int x = 42;
    printf("Inside block: x = %d\n", x);
  }
  printf("--- END TEST:\t defer ---\n");
  return 0;
}

int test_message_encryption() {
  printf("--- START TEST:\t message_encryption ---\n");

  allocator_storage_t a = allocator_storage_create(10);

  message_t *msg = message_create("Alice", "root", "Hello, Bob!");

  // Serialize
  size_t msg_serialized_len;
  byte_t *msg_serialized = message_serialize(msg, &msg_serialized_len);
  ALLOCATOR_STORE_OR_RETURN(&a, msg_serialized, 1);

  // Encrypt
  byte_t *key = (byte_t *)"password";
  size_t msg_serialized_encrypted_len;
  byte_t *msg_serialized_encrypted =
      aes_cbc(encryption, msg_serialized, msg_serialized_len, key,
              &msg_serialized_encrypted_len);
  ALLOCATOR_STORE_OR_RETURN(&a, msg_serialized_encrypted, 2);

  // Decrypt
  size_t msg_serialized_decrypted_len;
  byte_t *msg_serialized_decrypted =
      aes_cbc(decryption, msg_serialized_encrypted,
              msg_serialized_encrypted_len, key, &msg_serialized_decrypted_len);
  ALLOCATOR_STORE_OR_RETURN(&a, msg_serialized_decrypted, 3);

  // Deserialize
  message_t *msg_deserialized = message_deserialize(
      msg_serialized_decrypted, msg_serialized_decrypted_len);

  printf("Message: ");
  message_print(msg);
  printf("Serialized message: %s\n", btoh(msg_serialized, msg_serialized_len));
  printf("Serialized message: %s\n",
         btoh(msg_serialized_encrypted, msg_serialized_encrypted_len));
  printf("Serialized message: %s\n",
         btoh(msg_serialized_decrypted, msg_serialized_decrypted_len));
  printf("Deserialized message: ");
  message_print(msg_deserialized);

  // Message == Decrypted deserialized message
  if (strcmp(msg->sender, msg_deserialized->sender) != 0 ||
      strcmp(msg->content, msg_deserialized->content) != 0 ||
      msg->timestamp != msg_deserialized->timestamp) {
    printf("Decrypted+deserialized message does not match original "
           "message!\n");
    allocator_storage_free_inner_pointers(&a);
    message_inner_free(msg);
    message_inner_free(msg_deserialized);
    return 1;
  }

  // Enc != Dec
  if (memcmp(msg_serialized, msg_serialized_encrypted, msg_serialized_len) ==
      0) {
    printf("Encrypted message should not match original serialized message!\n");
    allocator_storage_free_inner_pointers(&a);
    message_inner_free(msg);
    message_inner_free(msg_deserialized);
    return 1;
  }

  allocator_storage_free_inner_pointers(&a);
  message_inner_free(msg);
  message_inner_free(msg_deserialized);
  printf("--- END TEST:\t compress_message ---\n");
  return 0;
}

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

  message_t *msg = message_create("Alice", "root", "Hello, Bob!");
  printf("Original message:\t");
  message_print(msg);

  size_t ser_len;
  byte_t *ser = message_serialize(msg, &ser_len);
  ALLOCATOR_STORE_OR_RETURN(&a, ser, 1);
  printf("Serialized message:\t%s\n", btoh(ser, ser_len));

  size_t ser_compressed_len;
  byte_t *ser_compressed = compress(ser, ser_len, &ser_compressed_len);
  ALLOCATOR_STORE_OR_RETURN(&a, ser_compressed, 1);
  printf("Compressed message:\t%s\n", btoh(ser_compressed, ser_compressed_len));

  size_t ser_decompressed_len;
  byte_t *ser_decompressed =
      decompress(ser_compressed, ser_compressed_len, &ser_decompressed_len);
  ALLOCATOR_STORE_OR_RETURN(&a, ser_decompressed, 1);
  printf("Decompressed message:\t%s\n",
         btoh(ser_decompressed, ser_decompressed_len));

  // message_t *deserialized_msg = message_deserialize((char*)ser_decompressed);
  // if (strcmp(msg->sender, deserialized_msg->sender) != 0 ||
  //     strcmp(msg->content, deserialized_msg->content) != 0 ||
  //     msg->timestamp != deserialized_msg->timestamp) {
  //     printf("Decompressed message does not match original message!\n");
  //     return 1;
  // }

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

int main() {
  int err = test_message();
  if (err != 0) {
    printf("test_message failed with error code %d\n", err);
    return 1;
  }

  err = test_compress();
  if (err != 0) {
    printf("test_compress failed with error code %d\n", err);
    return 1;
  }

  err = test_compress_message();
  if (err != 0) {
    printf("test_compress_message failed with error code %d\n", err);
    return 1;
  }

  err = test_message_encryption();
  if (err != 0) {
    printf("test_message_encryption failed with error code %d\n", err);
    return 1;
  }

  err = test_vector();
  if (err != 0) {
    printf("test_vector failed with error code %d\n", err);
    return 1;
  }

  err = test_defer();
  if (err != 0) {
    printf("test_defer failed with error code %d\n", err);
    return 1;
  }

  err = test_split_string();
  if (err != 0) {
    printf("test_split_string failed with error code %d\n", err);
    return 1;
  }

  printf("OK\n");
  return 0;
}
