#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>

#include "trie_node.h"

/**
 * Allocate and initialize a new TrieInnerNode representing a given
 * character.
 *
 * @param c character to represent
 * @return pointer to newly allocated struct
 */
TrieInnerNode *tinode_init(char c) {
  TrieInnerNode *node = calloc(1, sizeof(TrieInnerNode));
  node->key = c;

  return node;
}

/**
 * Allocate and initialize a new TrieNode.
 *
 * @return pointer to newly allocated struct
 */
TrieNode *tnode_init() {
  TrieNode *node = malloc(sizeof(TrieNode));

  node->tree_size = 0;
  node->string_len = 0;
  node->represents = false;

  return node;
}

/**
 * Free a TrieInnerNode and its underlying tree structure. This should
 * usually only be called on the root of a binary tree to free the entire
 * structure.
 *
 * @param node node whose tree to free
 */
void tinode_free_cascade(TrieInnerNode *node) {
  if (node->left != NULL) {
    tinode_free_cascade(node->left);
  }

  if (node->right != NULL) {
    tinode_free_cascade(node->right);
  }

  if (node->next != NULL) {
    tnode_free(node->next);
  }

  free(node);
}

/**
 * Free a TrieNode and its underlying tree structure.
 *
 * @param node node to free
 */
void tnode_free(TrieNode *node) {
  if (node->tree_size > 0) {
    tinode_free_cascade(node->tree);
  }

  // TODO properly free entry
  /* if (node->payload != NULL) { */
  /*     free(node->payload); */
  /* } */

  free(node);
}

/**
 * This function performs a lookup in the underlying binary tree of the given
 * TrieNode. If found, the return value is a pointer to the memory
 * location where the TrieInnerNode representing the given character
 * stores its `next` field. If not found, the return value is NULL, unless
 * `create` is true.
 *
 * NOTE: a non-NULL return value does not mean that the dereferenced value is
 * also not NULL. In particular, if `create` is set to true and the function had
 * to create the new node, the dereferenced value will always be NULL.
 *
 * @param node node to perform lookup in. If node is a full leaf, the return
 * value will always be NULL, regardless of the value of create.
 * @param create whether to create the TrieInnerNode if it isn't present
 * yet. If this is set to true, the function will never return NULL unless the
 * node represents a leaf with a string, because the struct and therefore the
 * address is created if it doesn't exist yet.
 */
TrieNode **tnode_search(TrieNode *node, const char c, bool create) {
  // It can happen that the node has no initialized root yet
  if (node->tree_size == 0) {
    if (create) {
      node->tree_size++;
      node->tree = tinode_init(c);

      return &node->tree->next;
    }

    return NULL;
  }

  TrieInnerNode *parent = node->tree;
  TrieInnerNode *child;

  // Iterate through the tree until we either find the character or realize it's
  // not present in the tree
  // FIXME don't use while (1)
  while (1) {
    if (parent->key == c) {
      return &parent->next;
    } else if (c < parent->key) {
      child = parent->left;
    } else {
      child = parent->right;
    }

    if (child == NULL) {
      break;
    }

    parent = child;
  };

  // child is NULL, meaning the character isn't in the binary tree yet.

  // If create is true, we create the new node so that we can still return a
  // non-NULL pointer.
  if (create) {
    TrieInnerNode *new_node = tinode_init(c);

    if (c < parent->key) {
      parent->left = new_node;
    } else {
      parent->right = new_node;
    }

    node->tree_size++;

    return &new_node->next;
  }

  return NULL;
}

/**
 * Split a remaining string leaf node in two. This function assumes it receives
 * a full leaf as its input.
 *
 * @param node node to split
 */
/* void tnode_split(TrieNode *node) { */
/*   TrieNode *new_node = tnode_init(); */
/*   char key = node->ptr.string[0]; */

/*   // There's a chance the remaining string was only 1 character, meaning the
 * new */
/*   // node doesn't have to store a string */
/*   if (node->ptr.string[1] != DELIMITER) { */
/*     tnode_set_string(new_node, node->ptr.string + 1); */
/*   } else { */
/*     new_node->type = 1; */
/*   } */

/*   new_node->entry = node->entry; */

/*   node->type = 0; */
/*   node->size = 0; */
/*   node->entry = NULL; */

/*   free(node->ptr.string); */
/*   node->ptr.string = NULL; */

/*   // Initialize node's binary tree with the correct character */
/*   TrieNode **node_ptr = tnode_search(node, key, true); */
/*   *node_ptr = new_node; */
/* } */

/*
 * Remove the given character from a TrieInnerNode's subtree. The
 * function assumes the character is indeed in the subtree.
 */
void tinode_remove(TrieInnerNode *node, const char c) {
  TrieInnerNode **to_remove_ptr = &node;

  // We use pointers to pointers here so we can later free the removed node
  // without having to know what its parent is
  while ((*to_remove_ptr)->key != c) {
    to_remove_ptr = (c < (*to_remove_ptr)->key) ? &(*to_remove_ptr)->left
                                                : &(*to_remove_ptr)->right;
  };

  // If the node isn't a leaf, we have to replace it with another
  if ((*to_remove_ptr)->left != NULL || (*to_remove_ptr)->right != NULL) {
    TrieInnerNode *to_replace = *to_remove_ptr;

    // Replace with its only right child
    if (to_replace->left == NULL) {
      TrieInnerNode *to_remove = to_replace->right;

      to_replace->key = to_remove->key;
      to_replace->next = to_remove->next;
      to_replace->left = to_remove->left;
      to_replace->right = to_remove->right;

      free(to_remove);
    }
    // Replace with its only left child
    else if (to_replace->right == NULL) {
      TrieInnerNode *to_remove = to_replace->left;

      to_replace->key = to_remove->key;
      to_replace->next = to_remove->next;
      to_replace->left = to_remove->left;
      to_replace->right = to_remove->right;

      free(to_remove);
    }
    // Node has two children, so replace with successor
    else {
      TrieInnerNode *to_remove_parent = to_replace;
      TrieInnerNode *to_remove = to_replace->right;

      while (to_remove->left != NULL) {
        to_remove_parent = to_remove;
        to_remove = to_remove->left;
      }

      to_replace->key = to_remove->key;
      to_replace->next = to_remove->next;

      if (to_remove_parent != to_replace) {
        to_remove_parent->left = to_remove->right;
      } else {
        to_remove_parent->right = to_remove->right;
      }

      free(to_remove);
    }
  }
  // We're the leaf, so we free ourselves
  else {
    free(*to_remove_ptr);
    *to_remove_ptr = NULL;
  }
}

/**
 * Remove the given character from a TrieNode, respecting the rules
 * of a binary search tree. This function assumes the character is in the search
 * tree.
 *
 * @param node node to remove character from
 * @param c character to remove
 */
void tnode_remove(TrieNode *node, const char c) {
  tinode_remove(node->tree, c);

  node->tree_size--;
}