/* See the DRL-LICENSE file for this file's software license. */

#include <arpa/inet.h>
#include <assert.h>
#include <inttypes.h>
#include <netinet/in.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <sys/types.h>
#include <time.h>

#include "common_accounting.h"
#include "standard.h"
#include "logging.h"

standard_flow_table standard_table_create(uint32_t (*hash_function)(const key_flow *key), common_accounting_t *common) {
    standard_flow_table table = malloc(sizeof(struct std_flow_table));

    if (table == NULL) {
        return NULL;
    }

    memset(table, 0, sizeof(struct std_flow_table));
    table->common = common;
    table->hash_function = hash_function;

    gettimeofday(&table->common->last_update, NULL);

    return table;
}

void standard_table_destroy(standard_flow_table table) {
    standard_flow *current, *next;

    if ((current = table->flows_head)) {
        while (current->next) {
            next = current->next;
            free(current);
            current = next;
        }
        free(current);
    }

    free(table);
}

/* Looks for the flow in the table.  If the flow isn't there, it allocates a
 * place for it. */
standard_flow *standard_table_lookup(standard_flow_table table, const key_flow *key) {
    uint32_t hash;
    standard_flow *flow;
    struct in_addr src, dst;
    char sip[22], dip[22];

    if (table == NULL) {
        return NULL;
    }

    hash = table->hash_function(key);

    /* Find the flow, if it's there. */
    for (flow = table->flows[hash]; flow; flow = flow->nexth) {
        if (flow->source_ip == key->source_ip &&
                flow->dest_ip == key->dest_ip &&
                flow->source_port == key->source_port &&
                flow->dest_port == key->dest_port &&
                flow->protocol == key->protocol) {
            break;
        }
    }

    if (flow == NULL) {
        flow = malloc(sizeof(standard_flow));
        if (flow == NULL) {
            printf("Malloc returned null.\n");
            printlog(LOG_CRITICAL, "ALLOC: Malloc returned NULL.\n");
            return NULL;
        }

        memset(flow, 0, sizeof(standard_flow));
        flow->protocol = key->protocol;
        flow->source_ip = key->source_ip;
        flow->dest_ip = key->dest_ip;
        flow->source_port = key->source_port;
        flow->dest_port = key->dest_port;
        flow->last_packet = key->packet_time;
        gettimeofday(&flow->last_update, NULL);

        /* Add the flow to the hash list. */
        flow->nexth = table->flows[hash];
        table->flows[hash] = flow;

        /* Add the flow to the linked list. */
        if (table->flows_tail) {
            flow->prev = table->flows_tail;
            table->flows_tail->next = flow;
            table->flows_tail = flow;
        } else {
            table->flows_head = table->flows_tail = flow;
            /* next and prev are already null due to memset above. */
        }

        src.s_addr = ntohl(flow->source_ip);
        dst.s_addr = ntohl(flow->dest_ip);
        strcpy(sip, inet_ntoa(src));
        strcpy(dip, inet_ntoa(dst));
        printlog(LOG_DEBUG, "ALLOC:%s:%hd -> %s:%hd\n", sip,
                flow->source_port, dip, flow->dest_port);
    }

    return flow;
}

int standard_table_sample(standard_flow_table table, const key_flow *key) {
    standard_flow *flow;

    assert(table != NULL);
    assert(table->common != NULL);

    /* Update aggregate. */
    table->common->bytes_since += key->packet_size;

    /* Update flow. */
    flow = standard_table_lookup(table, key);
    if (flow == NULL) {
        return 0;
    }

    flow->bytes_since += key->packet_size;
    flow->last_packet = key->packet_time;

    return 1;
}

