#include "util.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <assert.h>

void wm_treenode_init(TreeNode *node, NodeType type, TreeNode *parent)
{
    assert(node);

    node->parent = parent;
    node->type = type;
    node->children = (NodeArray) {0};
    wm_nodearray_new();
}

bool wm_treenode_is_empty(TreeNode *node)
{
    return (node->children.size == 0);
}

void wm_treenode_split_space(TreeNode *node, Rect *ret1, Rect *ret2)
{
    assert(node);
    assert(ret1);
    assert(ret2);

    switch (node->type) {
    NODE_VERTICAL:
        ret1->w = node->pos.w / 2;
        ret1->h = node->pos.h;
        ret2->w = node->pos.w / 2;
        ret2->h = node->pos.h;

        ret1->x = node->pos.x;
        ret1->y = node->pos.y;

        ret2->x = node->pos.x + ret1->w + 1;
        ret2->y = node->pos.y;
        return;
    NODE_HORIZONTAL:
        ret1->w = node->pos.w;
        ret1->h = node->pos.h / 2;
        ret2->w = node->pos.w;
        ret2->h = node->pos.h / 2;

        ret1->x = node->pos.x;
        ret1->y = node->pos.y;

        ret2->x = node->pos.x;
        ret2->y = node->pos.y + ret1->h + 1;
        return;
    default:
        fprintf(stderr, "wm: unhandled treenode type %d in %s, exiting.\n",
                node->type, __func__);
        abort();
    }
}

void wm_treenode_remove_client(TreeNode *root, Client *client)
{
    TreeNode *client_node = wm_treenode_ptr_find_client_node(root, client);

    if (client_node->parent == NULL) {
        client_node->client = NULL;
        wm_nodearray_clear(&client_node->children);
        return;
    }

    TreeNode *parent = client_node->parent;
    parent->type = NODE_CLIENT;
    Client *other_client = NULL;

    for (size_t i = 0; i < parent->children.size; i++) {
        if (parent->children.nodes[i].type == NODE_CLIENT &&
            parent->children.nodes[i].client != client) {
            other_client = parent->children.nodes[i].client;
        }
    }

    assert(other_client);
    parent->client = other_client;
    wm_nodearray_clear(&parent->children);
}

NodeArray wm_nodearray_new()
{
    NodeArray arr;

    arr.capacity = 10;
    arr.nodes = calloc(arr.capacity, sizeof(TreeNode));
    arr.size = 0;

    return arr;
}

void wm_nodearray_push(NodeArray *arr, TreeNode node)
{
    assert(arr);

    arr->size++;

    if (arr->size >= arr->capacity) {
        TreeNode* temp = calloc(arr->capacity, sizeof(TreeNode));
        memcpy(temp, arr->nodes, arr->capacity * sizeof(TreeNode));
        free(arr->nodes);
        size_t old_capacity = arr->capacity;
        arr->capacity = old_capacity * 2;

        arr->nodes = calloc(arr->capacity, sizeof(TreeNode));
        memcpy(arr->nodes, temp, old_capacity);

        free(temp);
        arr->nodes[arr->size - 1] = node;
        return;
    }

    arr->nodes[arr->size - 1] = node;
}

bool wm_nodearray_pop(NodeArray *arr, TreeNode *ret)
{
    assert(arr);
    assert(ret);
    
    if (arr->size == 0)
        return false;

    *ret = arr->nodes[arr->size - 1];
    arr->size--;
    return true;
}

bool wm_nodearray_pop_front(NodeArray *arr, TreeNode *ret)
{
    assert(arr);
    assert(ret);

    if (arr->size == 0)
        return false;

    *ret = arr->nodes[0];
    arr->size--;
    memcpy(arr->nodes, arr->nodes+1, arr->size * sizeof(TreeNode));

    return true;
}

void wm_nodearray_clear(NodeArray *arr)
{
    arr->size = 0;
    memset(arr->nodes, 0, arr->capacity * sizeof(TreeNode));
}

