#include "str8.h"
#include "aliases.h"
#include "mem_allocator.h"
#include <stdarg.h>
#include <stddef.h>
#include <stdio.h>
#include <string.h>
#include <stdbool.h>

#define STR8_BUF_ALLOC_SIZE(CAPACITY) (sizeof(Str8) + sizeof(c8) * CAPACITY)

internal Str8List node_to_list(Str8Node *node);

Str8 *wapp_str8_alloc_buf(const Allocator *allocator, u64 capacity) {
  Str8 *str = NULL;

  if (!allocator) {
    goto RETURN_STR8;
  }

  str = wapp_mem_allocator_alloc(allocator, STR8_BUF_ALLOC_SIZE(capacity));
  if (!str) {
    goto RETURN_STR8;
  }

  str->buf      = (u8 *)str + sizeof(Str8);
  str->size     = 0;
  str->capacity = capacity;

RETURN_STR8:
  return str;
}

Str8 *wapp_str8_alloc_cstr(const Allocator *allocator, const char *str) {
  Str8 *output = NULL;

  if (!allocator || !str) {
    goto RETURN_ALLOC_CSTR;
  }

  u64 length = strlen(str);
  output     = wapp_str8_alloc_buf(allocator, length * 2);
  if (!output) {
    goto RETURN_ALLOC_CSTR;
  }

  output->size = length;
  memcpy(output->buf, str, length);

RETURN_ALLOC_CSTR:
  return output;
}

Str8 *wapp_str8_alloc_str8(const Allocator *allocator, Str8RO *str) {
  Str8 *output = NULL;

  if (!allocator || !str) {
    goto RETURN_ALLOC_STR8;
  }

  output = wapp_str8_alloc_buf(allocator, str->capacity);
  if (!output) {
    goto RETURN_ALLOC_STR8;
  }

  output->size = str->size;
  memcpy(output->buf, str->buf, str->size);

RETURN_ALLOC_STR8:
  return output;
}

Str8 *wapp_str8_alloc_substr(const Allocator *allocator, Str8RO *str, u64 start, u64 end) {
  Str8 *output = NULL;
  if (!allocator || !str) {
    goto RETURN_ALLOC_SUBSTR;
  }

  if (start >= str->size || start >= end) {
    goto RETURN_ALLOC_SUBSTR;
  }

  if (end > str->size) {
    end = str->size;
  }

  output = wapp_str8_alloc_buf(allocator, str->capacity);
  if (!output) {
    goto RETURN_ALLOC_SUBSTR;
  }

  output->size = end - start;
  memcpy(output->buf, str->buf + start, output->size);

RETURN_ALLOC_SUBSTR:
  return output;
}

void wapp_str8_dealloc_buf(const Allocator *allocator, Str8 **str) {
  if (!allocator || !str || !(*str)) {
    return;
  }

  wapp_mem_allocator_free(allocator, (void **)str, STR8_BUF_ALLOC_SIZE((*str)->capacity));
}

c8 wapp_str8_get(const Str8 *str, u64 index) {
  if (index >= str->size) {
    return '\0';
  }

  return str->buf[index];
}

void wapp_str8_set(Str8 *str, u64 index, c8 c) {
  if (index >= str->size) {
    return;
  }

  str->buf[index] = c;
}

void wapp_str8_push_back(Str8 *str, c8 c) {
  if (!(str->size < str->capacity)) {
    return;
  }

  u64 index = (str->size)++;
  wapp_str8_set(str, index, c);
}

bool wapp_str8_equal(Str8RO *s1, Str8RO *s2) {
  if (s1->size != s2->size) {
    return false;
  }

  return wapp_str8_equal_to_count(s1, s2, s1->size);
}

bool wapp_str8_equal_to_count(Str8RO* s1, Str8RO* s2, u64 count) {
    if (!s1 || !s2) {
        return false;
    }

    return memcmp(s1->buf, s2->buf, count) == 0;
}

Str8 wapp_str8_slice(Str8RO *str, u64 start, u64 end) {
  if (start >= str->size || start >= end) {
    start = str->size;
    end   = str->size;
  }

  if (end > str->size) {
    end = str->size;
  }

  return (Str8RO){
    .capacity = end - start,
    .size     = end - start,
    .buf      = str->buf + start,
  };
}

