#include <assert.h>

#include "session_table.h"
#include "utils.h"

struct session_node
{
    uint64_t id;            /* first key */
    struct four_tuple addr; /* second key */

    void *value;
    fn_free_cb *free_cb;

    UT_hash_handle hh1; /* handle for first hash table */
    UT_hash_handle hh2; /* handle for second hash table */
};

struct session_table
{
    struct session_node *root_id;
    struct session_node *root_addr;
    uint64_t count;
};

static struct session_node *search_node_by_addr(struct session_table *table, const struct four_tuple *addr);

struct session_table *session_table_create()
{
    struct session_table *table = (struct session_table *)calloc(1, sizeof(struct session_table));
    assert(table);
    table->count = 0;

    return table;
}

void session_table_destory(struct session_table *table)
{
    if (table)
    {
        session_table_reset(table);

        free(table);
        table = NULL;
    }
}

void session_table_reset(struct session_table *table)
{
    if (table)
    {
        struct session_node *temp = NULL;
        struct session_node *node = NULL;
        HASH_ITER(hh1, table->root_id, node, temp)
        {
            HASH_DELETE(hh1, table->root_id, node);
            HASH_DELETE(hh2, table->root_addr, node);

            if (node->free_cb && node->value)
            {
                node->free_cb(node->value);
            }

            free(node);
            node = NULL;

            table->count--;
        }
    }
}

uint64_t session_table_count(struct session_table *table)
{
    if (table)
    {
        return table->count;
    }
    else
    {
        return 0;
    }
}

// addr  : deep copy
// value : shallow copy (malloc by user, free by free_cb)
// Note  : addr must be initialized by memset(0) before use !!!
int session_table_insert(struct session_table *table, uint64_t id, const struct four_tuple *addr, void *value, const fn_free_cb *free_cb)
{
    struct session_node *node = NULL;
    HASH_FIND(hh1, table->root_id, &id, sizeof(id), node);
    if (node)
    {
        return -1;
    }

    node = (struct session_node *)calloc(1, sizeof(struct session_node));
    assert(node);

    node->id = id;
    node->value = value;
    node->free_cb = free_cb;

    if (addr->src_port < addr->dst_port)
    {
        memcpy(&node->addr, addr, sizeof(struct four_tuple));
    }
    else
    {
        four_tuple_reverse(addr, &node->addr);
    }

    HASH_ADD(hh1, table->root_id, id, sizeof(node->id), node);
    HASH_ADD(hh2, table->root_addr, addr, sizeof(node->addr), node);

    table->count++;

    return 0;
}

int session_table_delete_by_id(struct session_table *table, uint64_t id)
{
    struct session_node *node = NULL;
    HASH_FIND(hh1, table->root_id, &id, sizeof(id), node);
    if (!node)
    {
        return -1;
    }

    HASH_DELETE(hh1, table->root_id, node);
    HASH_DELETE(hh2, table->root_addr, node);

    if (node->free_cb && node->value)
    {
        node->free_cb(node->value);
        node->value = NULL;
    }

    free(node);
    node = NULL;

    table->count--;

    return 0;
}

int session_table_delete_by_addr(struct session_table *table, const struct four_tuple *addr)
{
    struct session_node *node = search_node_by_addr(table, addr);
    if (node == NULL)
    {
        return -1;
    }

    HASH_DELETE(hh1, table->root_id, node);
    HASH_DELETE(hh2, table->root_addr, node);

    if (node->free_cb && node->value)
    {
        node->free_cb(node->value);
        node->value = NULL;
    }

    free(node);
    node = NULL;

    table->count--;

    return 0;
}

void *session_table_search_by_id(struct session_table *table, uint64_t id)
{
    struct session_node *node = NULL;
    HASH_FIND(hh1, table->root_id, &id, sizeof(id), node);
    if (!node)
    {
        return NULL;
    }

    return node->value;
}

void *session_table_search_by_addr(struct session_table *table, const struct four_tuple *addr)
{
    struct session_node *node = search_node_by_addr(table, addr);
    if (!node)
    {
        return NULL;
    }

    return node->value;
}

static struct session_node *search_node_by_addr(struct session_table *table, const struct four_tuple *addr)
{
    struct session_node *node = NULL;
    const struct four_tuple *first = NULL;
    const struct four_tuple *second = NULL;
    struct four_tuple reverse_addr;

    if (addr->src_port < addr->dst_port)
    {
        first = addr;
        second = NULL;
    }
    else
    {
        four_tuple_reverse(addr, &reverse_addr);
        first = &reverse_addr;
        second = addr;
    }

    HASH_FIND(hh2, table->root_addr, first, sizeof(struct four_tuple), node);
    if (node)
    {
        return node;
    }

    if (second == NULL)
    {
        four_tuple_reverse(addr, &reverse_addr);
        second = &reverse_addr;
    }
    HASH_FIND(hh2, table->root_addr, second, sizeof(struct four_tuple), node);
    if (node)
    {
        return node;
    }

    return NULL;
}