#include <stdio.h>
#include <sys/param.h>

#include "cJSON.h"

#include "osfp_common.h"

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

/*
 * 一个字段的值，最多能命中单个操作系统的，指纹总数的百分比
 * 例如：linux 总指纹数量 100，其中 ip_id 字段值为 1 的指纹数为80，那么 ip_id
 * 字段的得分最大为 FIELD_VALUE_DUP_RATE_MAX * 100 = 50
 */
#define FIELD_VALUE_DUP_RATE_MAX 0.5f

/*
 * 由于 FIELD_VALUE_DUP_RATE_MAX 设置了单个字段得分的最大百分比，所以在归一化时，应除以这个百分比
 */
#define PERFECT_SCORE_EXPECTED_RATE (FIELD_VALUE_DUP_RATE_MAX)

#define OSFP_SCORE_DB_FIELD_UINT_VALUE_MAX 65536

struct osfp_score_db_array_data {
    struct osfp_field_value_count *array_head;
    unsigned int    array_len;
};

struct osfp_score_db_hash_element {
    char *key;
    unsigned int keylen;
    struct osfp_field_value_count *fvc;
    UT_hash_handle hh;
};

struct osfp_score_db_hash_data {
    struct osfp_score_db_hash_element  *hash_head;
};

static int osfp_score_db_array_add(void *data, struct osfp_field_value_count *fvc, void *value, unsigned int len)
{
    int ret = -1, i;
    unsigned int index;
    struct osfp_score_db_array_data *array_data = (struct osfp_score_db_array_data *)data;

    if (array_data == NULL || fvc == NULL || value == NULL || len != sizeof(unsigned int)) {
        goto exit;
    }

    if (array_data->array_head == NULL || array_data->array_len == 0) {
        goto exit;
    }

    index = *(unsigned int *)value;

    if (index >= array_data->array_len) {
        goto exit;
    }

    for (i = 0; i < OSFP_OS_CLASS_MAX; i++) {
        array_data->array_head[index].counts[i] += fvc->counts[i];
    }

    return 0;
exit:
    return ret;
}

static struct osfp_field_value_count *osfp_score_db_array_match(void *data, void *value, unsigned int len)
{
    unsigned int index;
    struct osfp_score_db_array_data *array_data = (struct osfp_score_db_array_data *)data;

    if (array_data == NULL || value == NULL || len != sizeof(unsigned int)) {
        return NULL;
    }

    if (array_data->array_head == NULL || array_data->array_len == 0) {
        return NULL;
    }

    index = *(unsigned int *)value;

    if (index >= array_data->array_len) {
        return NULL;
    }

    return &((array_data->array_head)[index]);
}

static void *osfp_score_db_array_create(void)
{
    struct osfp_score_db_array_data *array_data = calloc(1, sizeof(struct osfp_score_db_array_data));
    if (array_data == NULL) {
        return NULL;
    }

    array_data->array_head = calloc(OSFP_SCORE_DB_FIELD_UINT_VALUE_MAX, sizeof(struct osfp_field_value_count));
    if (array_data->array_head == NULL) {
        free(array_data);
        return NULL;
    }

    array_data->array_len = OSFP_SCORE_DB_FIELD_UINT_VALUE_MAX;

    return (void *)array_data;
}

static void osfp_score_db_array_destroy(void *data) {
    struct osfp_score_db_array_data *array_data = (struct osfp_score_db_array_data *)data;

    if (array_data) {
        if (array_data->array_head) {
            free(array_data->array_head);
        }
        free(array_data);
    }
}

