#include "common.h"
#include <assert.h>
#include <ctype.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static bool ch_is_punctuation(char ch)
{
    return strchr(",.!?'", ch);
}

static bool ch_is_letters(char ch)
{
    return isalpha(ch);
}

size_t next_word(char line[], size_t line_size, FILE *file, int *ch)
{
    size_t len = 0;

    if (!*ch)
        *ch = fgetc(file);

    if (*ch != EOF && !ch_is_punctuation(*ch) && !ch_is_letters(*ch)) {
        assert(len < line_size - 1);
        line[len] = ' ';
        len += 1;

        int newlines = 0;

        while (*ch != EOF && !ch_is_punctuation(*ch) && !ch_is_letters(*ch)) {
            newlines += *ch == '\n';
            *ch = fgetc(file);
        }

        if (newlines >= 2) {
            line[len - 1] = '\n';
            return len;
        }
    }

    bool matched = false;

    while (*ch != EOF && ch_is_punctuation(*ch)) {
        matched = true;

        assert(len < line_size - 1);
        line[len] = *ch;
        len += 1;

        *ch = fgetc(file);
    }
    if (matched)
        return len;

    while (*ch != EOF && ch_is_letters(*ch)) {
        matched = true;

        assert(len < line_size - 1);
        line[len] = tolower(*ch);
        len += 1;

        *ch = fgetc(file);
    }

    return len;
}

void array_push(
    void **data, size_t *capacity, size_t *size, void *elem, size_t elem_size)
{
    if (!*capacity) {
        *capacity = 4;
    }
    if (!*data) {
        *data = malloc(*capacity * elem_size);
        *size = 0;
    }
    if (*size >= *capacity) {
        *capacity *= 2;
        *data = realloc(*data, *capacity * elem_size);
    }
    memcpy((unsigned char *)*data + *size * elem_size, elem, elem_size);
    *size += 1;
}

void array_push_n(void **data,
    size_t *capacity,
    size_t *size,
    void *elems,
    size_t elem_size,
    size_t elem_count)
{
    for (size_t i = 0; i < elem_count; ++i) {
        array_push(data,
            capacity,
            size,
            (unsigned char *)elems + elem_size * i,
            elem_size);
    }
}

typedef int (*CmpFn)(const void *, const void *);

static void *binary_search_internal(void *data,
    size_t begin,
    size_t end,
    const void *entry,
    size_t entry_size,
    size_t *out_idx,
    CmpFn cmp)
{
    if (begin == end) {
        if (out_idx)
            *out_idx = begin;
        return NULL;
    }
    size_t middle_idx = (end - begin) / 2 + begin;
    void *middle_entry = (unsigned char *)data + entry_size * middle_idx;
    int cmp_result = cmp(entry, middle_entry);
    if (cmp_result == 0) {
        if (out_idx)
            *out_idx = middle_idx;
        return middle_entry;
    }
    if (cmp_result < 0)
        end = middle_idx;
    else
        begin = middle_idx + 1;

    return binary_search_internal(
        data, begin, end, entry, entry_size, out_idx, cmp);
}

void *binary_search(void *data,
    size_t count,
    const void *entry,
    size_t entry_size,
    size_t *out_idx,
    CmpFn cmp)
{
    return binary_search_internal(
        data, 0, count, entry, entry_size, out_idx, cmp);
}

typedef struct RbTree RbTree;

RbTree *rbtree_new(size_t key_size, size_t val_size, CmpFn cmp);
void rbtree_free(RbTree *t);
void rbtree_insert(RbTree *t, const void *key, void *val);
void *rbtree_get(RbTree *t, const void *key);

typedef enum { Rb_Black, Rb_Red } RbColor;
typedef enum { Rb_Left, Rb_Right } RbDir;

typedef struct RbNode RbNode;

struct RbNode {
    RbNode *parent;
    union {
        struct {
            RbNode *left;
            RbNode *right;
        };
        RbNode *child[2];
    };
    uint8_t color;
    int8_t padding[7];
    uint8_t key_and_val[];
};

struct RbTree {
    size_t key_size;
    size_t val_size;
    CmpFn cmp;
    size_t node_key_size;
    size_t node_val_size;
    size_t size_of_node;
    RbNode *root;
};

RbTree *rbtree_new(size_t key_size, size_t val_size, CmpFn cmp)
{
    size_t node_key_size = (key_size + 7ul) & ~7ul;
    size_t node_val_size = (key_size + 7ul) & ~7ul;
    size_t size_of_node = sizeof(RbNode) + node_key_size + node_val_size;

    RbTree *t = malloc(sizeof(RbTree));
    *t = (RbTree) {
        .key_size = key_size,
        .val_size = val_size,
        .cmp = cmp,
        .node_key_size = node_key_size,
        .node_val_size = node_val_size,
        .size_of_node = size_of_node,
        .root = NULL,
    };
    return t;
}

static void t_free_node(RbNode *node)
{
    if (!node)
        return;

    t_free_node(node->child[0]);
    t_free_node(node->child[1]);

    free(node);
}

