#ifndef MESSAGE_C
#define MESSAGE_C

#include "aes.c"
#include "allocator.c"
#include "compress.c"
#include "const.c"
#include "sha256.c"
#include "types.c"
#include "utils.c"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct {
  char *sender;
  char *content;
  char *channel;
  timestamp_t timestamp;
} message_t;

message_t *message_create(char *sender, char *channel, char *content) {
  message_t *msg = malloc(sizeof(message_t));
  if (msg == NULL) {
    return NULL;
  }

  msg->sender = strndup(sender, MAX_SENDER_LEN);
  msg->channel = strndup(channel, MAX_CHANNEL_LEN);
  msg->content = strndup(content, MAX_CONTENT_LEN);
  msg->timestamp = get_current_timestamp();
  return msg;
}

void message_inner_free(message_t *msg) {
  if (msg->sender != NULL) {
    free(msg->sender);
  }

  if (msg->content != NULL) {
    free(msg->content);
  }

  if (msg->channel != NULL) {
    free(msg->channel);
  }

  if (msg != NULL) {
    free(msg);
  }
}

// malloc
byte_t *message_serialize(message_t *msg, size_t *out_len) {
  const size_t sender_len = strlen(msg->sender);
  const size_t content_len = strlen(msg->content);

  if (sender_len > MAX_SENDER_LEN) {
    return NULL;
  }
  if (content_len > MAX_CONTENT_LEN) {
    return NULL;
  }

  byte_t *buffer = malloc(
      (TIMESTAMP_LEN + MAX_SENDER_LEN + MAX_CONTENT_LEN + MAX_CHANNEL_LEN) *
      sizeof(byte_t));
  if (buffer == NULL) {
    return NULL;
  }
  memset(buffer, 0,
         TIMESTAMP_LEN + MAX_SENDER_LEN + MAX_CONTENT_LEN + MAX_CHANNEL_LEN);

  memcpy(buffer + TIMESTAMP_LEN, msg->sender, sender_len);
  memcpy(buffer + TIMESTAMP_LEN + MAX_SENDER_LEN, msg->content, content_len);
  memcpy(buffer + TIMESTAMP_LEN + MAX_SENDER_LEN + MAX_CONTENT_LEN,
         msg->channel, strlen(msg->channel));
  memcpy(buffer, &msg->timestamp, TIMESTAMP_LEN);
  *out_len = TIMESTAMP_LEN + MAX_SENDER_LEN + MAX_CONTENT_LEN + MAX_CHANNEL_LEN;
  return buffer;
}

// malloc
message_t *message_deserialize(byte_t *buffer, size_t buffer_len) {
  if (buffer_len !=
      TIMESTAMP_LEN + MAX_SENDER_LEN + MAX_CONTENT_LEN + MAX_CHANNEL_LEN) {
#ifdef DEBUG
    printf("<!> Invalid buffer length for deserialization: %zu not %zu\n",
           buffer_len,
           TIMESTAMP_LEN + MAX_SENDER_LEN + MAX_CONTENT_LEN + MAX_CHANNEL_LEN);
#endif
    return NULL;
  }

  message_t *msg = malloc(sizeof(message_t));
  if (msg == NULL) {
    return NULL;
  }

  msg->sender = malloc(sizeof(char) * MAX_SENDER_LEN);
  msg->content = malloc(sizeof(char) * MAX_CONTENT_LEN);
  msg->channel = malloc(sizeof(char) * MAX_CHANNEL_LEN);
  if (msg->sender == NULL) {
    free(msg);
    return NULL;
  }
  if (msg->content == NULL) {
    free(msg->sender);
    free(msg);
    return NULL;
  }
  if (msg->channel == NULL) {
    free(msg->sender);
    free(msg->content);
    free(msg);
    return NULL;
  }

  memcpy(msg->sender, buffer + TIMESTAMP_LEN, MAX_SENDER_LEN);
  memcpy(msg->content, buffer + TIMESTAMP_LEN + MAX_SENDER_LEN,
         MAX_CONTENT_LEN);
  memcpy(msg->channel,
         buffer + TIMESTAMP_LEN + MAX_SENDER_LEN + MAX_CONTENT_LEN,
         MAX_CHANNEL_LEN);
  memcpy(&msg->timestamp, buffer, TIMESTAMP_LEN);

  return msg;
}

