#include <assert.h>
#include <tfe_cmsg.h>
#include <tfe_utils.h>
#include <tfe_stream.h>
#include <tfe_resource.h>
#include <cjson/cJSON.h>
#include <MESA/maat.h>

#include "tcp_policy.h"

struct tcp_policy_enforcer
{
    struct maat *maat;
    int policy_table_id;
    int profile_table_id;
    void *logger;
};

struct side_conn_param
{
    int maxseg_enable;
    int maxseg_vaule;
    int nodelay;
    int keepalive;
    int keepcnt;
    int keepidle;
    int keepintvl;
    int ttl;
    int user_timeout;
};

struct tcp_profile_param
{
    int ref_cnt;
    int tcp_passthrough;
    int bypass_duplicated_packet; // TODO no use, need delete from CM

    struct side_conn_param client_side;
    struct side_conn_param server_side;
};

struct intercept_param
{
    uint64_t rule_id;
    int ref_cnt;
    int tcp_option_profile;
};

static void intercept_param_new_cb(const char *table_name, int table_id, const char *key, const char *table_line, void **ad, long argl, void *argp)
{
    size_t offset = 0;
    size_t len = 0;
    char *json_str = NULL;
    cJSON *json = NULL;
    cJSON *item = NULL;
    struct intercept_param *policy_param = NULL;
    struct tcp_policy_enforcer *enforcer = (struct tcp_policy_enforcer *)argp;

    if (maat_helper_read_column(table_line, 7, &offset, &len) < 0)
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept user region: %s", table_line);
        goto error_out;
    }

    json_str = ALLOC(char, len + 1);
    memcpy(json_str, table_line + offset, len);
    json = cJSON_Parse(json_str);
    if (json == NULL)
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: id = %s", key);
        goto error_out;
    }

    item = cJSON_GetObjectItem(json, "tcp_option_profile");
    if (item == NULL || item->type != cJSON_Number)
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid intercept parameter: %s invalid tcp_option_profile format", key);
        goto error_out;
    }

    policy_param = ALLOC(struct intercept_param, 1);
    policy_param->rule_id = atoll(key);
    policy_param->ref_cnt = 1;
    policy_param->tcp_option_profile = item->valueint;

    *ad = policy_param;
    TFE_LOG_INFO(enforcer->logger, "Add intercept policy: %lu", policy_param->rule_id);

error_out:
    if (json)
    {
        cJSON_Delete(json);
    }
    if (json_str)
    {
        free(json_str);
    }
}

static void intercept_param_free_cb(int table_id, void **ad, long argl, void *argp)
{
    struct tcp_policy_enforcer *enforcer = (struct tcp_policy_enforcer *)argp;
    struct intercept_param *policy_param = (struct intercept_param *)*ad;
    if (policy_param == NULL)
    {
        return;
    }

    if ((__sync_sub_and_fetch(&policy_param->ref_cnt, 1) == 0))
    {
        TFE_LOG_INFO(enforcer->logger, "Del intercept policy %lu", policy_param->rule_id);
        free(policy_param);
        *ad = NULL;
    }
}

static void intercept_param_free(struct intercept_param *policy_param)
{
    intercept_param_free_cb(0, (void **)&policy_param, 0, NULL);
}

static void intercept_param_dup_cb(int table_id, void **to, void **from, long argl, void *argp)
{
    struct intercept_param *policy_param = (struct intercept_param *)*from;
    if (policy_param)
    {
        __sync_add_and_fetch(&(policy_param->ref_cnt), 1);
        *to = policy_param;
    }
    else
    {
        *to = NULL;
    }
}