void rbtree_free(RbTree *t)
{
    t_free_node(t->root);
    free(t);
}

static RbNode *t_find_node(RbTree *t,
    RbNode *node,
    const void *key,
    RbNode **out_parent,
    RbDir *out_dir)
{
    int cmp_result = t->cmp(key, node->key_and_val);

    if (cmp_result == 0)
        return node;

    RbDir dir = cmp_result < 0 ? Rb_Left : Rb_Right;
    RbNode *next = node->child[dir];

    if (next)
        return t_find_node(t, next, key, out_parent, out_dir);

    if (out_dir)
        *out_dir = dir;
    if (out_parent)
        *out_parent = node;

    return NULL;
}

static RbNode *t_node_new(
    RbTree *t, RbNode *parent, RbColor color, const void *key, void *val)
{
    RbNode *node = malloc(t->size_of_node);
    *node = (RbNode) {
        .parent = parent,
        .child = { NULL, NULL },
        .color = color,
    };

    memcpy(node->key_and_val, key, t->key_size);
    memcpy(node->key_and_val + t->node_key_size, val, t->val_size);
    return node;
}

static RbNode *t_rotate_subtree(RbTree *t, RbNode *sub, RbDir dir)
{
    RbNode *sub_parent = sub->parent;
    RbNode *new_root = sub->child[1 - dir];
    RbNode *new_child = new_root->child[dir];

    sub->child[1 - dir] = new_child;

    if (new_child)
        new_child->parent = sub;

    new_root->child[dir] = sub;

    new_root->parent = sub_parent;
    sub->parent = new_root;
    if (sub_parent)
        sub_parent->child[sub == sub_parent->right] = new_root;
    else
        t->root = new_root;

    return new_root;
}

void rbtree_insert(RbTree *t, const void *key, void *val)
{
    // https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Insertion

    RbNode *node, *parent = NULL;
    RbDir dir = Rb_Left;

    if (t->root && (node = t_find_node(t, t->root, key, &parent, &dir))) {
        memcpy(node->key_and_val + t->node_key_size, val, t->val_size);
        return;
    }

    node = t_node_new(t, parent, Rb_Red, key, val);

    if (!t->root) {
        t->root = node;
        return;
    }

    parent->child[dir] = node;

    do {
        if (parent->color == Rb_Black)
            // case 1
            return;

        RbNode *grandparent = parent->parent;

        if (!grandparent) {
            // case 4
            parent->color = Rb_Black;
            return;
        }

        dir = parent == grandparent->left ? Rb_Left : Rb_Right;
        RbNode *uncle = grandparent->child[1 - dir];

        if (!uncle || uncle->color == Rb_Black) {
            if (node == parent->child[1 - dir]) {
                // case 5
                t_rotate_subtree(t, parent, dir);
                node = parent;
                parent = grandparent->child[dir];
            }

            // case 6
            t_rotate_subtree(t, grandparent, 1 - dir);
            parent->color = Rb_Black;
            grandparent->color = Rb_Red;
            return;
        }

        // case 2
        parent->color = Rb_Black;
        uncle->color = Rb_Black;
        grandparent->color = Rb_Red;
        node = grandparent;

    } while ((parent = node->parent));

    // case 3
}

void *rbtree_get(RbTree *t, const void *key)
{
    RbNode *node;
    if (!t->root || !(node = t_find_node(t, t->root, key, NULL, NULL)))
        return NULL;

    return node->key_and_val + t->node_key_size;
}

static bool t_red_node_children_valid(RbNode *node)
{
    if (!node)
        return true;
    if (node->parent && node->color == Rb_Red)
        return node->parent->color != Rb_Red;
    return t_red_node_children_valid(node->left)
        && t_red_node_children_valid(node->right);
}

static int t_walk_count_black_nodes(RbNode *node)
{
    if (!node)
        return 1;

    int left = t_walk_count_black_nodes(node->left);
    int right = t_walk_count_black_nodes(node->right);

    if (left != right)
        return 0;

    return left + (node->color == Rb_Black);
}

static bool rbtree_valid(RbTree *t)
{
    return t_red_node_children_valid(t->root)
        && t_walk_count_black_nodes(t->root) != 0;
}

static int int_cmp(const int *a, const int *b)
{
    return *a - *b;
}

void test_common(void)
{
    RbTree *t = rbtree_new(sizeof(int), sizeof(int), (CmpFn)int_cmp);

    int keys[] = { 3, 7, 5, 1, 4, 6 };
    int vals[] = { 0, 2, 4, 6, 8, 10 };
    int *result;

    result = rbtree_get(t, &keys[0]);
    assert(!result);

    for (size_t i = 0; i < sizeof(keys) / sizeof(keys[0]); ++i) {
        rbtree_insert(t, &keys[i], &vals[i]);
        assert(rbtree_valid(t));
    }

    for (size_t i = 0; i < sizeof(keys) / sizeof(keys[0]); ++i) {
        assert(*(int *)rbtree_get(t, &keys[i]) == vals[i]);
    }

    rbtree_free(t);
}