Str8 *wapp_str8_alloc_concat(const Allocator *allocator, Str8 *dst, Str8RO *src) {
  if (!allocator || !dst || !src) {
    return NULL;
  }

  Str8 *output  = NULL;
  u64 remaining = dst->capacity - dst->size;
  if (src->size <= remaining) {
    output = dst;
    goto COPY_STRING_STR8_CONCAT;
  }

  u64 capacity = dst->capacity + src->size;

  output = wapp_str8_alloc_buf(allocator, capacity);
  if (!output) {
    goto RETURN_STR8_CONCAT;
  }

  wapp_str8_concat_capped(output, dst);

COPY_STRING_STR8_CONCAT:
  wapp_str8_concat_capped(output, src);

RETURN_STR8_CONCAT:
  return output;
}

void wapp_str8_concat_capped(Str8 *dst, Str8RO *src) {
  if (!dst || !src) {
    return;
  }

  u64 remaining = dst->capacity - dst->size;
  u64 to_copy   = remaining < src->size ? remaining : src->size;

  memcpy(dst->buf + dst->size, src->buf, to_copy);
  dst->size += to_copy;
}

void wapp_str8_copy_cstr_capped(Str8 *dst, const char *src) {
  if (!dst || !src) {
    return;
  }

  u64 length  = strlen(src);
  u64 to_copy = length <= dst->capacity ? length : dst->capacity;

  memset(dst->buf, 0, dst->size);
  memcpy(dst->buf, src, to_copy);
  dst->size = to_copy;
}

void wapp_str8_copy_str8_capped(Str8 *dst, Str8RO *src) {
  if (!dst || !src) {
    return;
  }

  u64 to_copy = src->size <= dst->capacity ? src->size : dst->capacity;

  memset(dst->buf, 0, dst->size);
  memcpy(dst->buf, src->buf, to_copy);
  dst->size = to_copy;
}

void wapp_str8_copy_to_cstr(char *dst, Str8RO *src, u64 dst_capacity) {
  if (!dst || !src) {
    return;
  }

  u64 to_copy = src->size < dst_capacity ? src->size : dst_capacity - 1;

  memset(dst, 0, dst_capacity);
  memcpy(dst, src->buf, to_copy);
}

void wapp_str8_format(Str8 *dst, const char *format, ...) {
  if (!dst || !format) {
    return;
  }

  va_list args1;
  va_list args2;

  va_start(args1, format);
  va_copy(args2, args1);

  u64 total_size = vsnprintf(NULL, 0, format, args1);
  dst->size      = total_size <= dst->capacity ? total_size : dst->capacity;

  vsnprintf((char *)(dst->buf), dst->capacity, format, args2);

  va_end(args1);
  va_end(args2);
}

i64 wapp_str8_find(Str8RO *str, Str8RO substr) {
  if (!str || substr.size > str->size) {
    return -1;
  }

  // NOTE (Abdelrahman): Uses a while loop instead of a for loop to get rid of
  // MSVC Spectre mitigation warnings
  u64 char_index = 0;
  bool running   = true;
  while (running) {
    const c8 *sub = str->buf + char_index;
    if (memcmp(sub, substr.buf, substr.size) == 0) {
      return char_index;
    }

    ++char_index;
    running = char_index < str->size;
  }

  return -1;
}

i64 wapp_str8_rfind(Str8RO *str, Str8RO substr) {
  if (!str || substr.size > str->size) {
    return -1;
  }

  // NOTE (Abdelrahman): Uses a while loop instead of a for loop to get rid of
  // MSVC Spectre mitigation warnings
  i64 char_index = str->size - substr.size;
  bool running   = true;
  while (running) {
    const c8 *sub = str->buf + char_index;
    if (memcmp(sub, substr.buf, substr.size) == 0) {
      return char_index;
    }

    --char_index;
    running = char_index >= 0;
  }

  return -1;
}

