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

TreeNode wm_treenode_new(NodeType type, TreeNode *parent)
{
    DEBUG_PRINT("%s\n", __func__);
    TreeNode node;

    node.parent = parent;
    node.type = type;
    node.children = wm_nodearray_new();

    node_id++;
    node.id = node_id;
    node.client = NULL;

    return node;
}

bool wm_treenode_is_empty(TreeNode *node)
{
    DEBUG_PRINT("%s\n", __func__);
    return (node->children.size == 0);
}

void wm_treenode_split_space(TreeNode *node, Rect *ret1, Rect *ret2)
{
    DEBUG_PRINT("%s\n", __func__);
    assert(node);
    assert(ret1);
    assert(ret2);

    switch (node->type) {
    case 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;
        ret2->y = node->pos.y;
        return;
    case 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;
        return;
    default:
        fprintf(stderr, "wm: unhandled treenode type %d in %s, exiting.\n",
                node->type, __func__);
        abort();
    }
}

void wm_node_type_to_str(NodeType type, char *buf, size_t bufsize)
{
    switch (type) {
    case NODE_CLIENT:
        strncpy(buf, "NODE_CLIENT", bufsize);
        return;
    case NODE_VERTICAL:
        strncpy(buf, "NODE_VERTICAL", bufsize);
        return;
    case NODE_HORIZONTAL:
        strncpy(buf, "NODE_HORIZONTAL", bufsize);
        return;
    case NODE_TAB:
        strncpy(buf, "NODE_TAB", bufsize);
        return;
    case NODE_COUNT:
        return;
    }
}

void wm_tree_to_DOT(TreeNode *root, const char *filename)
{
    DEBUG_PRINT("%s\n", __func__);

    FILE *fp = fopen(filename, "w");
    assert(fp);

    fprintf(fp, "digraph {\n");

    PtrArray all_nodes = wm_all_nodes_to_ptrarray(root);
    NodeArray printed_nodes = wm_nodearray_new();

    const int bufsize = 100;
    char buf1[bufsize];
    char buf2[bufsize];

    for (size_t i = 0; i < all_nodes.size; i++) {
        TreeNode *node = all_nodes.ptrs[i];

        bool contains = false;
        for (size_t j = 0; j < printed_nodes.size; j++) {
            if (printed_nodes.nodes[i].id == node->id)
                contains = true;
        }

        if (contains)
            continue;

        for (size_t j = 0; j < node->children.size; j++) {
            memset(buf1, 0, bufsize);
            memset(buf2, 0, bufsize);

            wm_node_type_to_str(node->type, buf1, bufsize);
            wm_node_type_to_str(node->children.nodes[j].type, buf2, bufsize);

            fprintf(fp, "_%d_%s -> _%d_%s;\n", node->id, buf1,
                    node->children.nodes[j].id, buf2);
        }
        wm_nodearray_push(&printed_nodes, *node);
    }

    fprintf(fp, "}");

    fflush(fp);
    fclose(fp);

    wm_nodearray_free(&printed_nodes);
    wm_ptrarray_free(&all_nodes);
}

int wm_get_node_index(TreeNode *parent, unsigned int node_id)
{
    for (size_t i = 0; i < parent->children.size; i++) {
        if (node_id == parent->children.nodes[i].id)
            return i;
    }

    return -1;
}

void wm_treenode_remove_node(TreeNode *root, unsigned int node_id)
{
    DEBUG_PRINT("%s\n", __func__);
    assert(root);

    PtrArray all_nodes = wm_all_nodes_to_ptrarray(root);

    for (size_t i = 0; i < all_nodes.size; i++) {
        TreeNode *node = all_nodes.ptrs[i];
        if (node->id == node_id) {
            TreeNode *parent = node->parent;
            int index = wm_get_node_index(parent, node_id);
            assert(index >= 0);

            wm_nodearray_remove(&parent->children, index);
        }
    }
}

TreeNode* wm_treenode_split_get_sibling(TreeNode *node)
{
    for (size_t i = 0; i < node->parent->children.size; i++) {
        TreeNode *sibling_node = &node->parent->children.nodes[i];
        if (sibling_node->id != node->id) {
            return sibling_node;
        }
    }

    return NULL;
}