bool wm_nodearray_remove(NodeArray *arr, size_t index)
{
    assert(arr);

    if (!arr)
        return false;

    if (index >= arr->size)
        return false;

    if (index == arr->size - 1) {
        arr->size--;
        return true;
    }

    memmove(arr->nodes + index, arr->nodes + index+1, arr->size-1 * sizeof(TreeNode));

    return true;
}

void wm_nodearray_free(NodeArray *arr)
{
    assert(arr);

    arr->capacity = 0;
    arr->size = 0;
    free(arr->nodes);
}

// breadth-first search for finding client nodes in tree
ClientArray wm_nodearray_find_clients(TreeNode *root)
{
    assert(root);

    /* storing visited nodes in an array and searching is
       not optimal, but this is not a performance-critical function. */
    NodeArray visited = wm_nodearray_new();
    wm_nodearray_push(&visited, *root);

    NodeArray queue = wm_nodearray_new();
    ClientArray clients = wm_clientarray_new();

    wm_nodearray_push(&queue, *root);

    while (queue.size > 0) {
        TreeNode node;

        if (!wm_nodearray_pop_front(&queue, &node))
            goto ret;

        if (node.type == NODE_CLIENT)
            wm_clientarray_push(&clients, *node.client);

        for (int i = 0; i < node.children.size; i++) {
            TreeNode child_node = node.children.nodes[i];
            bool _visited = false;

            for (int j = 0; j < visited.size; j++) {
                // this comparison may not work
                if (visited.nodes[j].children.nodes == child_node.children.nodes) {
                    _visited = true;
                    break;
                }
            }

            if (!_visited) {
                wm_nodearray_push(&visited, child_node);
                wm_nodearray_push(&queue, child_node);
            }
        }
    }

ret:
    wm_nodearray_free(&visited);
    wm_nodearray_free(&queue);
    return clients;
}

PtrArray wm_ptrarray_new()
{
    PtrArray arr;

    arr.capacity = 10;
    arr.ptrs = calloc(arr.capacity, sizeof(TreeNode));
    arr.size = 0;

    return arr;
}

void wm_ptrarray_push(PtrArray *arr, void* ptr)
{
    assert(arr);

    arr->size++;

    if (arr->size >= arr->capacity) {
        TreeNode* temp = calloc(arr->capacity, sizeof(TreeNode));
        memcpy(temp, arr->ptrs, arr->capacity * sizeof(TreeNode));
        free(arr->ptrs);
        size_t old_capacity = arr->capacity;
        arr->capacity = old_capacity * 2;

        arr->ptrs = calloc(arr->capacity, sizeof(TreeNode));
        memcpy(arr->ptrs, temp, old_capacity);

        free(temp);
        arr->ptrs[arr->size - 1] = ptr;
        return;
    }

    arr->ptrs[arr->size - 1] = ptr;
}

bool wm_ptrarray_pop(PtrArray *arr, void **ret)
{
    assert(arr);
    assert(ret);

    if (arr->size == 0)
        return false;

    *ret = arr->ptrs[arr->size - 1];
    arr->size--;
    return true;
}

bool wm_ptrarray_pop_front(PtrArray *arr, void **ret)
{
    assert(arr);
    assert(ret);

    if (arr->size == 0)
        return false;

    *ret = arr->ptrs[0];
    arr->size--;
    memcpy(arr->ptrs, arr->ptrs+1, arr->size * sizeof(TreeNode));

    return true;
}

bool wm_ptrarray_remove(PtrArray *arr, size_t index)
{
    assert(arr);

    if (!arr)
        return false;

    if (index >= arr->size)
        return false;

    if (index == arr->size - 1) {
        arr->size--;
        return true;
    }

    memmove(arr->ptrs + index, arr->ptrs + index+1, arr->size-1 * sizeof(TreeNode));

    return true;
}

void wm_ptrarray_free(PtrArray *arr)
{
    assert(arr);

    arr->capacity = 0;
    arr->size = 0;
    free(arr->ptrs);
}

