#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <assert.h>

#include "cJSON.h"

#include "osfp.h"
#include "osfp_fingerprint.h"
#include "osfp_score_db.h"
#include "osfp_log.h"

#include "osfp_common.h"

#define DATA_FILE_PATH "./data.json"
#define DB_FILE_PATH "./db.json"
#define TEST_FILE_PATH "./test.json"
#define LOG_FILE_PATH "./osfp_test.log"

#define OSFP_OS_CLASS_MERGED_MAX (OSFP_OS_CLASS_MAX - 2)
#define EntryWidth 8

unsigned char *data_file_path = DATA_FILE_PATH;
unsigned char *db_file_path;
unsigned char *test_file_path;

FILE *data_file_ptr;
FILE *db_file_ptr;
FILE *test_file_ptr;
FILE *log_file_ptr;

unsigned int debug_enable;

//enum osfp_os_class_id {
//    OSFP_OS_CLASS_UNKNOWN,      // 未知
//    OSFP_OS_CLASS_WINDOWS,      // Windows
//    OSFP_OS_CLASS_LINUX,        // Linux
//    OSFP_OS_CLASS_MAC_OS,       // Mac OS
//    OSFP_OS_CLASS_IOS,          // iOS
//    OSFP_OS_CLASS_ANDROID,      // Android
//    OSFP_OS_CLASS_OTHERS,       // 其他
//    OSFP_OS_CLASS_MAX,
//};
// merged classes: unknown 0 windows-like 1 unix-like 2 apple-like 3 others 4

unsigned int testresult[OSFP_OS_CLASS_MAX][OSFP_OS_CLASS_MAX] = {0};
unsigned int testresult_merged[OSFP_OS_CLASS_MERGED_MAX][OSFP_OS_CLASS_MERGED_MAX] = {0};

static const char *class_to_merged_name(unsigned int class)
{
    switch (class) {
        case 0:
            return "Unknown";
        case 1:
            return "Windows-Like";
        case 2:
            return "Unix-Like";
        case 3:
            return "Apple-Like";
        case 4:
            return "Others";
    }
    return NULL;
}
static unsigned int class_to_merged_class(unsigned int class)
{
    switch (class) {
        case OSFP_OS_CLASS_UNKNOWN:
            return 0;
        case OSFP_OS_CLASS_WINDOWS:
            return 1;
        case OSFP_OS_CLASS_LINUX:
            return 2;
        case OSFP_OS_CLASS_MAC_OS:
            return 3;
        case OSFP_OS_CLASS_IOS:
            return 3;
        case OSFP_OS_CLASS_ANDROID:
            return 2;
        case OSFP_OS_CLASS_OTHERS:
            return 4;
    }
    return 0;
}

static void testresult_class_merge()
{
    int i,j;
    for (i = 0; i < OSFP_OS_CLASS_MAX; i++) {
        for (j = 0; j < OSFP_OS_CLASS_MAX; j++) {
            testresult_merged[class_to_merged_class(i)][class_to_merged_class(j)] += testresult[i][j];
        }
    }
}

static void print_confusion_matrix(unsigned int *result, unsigned int os_class_max, const char *(*get_name)(unsigned int))
{
    int i,j;
    int matched = 0, missed = 0;

    for (i = 0; i < os_class_max; i++) {
        printf("%*s(%c)", EntryWidth-3, " ", 'a' + i);
    }
    printf("    <-" " classified as" "\n");

    for (i = 0; i < os_class_max; i++) {
        printf("%*.*s", EntryWidth, EntryWidth-2, "----------");
    }

    printf("\n");


    for (i = 0; i < os_class_max; i++) {
        for (j = 0; j < os_class_max; j++) {
            if (i == j) {
                matched += *(result + os_class_max * i + j);
            } else {
                missed += *(result + os_class_max * i + j);
            }
            printf(" %*d", EntryWidth-1, *(result + os_class_max * i + j));
        }
        printf("    (%c): " "class" " %s\n", 'a' + i, get_name(i));
    }

    printf("miss rate: %u%%\n", 100 * missed / (matched + missed));
}

struct osfp_result *test_osfp_json_identify(struct osfp_db *db, const char *json_str)
{
    int ret = OSFP_EINVAL;
    struct osfp_fingerprint fp;
    struct osfp_os_class_score os_class_score;
    struct osfp_result *result;
    const char *matched;

    if (db == NULL || json_str == NULL) {
        goto exit;
    }

    ret = osfp_fingerprint_from_json(&fp, (char *)json_str);
    if (ret != 0) {
        goto exit;
    }

    matched = osfp_score_db_prefilter(db->score_db, &fp, &os_class_score);
    if (matched == NULL) {
        ret = osfp_score_db_score(db->score_db, 0, &fp, &os_class_score);
        if (ret != 0) {
            goto exit;
        }
    }

    result = osfp_result_build(&os_class_score, matched);
    if (result == NULL) {
        goto exit;
    }

    return result;
exit:
    return NULL;
}

