#include "common.c"
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>

/**
 * Represents a node of the binary tree contained within each non-leaf
 * TernaryTrieNode.
 */
typedef struct ttinode {
  struct ttinode *left;
  struct ttinode *right;
  struct ttnode *next;
  char key;
} TernaryTrieInnerNode;

/**
 * Represents a node inside a TernaryTrie. A node can be in one of three states:
 * - Internal node: a node that's part of a path to a leaf node. This node will
 *   always have a size greater than one, and an initialized root.
 * - Leaf: a node solely used to represent a string ending there. Its size is 0,
 *   its ptr is unitialized and represents is true.
 * - Full leaf: a leaf node that contains a string. This occurs when a string is
 *   added whose path is not fully in the tree yet, causing its remaining suffix
 *   to be stored as a single node. Its size will be zero, represents its true,
 *   and its string pointer is initialized.
 */
typedef struct ttnode {
  union {
    TernaryTrieInnerNode *root;
    char *string;
  } ptr;
  // What type of node this is
  // 0: regular non-representing node
  // 1: regular representing node
  // 2: full leaf
  uint8_t type;
  // Dependent on type
  // 0, 1: size of underlying binary tree
  // 2: length of string
  uint8_t size;
} TernaryTrieNode;

// Required for recursively freeing tree structure
void ttnode_free(TernaryTrieNode *node);

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

  return node;
}

/**
 * Allocate and initialize a new TernaryTrieNode.
 *
 * @return pointer to newly allocated struct
 */
TernaryTrieNode *ttnode_init() { return calloc(1, sizeof(TernaryTrieNode)); }

/**
 * Free a TernaryTrieInnerNode 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 ttinode_free_cascade(TernaryTrieInnerNode *node) {
  if (node->left != NULL) {
    ttinode_free_cascade(node->left);
  }

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

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

  free(node);
}

/**
 * Free a TernaryTrieNode and its underlying tree structure.
 *
 * @param node node to free
 */
void ttnode_free(TernaryTrieNode *node) {
  if (node->type == 2) {
    free(node->ptr.string);
  } else if (node->size != 0) {
    ttinode_free_cascade(node->ptr.root);
  }

  free(node);
}

/**
 * Add the string to the given node & set its type accordingely.
 *
 * @param node node to add string to
 * @param string string to add
 */
void ttnode_set_string(TernaryTrieNode *node, const char *string) {
  node->type = 2;
  node->size = strlen(string);
  node->ptr.string = my_strdup(string);
}

/**
 * This function performs a lookup in the underlying binary tree of the given
 * TernaryTrieNode. If found, the return value is a pointer to the memory
 * location where the TernaryTrieInnerNode 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 TernaryTrieInnerNode 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.
 */
TernaryTrieNode **ttnode_search(TernaryTrieNode *node, const char c,
                                bool create) {
  // Full leafs will always return NULL
  if (node->type == 2) {
    return NULL;
  }

  // It can happen that the node has no initialized root yet
  if (node->size == 0) {
    if (create) {
      node->size++;
      node->ptr.root = ttinode_init(c);

      return &node->ptr.root->next;
    }

    return NULL;
  }

  TernaryTrieInnerNode *parent = node->ptr.root;
  TernaryTrieInnerNode *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) {
    TernaryTrieInnerNode *new_node = ttinode_init(c);

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

    node->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 ttnode_split(TernaryTrieNode *node) {
  TernaryTrieNode *new_node = ttnode_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) {
    ttnode_set_string(new_node, node->ptr.string + 1);
  } else {
    new_node->type = 1;
  }

  node->type = 0;
  node->size = 0;

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

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

/*
 * Remove the given character from a TernaryTrieInnerNode's subtree. The
 * function assumes the character is indeed in the subtree.
 */
void ttinode_remove(TernaryTrieInnerNode *node, const char c) {
  TernaryTrieInnerNode **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) {
    TernaryTrieInnerNode *to_replace = *to_remove_ptr;

    // Replace with its only right child
    if (to_replace->left == NULL) {
      TernaryTrieInnerNode *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) {
      TernaryTrieInnerNode *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 {
      TernaryTrieInnerNode *to_remove_parent = to_replace;
      TernaryTrieInnerNode *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 TernaryTrieNode, 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 ttnode_remove(TernaryTrieNode *node, const char c) {
  ttinode_remove(node->ptr.root, c);

  node->size--;

  if (node->size == 0) {
    node->ptr.root = NULL;
  }
}