TreeNode* wm_treenode_remove_client(Wm *wm, TreeNode *root, Client *client)
{
    DEBUG_PRINT("%s\n", __func__);
    TreeNode *client_node = wm_treenode_ptr_find_client_node(root, client);

    wm_tree_to_DOT(root, "wm_output.dot");

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

    TreeNode *node = client_node->parent;

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

            node->type = NODE_CLIENT;
            node->client = node->children.nodes[i].client;
            wm_nodearray_clear(&node->children);
            return node;
        }
    }

    DEBUG_PRINT("%s: other_client was NULL!\n", __func__);

    TreeNode *sibling_node = wm_treenode_split_get_sibling(client_node);
    wm_nodearray_free(&client_node->children);

    // wm_nodearray_free(&node->children);

    // TODO
    size_t dock_y = 0;

    if (node->parent == NULL) {
        // parent is root node
        DEBUG_PRINT("parent is root node!\n");
        sibling_node->parent = NULL;
        client->ws->tree = *sibling_node;

        client->ws->tree.pos = (Rect) {
            .x = wm->config.border_width, .y = wm->config.border_width + dock_y ,
            .w = client->ws->monitor->info.width - wm->config.border_width*2,
            .h = client->ws->monitor->info.height - wm->config.border_width*2 - dock_y,
        };

        assert(client->ws->tree.children.size == 2);
        wm_treenode_split_space(&client->ws->tree, &client->ws->tree.children.nodes[0].pos,
                                &client->ws->tree.children.nodes[1].pos);
    } else
        *node = *sibling_node;

    assert(client->ws->tree.children.size >= 1);
    return &node->children.nodes[0];
}

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 * sizeof(TreeNode));

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

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

bool wm_nodearray_pop(NodeArray *arr, TreeNode *ret)
{
    DEBUG_PRINT("%s\n", __func__);
    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--;
    memmove(arr->nodes, arr->nodes+1, arr->size * sizeof(TreeNode));

    return true;
}

void wm_nodearray_clear(NodeArray *arr)
{
    DEBUG_PRINT("%s\n", __func__);
    arr->size = 0;
    memset(arr->nodes, 0, arr->capacity * sizeof(TreeNode));
}

bool wm_nodearray_remove(NodeArray *arr, size_t index)
{
    DEBUG_PRINT("%s\n", __func__);
    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)
{
    DEBUG_PRINT("%s\n", __func__);
    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)
{
    DEBUG_PRINT("%s\n", __func__);
    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 (size_t i = 0; i < node.children.size; i++) {
            TreeNode child_node = node.children.nodes[i];
            bool _visited = false;

            for (size_t j = 0; j < visited.size; j++) {
                if (visited.nodes[j].id == child_node.id) {
                    _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(void*));
    arr.size = 0;

    return arr;
}

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

    arr->size++;

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

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

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

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

bool wm_ptrarray_pop(PtrArray *arr, void **ret)
{
    DEBUG_PRINT("%s\n", __func__);
    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--;
    memmove(arr->ptrs, arr->ptrs+1, arr->size * sizeof(void*));

    return true;
}

bool wm_ptrarray_remove(PtrArray *arr, size_t index)
{
    DEBUG_PRINT("%s\n", __func__);
    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(void*));

    return true;
}

void wm_ptrarray_free(PtrArray *arr)
{
    DEBUG_PRINT("%s\n", __func__);
    assert(arr);

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

TreeNode* wm_treenode_ptr_find_focused_client_node(TreeNode *root)
{
    DEBUG_PRINT("%s\n", __func__);
    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 && node->client->focused) {
            ret = node;
            goto ret;
        }

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

            for (size_t j = 0; j < visited.size; j++) {
                if (visited.nodes[j].id == child_node->id) {
                    _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);
    assert(client);

    TreeNode *ret = NULL;
    PtrArray client_nodes = wm_treenode_find_client_nodes_ptr(root);

    for (size_t i = 0; i < client_nodes.size; i++) {
        TreeNode *node = (TreeNode*)client_nodes.ptrs[i];
        if (node->client == client) {
            ret = node;
            goto ret;
        }
    }

ret:
    wm_ptrarray_free(&client_nodes);
    return ret;
}

PtrArray wm_treenode_find_client_nodes_ptr(TreeNode *root)
{
    DEBUG_PRINT("%s\n", __func__);
    assert(root);

    PtrArray ret = wm_ptrarray_new();

    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 != NULL)
            wm_ptrarray_push(&ret, node);

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

            for (size_t j = 0; j < visited.size; j++) {
                if (visited.nodes[j].id == child_node->id) {
                    _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;
}

PtrArray wm_all_nodes_to_ptrarray(TreeNode *root)
{
    DEBUG_PRINT("%s\n", __func__);
    assert(root);

    PtrArray ret = wm_ptrarray_new();

    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;

        wm_ptrarray_push(&ret, node);

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

            for (size_t j = 0; j < visited.size; j++) {
                if (visited.nodes[j].id == child_node->id) {
                    _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)
{
    DEBUG_PRINT("%s\n", __func__);
    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 * sizeof(Client));

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

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

bool wm_clientarray_pop(ClientArray *arr, Client *ret)
{
    DEBUG_PRINT("%s\n", __func__);
    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)
{
    DEBUG_PRINT("%s\n", __func__);
    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)
{
    DEBUG_PRINT("%s\n", __func__);
    assert(arr);

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