static int parser_side_conn_param(const char *json_str, struct side_conn_param *out_val, void *logger)
{
    cJSON *json = NULL;
    cJSON *object = NULL;
    cJSON *item = NULL;

    json = cJSON_Parse(json_str);
    if (json == NULL)
    {
        TFE_LOG_ERROR(logger, "Invalid tcp option param %s", json_str);
        return -1;
    }

    object = cJSON_GetObjectItem(json, "tcp_maxseg");
    if (object)
    {
        item = cJSON_GetObjectItem(object, "enable");
        if (item && item->type == cJSON_Number)
        {
            out_val->maxseg_enable = item->valueint;
        }
        item = cJSON_GetObjectItem(object, "maxseg");
        if (item && item->type == cJSON_Number)
        {
            out_val->maxseg_vaule = item->valueint;
        }
    }

    item = cJSON_GetObjectItem(json, "nodelay");
    if (item && item->type == cJSON_Number)
    {
        out_val->nodelay = item->valueint;
    }

    object = cJSON_GetObjectItem(json, "keep_alive");
    if (object)
    {
        item = cJSON_GetObjectItem(object, "enable");
        if (item && item->type == cJSON_Number)
        {
            out_val->keepalive = item->valueint;
        }
        item = cJSON_GetObjectItem(object, "tcp_keepcnt");
        if (item && item->type == cJSON_Number)
        {
            out_val->keepcnt = item->valueint;
        }
        item = cJSON_GetObjectItem(object, "tcp_keepidle");
        if (item && item->type == cJSON_Number)
        {
            out_val->keepidle = item->valueint;
        }
        item = cJSON_GetObjectItem(object, "tcp_keepintvl");
        if (item && item->type == cJSON_Number)
        {
            out_val->keepintvl = item->valueint;
        }
    }

    item = cJSON_GetObjectItem(json, "ttl");
    if (item && item->type == cJSON_Number)
    {
        out_val->ttl = item->valueint;
    }

    item = cJSON_GetObjectItem(json, "user_timeout");
    if (item && item->type == cJSON_Number)
    {
        out_val->user_timeout = item->valueint;
    }
    cJSON_Delete(json);

    return 0;
}

static void profile_param_new_cb(const char *table_name, int table_id, const char *key, const char *table_line, void **ad, long argl, void *argp)
{
    int ret = 0;
    int profile_id = 0;
    int tcp_passthrough = 0;
    int bypass_duplicated_packet = 0;
    char client_side_conn_param[512] = {0};
    char server_side_conn_param[512] = {0};
    int is_valid = 0;
    struct tcp_profile_param *profile_param = NULL;
    struct tcp_policy_enforcer *enforcer = (struct tcp_policy_enforcer *)argp;

    ret = sscanf(table_line, "%d\t%d\t%d\t%s\t%s\t%d", &profile_id, &tcp_passthrough, &bypass_duplicated_packet, client_side_conn_param, server_side_conn_param, &is_valid);
    if (ret != 6)
    {
        TFE_LOG_ERROR(enforcer->logger, "Invalid tcp option profile: %s", table_line);
        goto error_out;
    }

    profile_param = ALLOC(struct tcp_profile_param, 1);
    profile_param->ref_cnt = 1;
    profile_param->tcp_passthrough = tcp_passthrough;
    profile_param->bypass_duplicated_packet = bypass_duplicated_packet;

    if (parser_side_conn_param(client_side_conn_param, &profile_param->client_side, enforcer->logger) == -1)
    {
        goto error_out;
    }

    if (parser_side_conn_param(server_side_conn_param, &profile_param->server_side, enforcer->logger) == -1)
    {
        goto error_out;
    }

    *ad = profile_param;
    TFE_LOG_INFO(enforcer->logger, "Add tcp option profile: %s", key);
    return;

error_out:
    if (profile_param)
    {
        free(profile_param);
    }
    return;
}

static void profile_param_free_cb(int table_id, void **ad, long argl, void *argp)
{
    struct tcp_profile_param *profile_param = (struct tcp_profile_param *)*ad;
    if (profile_param == NULL)
    {
        return;
    }

    if ((__sync_sub_and_fetch(&profile_param->ref_cnt, 1) == 0))
    {
        free(profile_param);
        *ad = NULL;
    }
}

static void profile_param_free(struct tcp_profile_param *profile_param)
{
    profile_param_free_cb(0, (void **)&profile_param, 0, NULL);
}

static void profile_param_dup_cb(int table_id, void **to, void **from, long argl, void *argp)
{
    struct tcp_profile_param *profile_param = (struct tcp_profile_param *)*from;
    if (profile_param)
    {
        __sync_add_and_fetch(&(profile_param->ref_cnt), 1);
        *to = profile_param;
    }
    else
    {
        *to = NULL;
    }
}