TreeNode* wm_treenode_ptr_find_focused_client_node(TreeNode *root)
{
    assert(root);

    TreeNode *ret = NULL;

    NodeArray visited = wm_nodearray_new();
    wm_nodearray_push(&visited, *root);

    PtrArray queue = wm_ptrarray_new();

    wm_ptrarray_push(&queue, root);

    while (queue.size > 0) {
        TreeNode *node;

        if (!wm_ptrarray_pop_front(&queue, (void**)&node))
            goto ret;

        if (node->type == NODE_CLIENT && node->client->focused) {
            ret = node;
            goto ret;
        }

        for (int i = 0; i < node->children.size; i++) {
            TreeNode child_node = node->children.nodes[i];
            bool _visited = false;

            for (int j = 0; j < visited.size; j++) {
                if (visited.nodes[j].children.nodes == child_node.children.nodes) {
                    _visited = true;
                    break;
                }
            }

            if (!_visited) {
                wm_nodearray_push(&visited, child_node);
                wm_ptrarray_push(&queue, &child_node);
            }
        }
    }

ret:
    wm_nodearray_free(&visited);
    wm_ptrarray_free(&queue);
    return ret;
}

TreeNode* wm_treenode_ptr_find_client_node(TreeNode *root, Client *client)
{
    assert(root);

    TreeNode *ret = NULL;

    NodeArray visited = wm_nodearray_new();
    wm_nodearray_push(&visited, *root);

    PtrArray queue = wm_ptrarray_new();

    wm_ptrarray_push(&queue, root);

    while (queue.size > 0) {
        TreeNode *node;

        if (!wm_ptrarray_pop_front(&queue, (void**)&node))
            goto ret;

        if (node->client == client) {
            ret = node;
            goto ret;
        }

        for (int i = 0; i < node->children.size; i++) {
            TreeNode child_node = node->children.nodes[i];
            bool _visited = false;

            for (int j = 0; j < visited.size; j++) {
                if (visited.nodes[j].children.nodes == child_node.children.nodes) {
                    _visited = true;
                    break;
                }
            }

            if (!_visited) {
                wm_nodearray_push(&visited, child_node);
                wm_ptrarray_push(&queue, &child_node);
            }
        }
    }

ret:
    wm_nodearray_free(&visited);
    wm_ptrarray_free(&queue);
    return ret;
}

ClientArray wm_clientarray_new()
{
    ClientArray arr;

    arr.capacity = 10;
    arr.clients = calloc(arr.capacity, sizeof(TreeNode));
    arr.size = 0;

    return arr;
}

void wm_clientarray_push(ClientArray *arr, Client node)
{
    assert(arr);

    arr->size++;

    if (arr->size >= arr->capacity) {
        Client* temp = calloc(arr->capacity, sizeof(Client));
        memcpy(temp, arr->clients, arr->capacity * sizeof(Client));
        free(arr->clients);
        size_t old_capacity = arr->capacity;
        arr->capacity = old_capacity * 2;

        arr->clients = calloc(arr->capacity, sizeof(Client));
        memcpy(arr->clients, temp, old_capacity);

        free(temp);
        arr->clients[arr->size - 1] = node;
        return;
    }

    arr->clients[arr->size - 1] = node;
}

bool wm_clientarray_pop(ClientArray *arr, Client *ret)
{
    assert(arr);

    if (!ret || !arr)
        return false;

    if (arr->size == 0)
        return false;

    *ret = arr->clients[arr->size - 1];
    arr->size--;
    return true;
}

bool wm_clientarray_remove(ClientArray *arr, size_t index)
{
    assert(arr);

    if (!arr)
        return false;

    if (index >= arr->size)
        return false;

    if (index == arr->size - 1) {
        arr->size--;
        return true;
    }

    memmove(arr->clients + index, arr->clients + index+1, arr->size-1 * sizeof(Client));

    return true;
}

void wm_clientarray_free(ClientArray *arr)
{
    assert(arr);

    arr->capacity = 0;
    arr->size = 0;
    free(arr->clients);
}