static int osfp_score_db_hash_add(void *data, struct osfp_field_value_count *fvc, void *value, unsigned int len)
{
    int ret = -1, i;
    struct osfp_score_db_hash_data *hash_data = (struct osfp_score_db_hash_data *)data;
    struct osfp_score_db_hash_element *element = NULL;

    if (hash_data == NULL || fvc == NULL || value == NULL || len == 0) {
        goto exit;
    }

    HASH_FIND(hh, hash_data->hash_head, value, len, element);
    if (element == NULL) {
        element = (struct osfp_score_db_hash_element *)calloc(1, sizeof(struct osfp_score_db_hash_element));
        if (element == NULL) {
            goto exit;
        }
        element->key = strdup(value);
        element->keylen = len;
        element->fvc = (struct osfp_field_value_count *)calloc(1, sizeof(struct osfp_field_value_count));
        if (element->fvc == NULL) {
            free(element);
            element = NULL;
            goto exit;
        }
        HASH_ADD_KEYPTR(hh, hash_data->hash_head, element->key, element->keylen, element);
    }

    for (i = 0; i < OSFP_OS_CLASS_MAX; i++) {
        element->fvc->counts[i] += fvc->counts[i];
    }

    return 0;
exit:
    return ret;
}

static struct osfp_field_value_count *osfp_score_db_hash_match(void *data, void *value, unsigned int len)
{
    struct osfp_score_db_hash_data *hash_data = (struct osfp_score_db_hash_data *)data;
    struct osfp_score_db_hash_element *element = NULL;

    if (data == NULL || value == NULL || len == 0) {
        return NULL;
    }

    if (hash_data->hash_head == NULL) {
        return NULL;
    }

    HASH_FIND(hh, hash_data->hash_head, value, len, element);
    if (element == NULL) {
        return NULL;
    }

    return element->fvc;
}

static void *osfp_score_db_hash_create(void)
{
    return (void*)calloc(1, sizeof(struct osfp_score_db_hash_data));
}

static void osfp_score_db_hash_destroy(void *data) {
    struct osfp_score_db_hash_data *hash_data = (struct osfp_score_db_hash_data *)data;
    struct osfp_score_db_hash_element *element = NULL;
    struct osfp_score_db_hash_element *tmp = NULL;

    if (hash_data) {
        if (hash_data->hash_head) {
            HASH_ITER(hh, hash_data->hash_head, element, tmp) {
                HASH_DELETE(hh, hash_data->hash_head, element);
                if (element) {
                    if (element->key) {
                        free(element->key);
                    }
                    if (element->fvc) {
                        free(element->fvc);
                    }
                    free(element);
                }
            }
        }
        free(hash_data);
    }
}

void osfp_score_db_debug_print(struct osfp_score_db *score_db)
{
    int i;
    printf("score_db:\n");
    printf("entry_count: %u\n", score_db->entry_count);
    printf("total_weight: %u\n", score_db->total_weight);

    for (i = 0; i < OSFP_OS_CLASS_MAX; i++) {
        const char *name = osfp_os_class_id_to_name(i);
        printf("os class %s ", name);

        printf("entry_count: %u\n", score_db->os_entry_count[i]);
        printf("os class %s entry_count: %u\n", osfp_os_class_id_to_name(i), score_db->os_entry_count[i]);
    }

    for (i = 0; i < OSFP_FIELD_MAX; i++) {
        printf("field %s enabled: %u\n", osfp_fingerprint_get_field_name(i), score_db->field_score_dbs[i].enabled);
        printf("field %s type: %u\n", osfp_fingerprint_get_field_name(i), score_db->field_score_dbs[i].type);
        printf("field %s entry_count: %u\n", osfp_fingerprint_get_field_name(i), score_db->field_score_dbs[i].entry_count);
    }
}

const char *osfp_score_db_prefilter(struct osfp_score_db *score_db, struct osfp_fingerprint *fp, struct osfp_os_class_score *result_score)
{
    int ret, i;
    unsigned int value_buffer_used = 0;
    char value_buffer[OSFP_FINGERPRINT_VALUE_BUFFER_MAX];
    struct osfp_prefilter_hash_element *element = NULL;

    if (score_db->prefilter_head == NULL) {
        return NULL;
    }

    for (i = 0; i < OSFP_FIELD_OS; i++) {
        if (0 == osfp_fingerprint_get_field_enabled(i)) {
            continue;
        }
        if (fp->fields[i].value && fp->fields[i].value_len != 0) {
            memcpy(value_buffer + value_buffer_used, fp->fields[i].value, fp->fields[i].value_len);
            value_buffer_used += fp->fields[i].value_len;
        }
    }

    HASH_FIND(hh, score_db->prefilter_head, value_buffer, value_buffer_used, element);
    if (element == NULL) {
        return NULL;
    }

    if (element->repeated) {
        return NULL;
    }

    memset(result_score, 0, sizeof(struct osfp_os_class_score));
    result_score->scores[element->os_class] = OSFP_PERCENTILE;

    return (const char *)element->fp_json;
}

