#include <stdlib.h>
#include <stdio.h>
#include <bpf/bpf.h>
#include <bpf/libbpf.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include "common.h"
#include <unistd.h>
#include <errno.h>
#include "ic.h"

#define BUF_SIZE 300        
#define BUCKET_SIZE 16      // Bucket size for InfluxDB
#define INFLUXDB_SIZE 64    // Host and ordID size for InfluxDB
#define TOKEN_SIZE 128      // Token size for InfluxDB
#define CNT_ARGS 5          // Number of args


static void clean_obj(struct bpf_object *obj){
    printf("Cleaning\n");
    bpf_object__close(obj);
}
static void usage(char *app){
    printf("Usage: %s <host> <ordID> <token> <bucket>\n", app);
}
static int check_arguments(int argc, char *argv[]){
    if (argc < 1){
        usage(argv[0]);
        return -1;
    }

    if (strcmp(argv[1], "--no-influxdb") == 0)
        return 0;

    if (argc < CNT_ARGS){
        usage(argv[0]);
        return -1;
    }

    return argc; // Return the number of arguments
}
int main(int argc, char *argv[]){
    const char *fileObj = "tp_tcp.o";
    struct bpf_object *obj;
    struct bpf_program *program;
    struct bpf_map *map;
    struct reset s_reset;
    int err;
    int map_fd;
    int map_fd_filter_family;
    int map_fd_filter_sport;
    int map_fd_index;
    int keys = 0;
    int indexPackets = 0;
    int index = 0;
    char buf[BUF_SIZE];
    char host[INFLUXDB_SIZE];
    char orgID[INFLUXDB_SIZE];
    char token[TOKEN_SIZE];
    char bucket[BUCKET_SIZE];
    int use_influxdb = 0;
    int debug = 1;

    // We get args
    err = check_arguments(argc, argv);
    if (err == -1)
        return -1;

    if (err == CNT_ARGS){
        strncpy(host, argv[1], INFLUXDB_SIZE);
        strncpy(orgID, argv[2], INFLUXDB_SIZE);
        strncpy(token, argv[3], TOKEN_SIZE);
        strncpy(bucket, argv[4], BUCKET_SIZE);
        use_influxdb = 1;
    }

    // Connect to InfluxDB
    if (use_influxdb) {
        ic_influx_database(host, 8086, bucket);
        ic_influx_orgID(orgID);
        ic_influx_token(token);
    }

    obj = bpf_object__open_file(fileObj, NULL);
    if (!obj){
        printf("Failed to open %s\n", fileObj);
        return -1;
    }

    //map_fd = bpf_create_map(BPF_MAP_TYPE_ARRAY, sizeof(int), sizeof(struct reset), 4096, BPF_ANY);

    err = bpf_object__load(obj);
    if (err){
        printf("%s\n", strerror(errno));
        printf("Failed to load object\n");
        clean_obj(obj);
        return -1;
    }

    program = bpf_object__find_program_by_name(obj, "tcp_rst_stats");
    if (!program){
        printf("Failed to find the program\n");
        clean_obj(obj);
        return -1;
    }

    map = bpf_object__find_map_by_name(obj, "tcp_reset_stats");
    if (!map){
        printf("Failed to get the map\n");
        clean_obj(obj);
        return -1;
    }

    map_fd = bpf_object__find_map_fd_by_name(obj, "tcp_reset_stats");
    if (map_fd < 0){
        printf("Failed to find the map 'tcp_reset_stats'\n");
        clean_obj(obj);
        return -1;
    }

    map_fd_filter_family = bpf_object__find_map_fd_by_name(obj, "filter_family");
    if (map_fd_filter_family < 0){
        printf("Failed to find the map 'filter_family'\n");
        clean_obj(obj);
        return -1;
    }

    map_fd_index = bpf_object__find_map_fd_by_name(obj, "tcp_stats_index");
    if (map_fd_index < 0){
        printf("Failed to find the map 'tcp_stats_index'\n");
        clean_obj(obj);
        return -1;
    }

    struct bpf_link *link = bpf_program__attach(program);
    if (!link){
        printf("Failed to attach the program\n");
        return -1;
    }

    // Sepcify our filters
    /*
     * IPv4: AF_INET  -> 2
     * IPv6: AF_INET6 -> 10
     */
    __s16 f = AF_INET;
    err = bpf_map_update_elem(map_fd_filter_family, &keys, &f, BPF_ANY);

    while(1){
        err = bpf_map_lookup_elem(map_fd_index, &keys, &indexPackets);

        // We have a new packet
        if (indexPackets > index){
            index = indexPackets;

            err = bpf_map_lookup_elem(map_fd, &keys, &s_reset);
            if (err == 0){
                struct in_addr *src = (struct in_addr*)&s_reset.saddr;
                struct in_addr *dest = (struct in_addr*)&s_reset.daddr;
                char *s = inet_ntoa(*src);
                char *d;
                char tmp[35];
                int lastValue;
                memcpy(tmp, s, 35);
                d = inet_ntoa(*dest);

                if (debug)
                    printf("Sport: %d; dport: %d %d %d %s - %s\n", s_reset.sport, s_reset.dport, s_reset.family, s_reset.proto, tmp, d);

                if (use_influxdb) {
                    printf("Send data to influx\n");
                    // Send data to InfluxDB
                    snprintf(buf, BUF_SIZE, "host=%s", s);
                    ic_tags(buf);

                    ic_measure("tcp_reset");
                    ic_long("value", 1);
                    ic_measureend();
                    ic_push();
            
                    memset(buf, 0, BUF_SIZE);
                }
            }
            memset(&s_reset, 0, sizeof(struct reset));
        }
    }

    return 0;
}