Str8List *wapp_str8_split_with_max(const Allocator *allocator, Str8RO *str, Str8RO *delimiter, i64 max_splits) {
  if (!allocator || !str || !delimiter) {
    return NULL;
  }

  Str8List *output = wapp_mem_allocator_alloc(allocator, sizeof(Str8List));

  if (delimiter->size > str->size) {
    Str8 *full     = wapp_str8_alloc_str8(allocator, str);
    Str8Node *node = wapp_mem_allocator_alloc(allocator, sizeof(Str8Node));
    if (node) {
      node->string = full;
      wapp_str8_list_push_back(output, node);
    }

    goto RETURN_STR8_SPLIT;
  }

  i64 start   = 0;
  i64 end     = 0;
  i64 splits  = 0;
  Str8 *rest = wapp_str8_alloc_str8(allocator, str);
  Str8 *before_str;

  while ((end = wapp_str8_find(rest, *delimiter)) != -1) {
    if (max_splits > 0 && splits >= max_splits) {
      break;
    }

    before_str     = wapp_str8_alloc_substr(allocator, str, start, start + end);
    Str8Node *node = wapp_mem_allocator_alloc(allocator, sizeof(Str8Node));
    if (node) {
      node->string = before_str;
      wapp_str8_list_push_back(output, node);
    }

    wapp_mem_allocator_free(allocator, (void **)&rest, sizeof(Str8));
    rest   = wapp_str8_alloc_substr(allocator, str, start + end + delimiter->size, str->size);
    start += end + delimiter->size;

    ++splits;
  }

  // Ensure the last part of the string after the delimiter is added to the list
  rest           = wapp_str8_alloc_substr(allocator, str, start, str->size);
  Str8Node *node = wapp_mem_allocator_alloc(allocator, sizeof(Str8Node));
  if (node) {
    node->string = rest;
    wapp_str8_list_push_back(output, node);
  }

RETURN_STR8_SPLIT:
  return output;
}

Str8List *wapp_str8_rsplit_with_max(const Allocator *allocator, Str8RO *str, Str8RO *delimiter, i64 max_splits) {
  if (!allocator || !str || !delimiter) {
    return NULL;
  }

  Str8List *output = wapp_mem_allocator_alloc(allocator, sizeof(Str8List));

  if (delimiter->size > str->size) {
    Str8 *full     = wapp_str8_alloc_str8(allocator, str);
    Str8Node *node = wapp_mem_allocator_alloc(allocator, sizeof(Str8Node));
    if (node) {
      node->string = full;
      wapp_str8_list_push_back(output, node);
    }

    goto RETURN_STR8_SPLIT;
  }

  i64 end     = 0;
  i64 splits  = 0;
  Str8 *rest = wapp_str8_alloc_str8(allocator, str);
  Str8 *after_str;

  while ((end = wapp_str8_rfind(rest, *delimiter)) != -1) {
    if (max_splits > 0 && splits >= max_splits) {
      break;
    }

    after_str     = wapp_str8_alloc_substr(allocator, rest, end + delimiter->size, str->size);
    Str8Node *node = wapp_mem_allocator_alloc(allocator, sizeof(Str8Node));
    if (node) {
      node->string = after_str;
      wapp_str8_list_push_front(output, node);
    }

    wapp_mem_allocator_free(allocator, (void **)&rest, sizeof(Str8));
    rest   = wapp_str8_alloc_substr(allocator, rest, 0, end);

    ++splits;
  }

  rest           = wapp_str8_alloc_substr(allocator, str, 0, rest->size);
  Str8Node *node = wapp_mem_allocator_alloc(allocator, sizeof(Str8Node));
  if (node) {
    node->string = rest;
    wapp_str8_list_push_front(output, node);
  }

RETURN_STR8_SPLIT:
  return output;
}

Str8 *wapp_str8_join(const Allocator *allocator, const Str8List *list, Str8RO *delimiter) {
  if (!allocator || !list || !delimiter) {
    return NULL;
  }

  u64 capacity = list->total_size + (delimiter->size * (list->node_count - 1));
  Str8 *output = wapp_str8_alloc_buf(allocator, capacity * 2);

  // NOTE (Abdelrahman): Uses a while loop instead of a for loop to get rid of
  // MSVC Spectre mitigation warnings
  Str8Node *node;
  u64 node_index = 0;
  bool running   = true;
  while (running) {
    node = wapp_str8_list_get(list, node_index);
    wapp_str8_concat_capped(output, node->string);
    if (node_index + 1 < list->node_count) {
      wapp_str8_concat_capped(output, delimiter);
    }

    ++node_index;
    running = node_index < list->node_count;
  }

  return output;
}

Str8Node *wapp_str8_list_get(const Str8List *list, u64 index) {
  if (index >= list->node_count) {
    return NULL;
  }

  Str8Node *output  = NULL;
  Str8Node *current = list->first;
  for (u64 i = 1; i <= index; ++i) {
    current = current->next;
  }

  output = current;

  return output;
}