struct tcp_policy_enforcer *tcp_policy_enforcer_create(void *logger)
{
    int ret = 0;
    struct tcp_policy_enforcer *enforcer = ALLOC(struct tcp_policy_enforcer, 1);
    enforcer->maat = (struct maat *)tfe_bussiness_resouce_get(STATIC_MAAT);
    enforcer->logger = logger;
    enforcer->policy_table_id = maat_get_table_id(enforcer->maat, "TSG_SECURITY_COMPILE");
    if (enforcer->policy_table_id < 0)
    {
        TFE_LOG_ERROR(enforcer->logger, "failed at register table of TSG_SECURITY_COMPILE, ret = %d", enforcer->policy_table_id);
        goto error_out;
    }
    enforcer->profile_table_id = maat_get_table_id(enforcer->maat, "PXY_PROFILE_TCP_OPTION");
    if (enforcer->profile_table_id < 0)
    {
        TFE_LOG_ERROR(enforcer->logger, "failed at register table of PXY_PROFILE_TCP_OPTION, ret = %d", enforcer->profile_table_id);
        goto error_out;
    }

    ret = maat_plugin_table_ex_schema_register(enforcer->maat, "TSG_SECURITY_COMPILE",
                                               intercept_param_new_cb,
                                               intercept_param_free_cb,
                                               intercept_param_dup_cb,
                                               0, enforcer);
    if (ret < 0)
    {
        TFE_LOG_ERROR(enforcer->logger, "failed at register callback of TSG_SECURITY_COMPILE, ret = %d", ret);
        goto error_out;
    }
    ret = maat_plugin_table_ex_schema_register(enforcer->maat, "PXY_PROFILE_TCP_OPTION",
                                               profile_param_new_cb,
                                               profile_param_free_cb,
                                               profile_param_dup_cb,
                                               0, enforcer);
    if (ret < 0)
    {
        TFE_LOG_ERROR(enforcer->logger, "failed at register callback of PXY_PROFILE_TCP_OPTION, ret = %d", ret);
        goto error_out;
    }
    return enforcer;

error_out:
    tcp_policy_enforcer_destory(enforcer);
    return NULL;
}

void tcp_policy_enforcer_destory(struct tcp_policy_enforcer *enforcer)
{
    if (enforcer)
    {
        free(enforcer);
        enforcer = NULL;
    }
}