void standard_table_remove(standard_flow_table table, standard_flow *flow) {
    key_flow key;
    uint32_t hash;
    standard_flow *current, *prev;

    assert(flow);

    /* Remove the flow from the hash list. */
    key.source_ip = flow->source_ip;
    key.dest_ip = flow->dest_ip;
    key.source_port = flow->source_port;
    key.dest_port = flow->dest_port;
    key.protocol = flow->protocol;

    hash = table->hash_function(&key);

    assert(table->flows[hash]);

    if (table->flows[hash] == flow) {
        /* It's the head of the hash list. */
        table->flows[hash] = flow->nexth;
    } else {
        prev = table->flows[hash];
        current = table->flows[hash]->nexth;

        while (current != NULL) {
            if (current == flow) {
                prev->nexth = flow->nexth;
                break;
            } else {
                prev = current;
                current = current->next;
            }
        }

        assert(current != NULL);
    }

    /* Remove the flow from the linked list. */
    if (flow->prev == NULL && flow->next == NULL) {
        /* It's the head, tail, and only element of the list. */
        assert(table->flows_head == flow);
        assert(table->flows_tail == flow);

        table->flows_head = NULL;
        table->flows_tail = NULL;
    } else if (flow->prev == NULL) {
        /* It's the head of the list. */
        assert(table->flows_head == flow);

        table->flows_head = flow->next;

        if (table->flows_head != NULL) {
            table->flows_head->prev = NULL;
        }
    } else if (flow->next == NULL) {
        /* It's the tail of the list. */
        assert(table->flows_tail == flow);

        table->flows_tail = flow->prev;

        table->flows_tail->next = NULL;
    } else {
        /* Not the head or tail of the list. */
        assert(table->flows_head != flow);

        flow->prev->next = flow->next;

        if (flow->next != NULL) {
            flow->next->prev = flow->prev;
        }
    }

    memset(flow, 0, sizeof(standard_flow));

    /* Free the flow. */
    free(flow);
}

int standard_table_cleanup(standard_flow_table table) {
    standard_flow *current = table->flows_head;
    standard_flow *remove;
    time_t now = time(NULL);

    while (current != NULL) {
        if (current->last_packet + FLOW_IDLE_TIME <= now) {
            /* Flow hasn't received a packet in the time limit - kill it. */
            remove = current;
            current = current->next;

            standard_table_remove(table, remove);
        } else {
            current = current->next;
        }
    }

    return 0;
}

void standard_table_update_flows(standard_flow_table table, struct timeval now, double ewma_weight) {
    uint32_t maxflowrate = 0;
    double time_delta;
    double unweighted_rate;
    standard_flow *current;
    struct in_addr src, dst;
    char sip[22], dip[22];

    time_delta = timeval_subtract(now, table->common->last_update);

    if (time_delta <= 0) {
        unweighted_rate = 0;
    } else {
        unweighted_rate = table->common->bytes_since / time_delta;
    }

    table->common->last_inst_rate = table->common->inst_rate;
    table->common->inst_rate = unweighted_rate;

    table->common->last_rate = table->common->rate;

    /* If the rate is zero, then we don't know anything yet.  Don't apply EWMA
     * in that case. */
    if (table->common->rate == 0) {
        table->common->rate = unweighted_rate;
    } else {
        table->common->rate = table->common->rate * ewma_weight +
                              unweighted_rate * (1 - ewma_weight);
    }

    table->common->bytes_since = 0;
    table->common->last_update = now;

    //printf("Flows: ");

    /* Update per-flow information. */
    for (current = table->flows_head; current; current = current->next) {
        time_delta = timeval_subtract(now, current->last_update);

        if (time_delta <= 0) {
            unweighted_rate = 0;
        } else {
            unweighted_rate = current->bytes_since / time_delta;
        }

        current->last_rate = current->rate;

        if (current->rate == 0) {
            current->rate = unweighted_rate;
        } else {
            current->rate = current->rate * ewma_weight +
                            unweighted_rate * (1 - ewma_weight);
        }

        current->bytes_since = 0;
        current->last_update = now;

        if (current->rate > maxflowrate) {
            maxflowrate = current->rate;
        }

        //printf("%d, ", current->rate);

        src.s_addr = ntohl(current->source_ip);
        dst.s_addr = ntohl(current->dest_ip);
        strcpy(sip, inet_ntoa(src));
        strcpy(dip, inet_ntoa(dst));
        printlog(LOG_DEBUG, "FLOW: (%p)  %s:%d -> %s:%d at %d\n", current,
                sip, current->source_port,
                dip, current->dest_port,
                current->rate);
    }

    //printf("\n");
    printlog(LOG_DEBUG, "FLOW:--\n--\n");

    table->common->max_flow_rate = maxflowrate;
}