int osfp_score_db_score(struct osfp_score_db *score_db, unsigned int flags, struct osfp_fingerprint *fp, struct osfp_os_class_score *result_score)
{
    int ret = -1, i, j;

    unsigned int coefficient;
    unsigned int matched_count;
    unsigned int os_entry_count;
    unsigned int total_weight;
    unsigned int field_weight;

    struct osfp_fingerprint_field *field;
    struct osfp_field_value_count *fvc; // field_value_count
    struct osfp_field_score_db *fdb;

    if (score_db == NULL || fp == NULL || result_score == NULL) {
        goto exit;
    }

    // score
    memset(result_score, 0, sizeof(struct osfp_os_class_score));

    total_weight = score_db->total_weight;
    if (total_weight == 0) {
        goto exit;
    }

    for (i = 0; i < OSFP_FIELD_MAX; i++) {
        fdb = &score_db->field_score_dbs[i];
        if (!fdb->enabled) {
            continue;
        }

        field = &fp->fields[i];
        if (!field->enabled) {
            continue;
        }

        fvc = fdb->match(fdb->data, field->value, field->value_len);
        if (fvc == NULL) {
            continue;
        }

        coefficient = fdb->coefficient;

        for (j = 0; j < OSFP_OS_CLASS_MAX; j++) {
            os_entry_count = score_db->os_entry_count[j];
            matched_count = MIN(fvc->counts[j], os_entry_count * FIELD_VALUE_DUP_RATE_MAX);

            if (os_entry_count == 0 || matched_count == 0) {
                continue;
            }

            if (0 == flags || flags & OSFP_BIT_U32(j)) {
                result_score->scores[j] += coefficient * matched_count / os_entry_count;
            }
        }

        // if tcp options matched tcp options ordered is not needed
        if (i == OSFP_FIELD_TCP_OPTIONS) {
            i++;
        }
    }

    return 0;
exit:
    return ret;
}

static int osfp_score_db_load_field(struct osfp_field_score_db *db, cJSON *field, enum osfp_os_class_id os_class)
{
    int ret = -1;
    struct osfp_field_value_count field_value_count = {0};

    void *value_ptr;
    unsigned int value_len;

    switch (field->type) {
        case cJSON_Number:
            value_ptr = (void *)&field->valueint;
            value_len = sizeof(field->valueint);
            break;
        case cJSON_String:
            value_ptr = (void *)field->valuestring;
            value_len = strlen(field->valuestring) + 1;
            break;
        case cJSON_NULL:
            ret = 0;
            goto exit;
        default:
            goto exit;
    }

    field_value_count.counts[os_class] = 1;

    ret = db->add(db->data, &field_value_count, value_ptr, value_len);
    if (ret != 0) {
        goto exit;
    }

    db->entry_count++;

    return 0;
exit:
    return ret;
}