char *message_to_transmittable(message_t *msg, bool_t do_compression,
                               char *password, size_t *out_len) {
  // Serialize
  size_t serialized_msg_len = 0;
  defer(free_checked) byte_t *serialized_msg =
      message_serialize(msg, &serialized_msg_len);
  if (serialized_msg == NULL) {
#ifdef DEBUG
    printf("<!> Failed to serialize message\n");
#endif
    return NULL;
  }

  // Compress
  byte_t *compressed_msg;
  size_t compressed_msg_len = 0;
  if (do_compression) {
    compressed_msg =
        compress(serialized_msg, serialized_msg_len, &compressed_msg_len);
    if (compressed_msg == NULL) {
#ifdef DEBUG
      printf("<!> Failed to compress message\n");
#endif
      return NULL;
    }
  } else {
    compressed_msg = serialized_msg;
    compressed_msg_len = serialized_msg_len;
  }

  // Encrypt
  defer(free_checked) byte_t *hashed_password =
      sha256((byte_t *)password, strlen(password));
  if (hashed_password == NULL) {
#ifdef DEBUG
    printf("<!> Failed to hash password\n");
#endif
    return NULL;
  }

  size_t encrypted_msg_len = 0;
  defer(free_checked) byte_t *encrypted_msg =
      aes_cbc(encryption, compressed_msg, compressed_msg_len, hashed_password,
              &encrypted_msg_len);
  if (encrypted_msg == NULL) {
#ifdef DEBUG
    printf("<!> Failed to encrypt message\n");
#endif
    return NULL;
  }

  if (do_compression) {
    free(compressed_msg);
  }

  // Checksum
  defer(free_checked) byte_t *checksummed_msg =
      malloc(sizeof(byte_t) * (encrypted_msg_len + 1));
  if (checksummed_msg == NULL) {
#ifdef DEBUG
    printf("<!> Failed to allocate checksummed message\n");
#endif
    return NULL;
  }
  memcpy(checksummed_msg, encrypted_msg, encrypted_msg_len);
  checksummed_msg[encrypted_msg_len] =
      checksum(encrypted_msg, encrypted_msg_len);

  // Hex
  char *msg_transmittable = btoh(checksummed_msg, encrypted_msg_len + 1);
  if (msg_transmittable == NULL) {
#ifdef DEBUG
    printf("<!> Failed to convert message to hex\n");
#endif
    return NULL;
  }

  *out_len = encrypted_msg_len + 1;
  return msg_transmittable;
}

message_t *message_from_transmittable(char *transmittable,
                                      bool_t do_compression, char *password) {
  // Hex to bin
  size_t checked_msg_len = 0;
  byte_t *checked_msg = htob(transmittable, &checked_msg_len);
  if (checked_msg == NULL) {
#ifdef DEBUG
    printf("<!> Failed to convert transmittable from hex\n");
#endif
    return NULL;
  }

  // Checksum
  byte_t sum = checksum(checked_msg, checked_msg_len - 1);
  if (sum != checked_msg[checked_msg_len - 1]) {
#ifdef DEBUG
    printf("<!> Checksum mismatch for transmittable message\n");
#endif
    return NULL;
  }

  // Decrypt
  byte_t *hashed_password = sha256((byte_t *)password, strlen(password));
  if (hashed_password == NULL) {
#ifdef DEBUG
    printf("<!> Failed to hash password\n");
#endif
    return NULL;
  }

  size_t decrypted_msg_len = 0;
  byte_t *decrypted_msg = aes_cbc(decryption, checked_msg, checked_msg_len - 1,
                                  hashed_password, &decrypted_msg_len);
  if (decrypted_msg == NULL) {
#ifdef DEBUG
    printf("<!> Failed to decrypt message\n");
#endif
    return NULL;
  }

  // Decompress
  byte_t *decompressed_msg;
  size_t decompressed_msg_len = 0;

  if (do_compression) {
    decompressed_msg =
        decompress(decrypted_msg, decrypted_msg_len, &decompressed_msg_len);
    if (decompressed_msg == NULL) {
#ifdef DEBUG
      printf("<!> Failed to decompress message\n");
#endif
      return NULL;
    }
  } else {
    decompressed_msg = decrypted_msg;
    decompressed_msg_len = decrypted_msg_len;
  }

  // Deserialize
  message_t *msg = message_deserialize(decompressed_msg, decompressed_msg_len);
  if (msg == NULL) {
#ifdef DEBUG
    printf("<!> Failed to deserialize message\n");
#endif
    return NULL;
  }

  return msg;
}

void message_print(message_t *msg) {
  printf("[%llu] %s#%s: %s\n", msg->timestamp, msg->sender, msg->channel,
         msg->content);
}

int test_message() {
  printf("--- START TEST:\t message ---\n");

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

  size_t out_transmittable_len;
  char *transmittable =
      message_to_transmittable(msg, false, "password", &out_transmittable_len);
  if (transmittable == NULL) {
    return 1;
  }
  printf("Transmittable message:\t%s\n", transmittable);

  message_t *from_transmittable_msg =
      message_from_transmittable(transmittable, false, "password");
  printf("From transmittable message:\t");
  message_print(from_transmittable_msg);

  if (strcmp(msg->sender, from_transmittable_msg->sender) != 0 ||
      strcmp(msg->content, from_transmittable_msg->content) != 0 ||
      strcmp(msg->channel, from_transmittable_msg->channel) != 0 ||
      msg->timestamp != from_transmittable_msg->timestamp) {

    return 4;
  }

  printf("--- END TEST:\t message ---\n");
  return 0;
}

#endif