void wapp_str8_list_push_front(Str8List *list, Str8Node *node) {
  if (!list || !node || !(node->string)) {
    return;
  }

  Str8List node_list = node_to_list(node);

  if (list->node_count == 0) {
    *list = node_list;
    return;
  }

  list->total_size += node_list.total_size;
  list->node_count += node_list.node_count;

  Str8Node *first = list->first;
  if (first) {
    first->prev = node_list.last;
  }

  list->first          = node_list.first;
  node_list.last->next = first;
}

void wapp_str8_list_push_back(Str8List *list, Str8Node *node) {
  if (!list || !node || !(node->string)) {
    return;
  }

  Str8List node_list = node_to_list(node);

  if (list->node_count == 0) {
    *list = node_list;
    return;
  }

  list->total_size += node_list.total_size;
  list->node_count += node_list.node_count;

  Str8Node *last = list->last;
  if (last) {
    last->next = node_list.first;
  }

  list->last            = node_list.last;
  node_list.first->prev = last;
}

void wapp_str8_list_insert(Str8List *list, Str8Node *node, u64 index) {
  if (!list || !node || !(node->string)) {
    return;
  }

  if (index == 0) {
    wapp_str8_list_push_front(list, node);
    return;
  } else if (index == list->node_count) {
    wapp_str8_list_push_back(list, node);
    return;
  }

  Str8Node *dst_node = wapp_str8_list_get(list, index);
  if (!dst_node) {
    return;
  }

  Str8List node_list = node_to_list(node);

  list->total_size += node_list.total_size;
  list->node_count += node_list.node_count;

  Str8Node *prev = dst_node->prev;

  dst_node->prev = node_list.last;
  prev->next     = node_list.first;

  node_list.first->prev = prev;
  node_list.last->next  = dst_node;
}

Str8Node *wapp_str8_list_pop_front(Str8List *list) {
  Str8Node *output = NULL;

  if (!list || list->node_count == 0) {
    goto RETURN_STR8_LIST_POP_FRONT;
  }

  output = list->first;

  if (list->node_count == 1) {
    *list = (Str8List){0};
    goto RETURN_STR8_LIST_POP_FRONT;
  }

  --(list->node_count);
  list->total_size -= output->string->size;
  list->first       = output->next;

  output->prev = output->next = NULL;

RETURN_STR8_LIST_POP_FRONT:
  return output;
}

Str8Node *wapp_str8_list_pop_back(Str8List *list) {
  Str8Node *output = NULL;

  if (!list || list->node_count == 0) {
    goto RETURN_STR8_LIST_POP_BACK;
  }

  output = list->last;

  if (list->node_count == 1) {
    *list = (Str8List){0};
    goto RETURN_STR8_LIST_POP_BACK;
  }

  --(list->node_count);
  list->total_size -= output->string->size;
  list->last        = output->prev;

  output->prev = output->next = NULL;

RETURN_STR8_LIST_POP_BACK:
  return output;
}

Str8Node *wapp_str8_list_remove(Str8List *list, u64 index) {
  Str8Node *output = NULL;
  if (!list) {
    goto RETURN_STR8_LIST_REMOVE;
  }

  if (index == 0) {
    output = wapp_str8_list_pop_front(list);
    goto RETURN_STR8_LIST_REMOVE;
  } else if (index == list->node_count) {
    output = wapp_str8_list_pop_back(list);
    goto RETURN_STR8_LIST_REMOVE;
  }

  output = wapp_str8_list_get(list, index);
  if (!output) {
    goto RETURN_STR8_LIST_REMOVE;
  }

  output->prev->next = output->next;
  output->next->prev = output->prev;

  --(list->node_count);
  list->total_size -= output->string->size;

  output->prev = output->next = NULL;

RETURN_STR8_LIST_REMOVE:
  return output;
}

void wapp_str8_list_empty(Str8List *list) {
  if (!list) {
    return;
  }

  u64 count = list->node_count;
  for (u64 i = 0; i < count; ++i) {
    wapp_str8_list_pop_back(list);
  }
}

internal Str8List node_to_list(Str8Node *node) {
  Str8List output = {.first = node, .last = node, .total_size = node->string->size, .node_count = 1};

  while (output.first->prev != NULL) {
    output.total_size += output.first->prev->string->size;
    output.first = output.first->prev;
    ++(output.node_count);
  }

  while (output.last->next != NULL) {
    output.total_size += output.last->next->string->size;
    output.last = output.last->next;
    ++(output.node_count);
  }

  return output;
}