static int osfp_score_db_load_entry(struct osfp_score_db *score_db, cJSON *entry)
{
    int ret = -1, i;
    cJSON *field = NULL;
    struct osfp_field_score_db *db;
    enum osfp_os_class_id os_class;

    if (score_db == NULL || entry == NULL) {
        goto exit;
    }

    // get os class
    field = cJSON_GetObjectItem(entry, osfp_fingerprint_get_field_name(OSFP_FIELD_OS));
    if (field == NULL || field->valuestring == NULL) {
        goto exit;
    }

    os_class = osfp_os_class_name_to_id(field->valuestring);
    if (os_class >= OSFP_OS_CLASS_MAX) {
        goto exit;
    }

    // prefileter
    struct osfp_prefilter_hash_element *element = NULL;
    struct osfp_fingerprint *fp = osfp_fingerprint_from_cjson(entry);
    if (fp == NULL) {
        goto exit;
    }
    HASH_FIND(hh, score_db->prefilter_head, fp->value_buffer, fp->value_buffer_used, element);
    if (element == NULL) {
        element = (struct osfp_prefilter_hash_element *)calloc(1, sizeof(struct osfp_prefilter_hash_element));
        if (element == NULL) {
            free(fp);
            goto exit;
        }

        element->fp_json = cJSON_Print(entry);
        element->fp = fp;
        element->os_class = os_class;
        HASH_ADD_KEYPTR(hh, score_db->prefilter_head, fp->value_buffer, fp->value_buffer_used, element);
    } else {
        // TODO: same fingerprints with different os should not insert into prefilter hash table, now just tag
        element->repeated++;
        free(fp);
    }

    // field score db
    for (i = 0; i < OSFP_FIELD_OS; i++) {
        db = &score_db->field_score_dbs[i];
        if (db == NULL) {
            goto exit;
        }

        if (!db->enabled) {
            continue;
        }

        field = cJSON_GetObjectItem(entry, osfp_fingerprint_get_field_name(i));
        if (field == NULL) {
            osfp_log_info("json entry missing field: %s\n%s\n",
                        osfp_fingerprint_get_field_name(i), cJSON_Print(entry));
            continue;
        }

        ret = osfp_score_db_load_field(db, field, os_class);
        if (ret != 0) {
            osfp_log_info("json entry field load failed. field: %s\n%s\n",
                        osfp_fingerprint_get_field_name(i), cJSON_Print(entry));
            continue;
        }
    }

    score_db->entry_count++;
    score_db->os_entry_count[os_class]++;

    return 0;
exit:
    return ret;
}

int osfp_score_db_load(struct osfp_score_db *score_db, char *fp_file)
{
    int ret = -1, i, count;
    char *file_buffer;
    cJSON *entry;
    cJSON *root = NULL;
    struct osfp_field_score_db *field_score_db;

    if (score_db == NULL || fp_file == NULL) {
        goto exit;
    }

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

    root = cJSON_Parse(file_buffer);
    if (root == NULL) {
        osfp_log_error("parse json: '%s'\n", fp_file);
        goto exit;
    }

    count = cJSON_GetArraySize(root);
    for (i = 0; i < count; i++) {
        entry = cJSON_GetArrayItem(root, i);
        ret = osfp_score_db_load_entry(score_db, entry);
        if (ret != 0) {
            osfp_log_debug("json entry load failed.\n%s\n", cJSON_Print(entry));
            continue;
        }
    }

    ret = 0;
exit:
    if (root) {
        cJSON_Delete(root);
    }
    if (file_buffer) {
        free(file_buffer);
    }
    return ret;
}

struct osfp_score_db *osfp_score_db_create(void)
{
    int i;
    struct osfp_score_db *score_db;
    struct osfp_field_score_db *db;

    score_db = calloc(1, sizeof(struct osfp_score_db));
    if (score_db == NULL) {
        goto exit;
    }

    for (i = 0; i < OSFP_FIELD_MAX; i++) {
        db = &score_db->field_score_dbs[i];

        db->enabled = osfp_fingerprint_get_field_enabled(i);
        if (!db->enabled) {
            osfp_log_info("fingerprint field disabled: %s", osfp_fingerprint_get_field_name(i));
            continue;
        }

        db->weight = osfp_fingerprint_get_field_importance(i);
        // tcp options ordered and tcp options overlap
        if (i != OSFP_FIELD_TCP_OPTIONS_ORDERED) {
            score_db->total_weight += db->weight;
        }

        db->type = osfp_fingerprint_get_field_type(i);
        switch (db->type) {
            case OSFP_FIELD_TYPE_UINT:
                db->create = osfp_score_db_array_create;
                db->destroy = osfp_score_db_array_destroy;
                db->add = osfp_score_db_array_add;
                db->match = osfp_score_db_array_match;
                break;
            case OSFP_FIELD_TYPE_STRING:
                db->create = osfp_score_db_hash_create;
                db->destroy = osfp_score_db_hash_destroy;
                db->add = osfp_score_db_hash_add;
                db->match = osfp_score_db_hash_match;
                break;
            default:
                osfp_log_error("fingerprint field unsupported type: %u", db->type);
                goto exit;
        }

        db->data = db->create();
        if (db->data == NULL) {
            osfp_log_error("field db create failed. field: %s", osfp_fingerprint_get_field_name(i));
            goto exit;
        }
    }