void tcp_policy_enforce(struct tcp_policy_enforcer *enforcer, struct tfe_cmsg *cmsg, uint64_t rule_id)
{
    char rule_id_str[16] = {0};
    char profile_id_str[16] = {0};

    snprintf(rule_id_str, sizeof(rule_id_str), "%lu", rule_id);
    struct intercept_param *policy_param = (struct intercept_param *)maat_plugin_table_get_ex_data(enforcer->maat, enforcer->policy_table_id, rule_id_str);
    if (policy_param == NULL)
    {
        TFE_LOG_INFO(enforcer->logger, "Failed to get intercept parameter of policy %lu.", rule_id);
        return;
    }

    snprintf(profile_id_str, sizeof(profile_id_str), "%d", policy_param->tcp_option_profile);
    struct tcp_profile_param *profile_param = (struct tcp_profile_param *)maat_plugin_table_get_ex_data(enforcer->maat, enforcer->profile_table_id, profile_id_str);
    if (profile_param == NULL)
    {
        TFE_LOG_INFO(enforcer->logger, "Failed to get tcp option parameter of profile %lu.", rule_id);
        intercept_param_free(policy_param);
        return;
    }

    tfe_cmsg_set(cmsg, TFE_CMSG_TCP_PASSTHROUGH, (unsigned char *)&profile_param->tcp_passthrough, sizeof(profile_param->tcp_passthrough));

    struct side_conn_param *client_side = &profile_param->client_side;
    tfe_cmsg_set(cmsg, TFE_CMSG_DOWNSTREAM_TCP_MSS_ENABLE, (unsigned char *)&client_side->maxseg_enable, sizeof(client_side->maxseg_enable));
    tfe_cmsg_set(cmsg, TFE_CMSG_DOWNSTREAM_TCP_MSS_VALUE, (unsigned char *)&client_side->maxseg_vaule, sizeof(client_side->maxseg_vaule));
    tfe_cmsg_set(cmsg, TFE_CMSG_DOWNSTREAM_TCP_NODELAY, (unsigned char *)&client_side->nodelay, sizeof(client_side->nodelay));
    tfe_cmsg_set(cmsg, TFE_CMSG_DOWNSTREAM_TCP_TTL, (unsigned char *)&client_side->ttl, sizeof(client_side->ttl));
    tfe_cmsg_set(cmsg, TFE_CMSG_DOWNSTREAM_TCP_KEEPALIVE, (unsigned char *)&client_side->keepalive, sizeof(client_side->keepalive));
    tfe_cmsg_set(cmsg, TFE_CMSG_DOWNSTREAM_TCP_KEEPCNT, (unsigned char *)&client_side->keepcnt, sizeof(client_side->keepcnt));
    tfe_cmsg_set(cmsg, TFE_CMSG_DOWNSTREAM_TCP_KEEPIDLE, (unsigned char *)&client_side->keepidle, sizeof(client_side->keepidle));
    tfe_cmsg_set(cmsg, TFE_CMSG_DOWNSTREAM_TCP_KEEPINTVL, (unsigned char *)&client_side->keepidle, sizeof(client_side->keepintvl));
    tfe_cmsg_set(cmsg, TFE_CMSG_DOWNSTREAM_TCP_USER_TIMEOUT, (unsigned char *)&client_side->user_timeout, sizeof(client_side->user_timeout));

    struct side_conn_param *server_side = &profile_param->server_side;
    tfe_cmsg_set(cmsg, TFE_CMSG_UPSTREAM_TCP_MSS_ENABLE, (unsigned char *)&server_side->maxseg_enable, sizeof(server_side->maxseg_enable));
    tfe_cmsg_set(cmsg, TFE_CMSG_UPSTREAM_TCP_MSS_VALUE, (unsigned char *)&server_side->maxseg_vaule, sizeof(server_side->maxseg_vaule));
    tfe_cmsg_set(cmsg, TFE_CMSG_UPSTREAM_TCP_NODELAY, (unsigned char *)&server_side->nodelay, sizeof(server_side->nodelay));
    tfe_cmsg_set(cmsg, TFE_CMSG_UPSTREAM_TCP_TTL, (unsigned char *)&server_side->ttl, sizeof(server_side->ttl));
    tfe_cmsg_set(cmsg, TFE_CMSG_UPSTREAM_TCP_KEEPALIVE, (unsigned char *)&server_side->keepalive, sizeof(server_side->keepalive));
    tfe_cmsg_set(cmsg, TFE_CMSG_UPSTREAM_TCP_KEEPCNT, (unsigned char *)&server_side->keepcnt, sizeof(server_side->keepcnt));
    tfe_cmsg_set(cmsg, TFE_CMSG_UPSTREAM_TCP_KEEPIDLE, (unsigned char *)&server_side->keepidle, sizeof(server_side->keepidle));
    tfe_cmsg_set(cmsg, TFE_CMSG_UPSTREAM_TCP_KEEPINTVL, (unsigned char *)&server_side->keepintvl, sizeof(server_side->keepintvl));
    tfe_cmsg_set(cmsg, TFE_CMSG_UPSTREAM_TCP_USER_TIMEOUT, (unsigned char *)&server_side->user_timeout, sizeof(server_side->user_timeout));

    TFE_LOG_INFO(enforcer->logger, "hit rule_id %lu tcp_option_profile %d tcp_passthrough %d "
                                   "client_side={maxseg_enable:%d, maxseg_vaule:%d, nodelay:%d, ttl:%d, keepalive:%d, keepcnt:%d, keepidle:%d, keepintvl:%d, user_timeout:%d} "
                                   "server_side={maxseg_enable:%d, maxseg_vaule:%d, nodelay:%d, ttl:%d, keepalive:%d, keepcnt:%d, keepidle:%d, keepintvl:%d, user_timeout:%d} ",
                 rule_id, policy_param->tcp_option_profile, profile_param->tcp_passthrough,
                 client_side->maxseg_enable, client_side->maxseg_vaule, client_side->nodelay, client_side->ttl, client_side->keepalive, client_side->keepcnt, client_side->keepidle, client_side->keepidle, client_side->user_timeout,
                 server_side->maxseg_enable, server_side->maxseg_vaule, server_side->nodelay, server_side->ttl, server_side->keepalive, server_side->keepcnt, server_side->keepidle, server_side->keepidle, server_side->user_timeout);
    profile_param_free(profile_param);
    intercept_param_free(policy_param);
}