void test_data_prepare()
{
    char *file_buffer;

    if (test_file_path == NULL) {
        file_buffer = osfp_read_file(data_file_path);
        if (file_buffer == NULL) {
            osfp_log_error("read file: '%s'\n", data_file_path);
            exit(1);
        }

        cJSON *data;
        data = cJSON_Parse(file_buffer);
        if (data == NULL) {
            exit(1);
        }

        cJSON *test;
        test = cJSON_CreateArray();

        unsigned int data_count = cJSON_GetArraySize(data);
        unsigned int db_count = data_count * 8 / 10;
        unsigned int test_count = data_count - db_count;
        srand(time(0));

        int i;
        for (i = 0; i < test_count; i++) {
            cJSON_AddItemToArray(test, cJSON_DetachItemFromArray(data, rand() % data_count));
            data_count--;
        }

        db_file_ptr = fopen(DB_FILE_PATH, "w");
        test_file_ptr = fopen(TEST_FILE_PATH, "w");

        fprintf(db_file_ptr, "%s", cJSON_Print(data));
        fprintf(test_file_ptr, "%s", cJSON_Print(test));

        fflush(db_file_ptr);
        fflush(test_file_ptr);

        db_file_path = DB_FILE_PATH;
        test_file_path = TEST_FILE_PATH;
    }else{
        db_file_path = data_file_path;
    }

}

void test_miss_rate()
{
    int i;
    unsigned int other_count = 0;
    unsigned int unknown_count = 0;
    unsigned int identify_failed_count = 0;
    unsigned int prefiltered_count = 0;
    unsigned int prefiltered_wrong_count = 0;
    unsigned int wrong_count = 0;
    unsigned int verified_count = 0;
    unsigned int fingerprint_count = 0;
    char *file_buffer;

    file_buffer = osfp_read_file(test_file_path);
    if (file_buffer == NULL) {
        osfp_log_error("read file: '%s'\n", test_file_path);
        exit(1);
    }

    cJSON *root = cJSON_Parse(file_buffer);
    if (root == NULL) {
        exit(1);
    }

    fingerprint_count = cJSON_GetArraySize(root);

    if (debug_enable) {
        osfp_log_level_set(OSFP_LOG_LEVEL_DEBUG);
    }

    struct osfp_db *osfp_db = osfp_db_new(db_file_path);
    if (osfp_db == NULL) {
        printf("could not create osfp context. fingerprints file: %s\n", db_file_path);
        exit(1);
    }

    //osfp_score_db_debug_print(osfp_db->score_db);

    FILE *log_file_ptr = fopen(LOG_FILE_PATH, "w");

    for (i = 0; i < fingerprint_count; i++) {
        cJSON *entry = cJSON_GetArrayItem(root, i);
        if (entry) {
            cJSON *os_class_json = cJSON_GetObjectItem(entry, "os");
            int os_class = osfp_os_class_name_to_id(os_class_json->valuestring);

            const char *fp_str = cJSON_PrintUnformatted(entry);

            struct osfp_fingerprint fp = {0};
            osfp_fingerprint_from_json(&fp, (char*)fp_str);
            char str_buf[2048] = "";
            osfp_fingerprint_to_json_buf(&fp, str_buf, 2048, 0);
            fprintf(log_file_ptr, "%s\n", str_buf);

            struct osfp_result *result = test_osfp_json_identify(osfp_db, fp_str);
            if (result == NULL) {
                identify_failed_count++;
                continue;
            }

            testresult[result->likely_os_class][os_class]++;

            if (result->details[result->likely_os_class].score == 100) {
                prefiltered_count++;
            }

            if (os_class == result->likely_os_class) {
                verified_count++;
                osfp_result_free(result);
                continue;
            }

            if (result->details[result->likely_os_class].score == 100) {
                prefiltered_wrong_count++;
            }
            wrong_count++;

            if (result->likely_os_class == OSFP_OS_CLASS_OTHERS) {
                other_count++;
            }

            if (result->likely_os_class == OSFP_OS_CLASS_UNKNOWN) {
                unknown_count++;
            }

            fprintf(log_file_ptr, "expect: %s, result: %s, \n", os_class_json->valuestring, osfp_result_os_name_get(result));

            char *result_json = osfp_result_score_detail_export(result);
            if (result_json) {
                fprintf(log_file_ptr, "%s\n", result_json);
                free(result_json);
            } else {
                fprintf(log_file_ptr, "result detail error:%p\n", result);
            }
            fflush(log_file_ptr);
            osfp_result_free(result);
        }
    }

    printf("total %u, failed %u, pass %u, prefiltered %u (wrong: %u), wrong %u, other %u, unknown %u\n",
                fingerprint_count, identify_failed_count, verified_count, prefiltered_count, prefiltered_wrong_count,  wrong_count, other_count, unknown_count);

    //printf("miss rate: %d%%\n", 100 - (verified_count * 100 / fingerprint_count));

    testresult_class_merge();

    print_confusion_matrix((unsigned int *)testresult, OSFP_OS_CLASS_MAX, osfp_os_class_id_to_name);

    print_confusion_matrix((unsigned int *)testresult_merged, OSFP_OS_CLASS_MERGED_MAX, class_to_merged_name);

    printf("details in: %s\n", LOG_FILE_PATH);

    osfp_db_free(osfp_db);
}

int main(int argc, char **argv)
{
    int r;

    while ((r = getopt(argc, argv, "+f:t:o:d")) != -1) {
        switch(r) {
            case 'f':
                data_file_path = (unsigned char*)optarg;
                break;
            case 't':
                test_file_path = (unsigned char*)optarg;
                break;
            case 'd':
                debug_enable = 1;
                break;
            default:
                break;
        }
    }

    assert(0 == test_osfp_fingerprinting_ipv4());
    assert(0 == test_osfp_fingerprinting_ipv6());
    assert(0 == test_osfp_fingerprinting_tcp_option());
    assert(0 == test_osfp_score_db());

    test_data_prepare();
    test_miss_rate();

    return 0;
}