    for (i = 0; i < OSFP_FIELD_MAX; i++) {
        db = &score_db->field_score_dbs[i];
        db->coefficient = (OSFP_PERCENTILE / PERFECT_SCORE_EXPECTED_RATE) * ((float)db->weight / (float)score_db->total_weight);
    }

    return score_db;
exit:
    if (score_db) {
        osfp_score_db_destroy(score_db);
    }
    return NULL;
}

void osfp_score_db_destroy(struct osfp_score_db *score_db)
{
    int i;
    struct osfp_field_score_db *db;
    struct osfp_prefilter_hash_element *element = NULL;
    struct osfp_prefilter_hash_element *tmp = NULL;

    if (score_db) {
        // prefilter
        if (score_db->prefilter_head) {
            HASH_ITER(hh, score_db->prefilter_head, element, tmp) {
                HASH_DELETE(hh, score_db->prefilter_head, element);
                if (element) {
                    if (element->fp) {
                        free(element->fp);
                    }
                    if (element->fp_json) {
                        free(element->fp_json);
                    }
                    free(element);
                }
            }
        }

        // field score db
        for (i = 0; i < OSFP_FIELD_MAX; i++) {
            db = &score_db->field_score_dbs[i];
            if (db->destroy && db->data) {
                db->destroy(db->data);
                db->data = NULL;
            }
        }
        free(score_db);
    }
}

int test_osfp_score_db(void)
{
    int ret, i;
    struct osfp_score_db *db;
    const char *test_file_path = "./.fp.json";
    const char *fingerprint_file_content = "["
        "  {"
        "    \"tcp_options\": \"M1432,S,T,N,W9,\","
        "    \"tcp_options_ordered\": \"MSTNW\","
        "    \"ip_total_length\": 60,"
        "    \"tcp_off\": 10,"
        "    \"tcp_window_scaling\": 9,"
        "    \"tcp_window_size\": 65535,"
        "    \"ip_ttl\": 64,"
        "    \"ip_id\": 1,"
        "    \"tcp_timestamp\": 1,"
        "    \"tcp_timestamp_echo_reply\": 0,"
        "    \"tcp_mss\": 1432,"
        "    \"tcp_flags\": 2,"
        "    \"ip_tos\": 0,"
        "    \"os\": \"Android\""
        "  }"
        "]";

    db = (void *)osfp_score_db_create();
    if (db == NULL) {
        goto exit;
    }

    FILE *fingerprint_file_ptr = fopen(test_file_path, "w");
    if (fingerprint_file_ptr == NULL) {
        goto exit;
    }

    fprintf(fingerprint_file_ptr, "%s", fingerprint_file_content);
    fflush(fingerprint_file_ptr);
    fclose(fingerprint_file_ptr);

    ret = osfp_score_db_load(db, (char *)test_file_path);
    remove(test_file_path);
    if (ret != 0) {
        goto exit;
    }

    if (db->entry_count != 1) {
        goto exit;
    }

    if (db->os_entry_count[OSFP_OS_CLASS_ANDROID] != 1) {
        goto exit;
    }

    for (i = 0; i < OSFP_FIELD_MAX; i++) {
        if (db->field_score_dbs[i].enabled != osfp_fingerprint_get_field_enabled(i)) {
            goto exit;
        }
        if (db->field_score_dbs[i].enabled && db->field_score_dbs[i].type != osfp_fingerprint_get_field_type(i)) {
            goto exit;
        }
        if (db->field_score_dbs[i].enabled && db->field_score_dbs[i].entry_count != 1) {
            goto exit;
        }
    }

    osfp_score_db_destroy(db);

    return 0;
exit:
    return -1;
}