#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>

#include "mpack.h"
#include "tfe_cmsg.h"
#include "tfe_utils.h"
#include "tfe_ctrl_packet.h"

enum ctr_pkt_index
{
    INDEX_TSYNC = 0,
    INDEX_SESSION_ID,
    INDEX_STATE,
    INDEX_METHOD,
    INDEX_KEY_SCE,
    INDEX_VALUE_SCE,
    INDEX_KEY_SHAPER,
    INDEX_VALUE_SHAPER,
    INDEX_KEY_PROXY,
    INDEX_VALUE_PROXY,
    INDEX_MAX
};

enum {
    MPACK_ARRAY_INIT,
    MPACK_ARRAY_FQDN_IDS,
    MPACK_ARRAY_SEQ_SIDS,
    MPACK_ARRAY_ACK_SIDS,
    MPACK_ARRAY_SEQ_ROUTE_CTX,
    MPACK_ARRAY_ACK_ROUTE_CTX,
};

struct mpack_mmap_id2type
{
    int id;
    enum tfe_cmsg_tlv_type type;
    const char *str_name;
    int size;
    int array_index;
}mpack_table[] = {
    {.id = 0, .type = TFE_CMSG_TCP_RESTORE_SEQ, .str_name = "TFE_CMSG_TCP_RESTORE_SEQ", .size = 4, .array_index = MPACK_ARRAY_INIT},
    {.id = 1, .type = TFE_CMSG_TCP_RESTORE_ACK, .str_name = "TFE_CMSG_TCP_RESTORE_ACK", .size = 4, .array_index = MPACK_ARRAY_INIT},
    {.id = 2, .type = TFE_CMSG_TCP_RESTORE_MSS_CLIENT, .str_name = "TFE_CMSG_TCP_RESTORE_MSS_CLIENT", .size = 2, .array_index = MPACK_ARRAY_INIT},
    {.id = 3, .type = TFE_CMSG_TCP_RESTORE_MSS_SERVER, .str_name = "TFE_CMSG_TCP_RESTORE_MSS_SERVER", .size = 2, .array_index = MPACK_ARRAY_INIT},
    {.id = 4, .type = TFE_CMSG_TCP_RESTORE_WSACLE_CLIENT, .str_name = "TFE_CMSG_TCP_RESTORE_WSACLE_CLIENT", .size = 1, .array_index = MPACK_ARRAY_INIT},
    {.id = 5, .type = TFE_CMSG_TCP_RESTORE_WSACLE_SERVER, .str_name = "TFE_CMSG_TCP_RESTORE_WSACLE_SERVER", .size = 1, .array_index = MPACK_ARRAY_INIT},
    {.id = 6, .type = TFE_CMSG_TCP_RESTORE_SACK_CLIENT, .str_name = "TFE_CMSG_TCP_RESTORE_SACK_CLIENT", .size = 1, .array_index = MPACK_ARRAY_INIT},
    {.id = 7, .type = TFE_CMSG_TCP_RESTORE_SACK_SERVER, .str_name = "TFE_CMSG_TCP_RESTORE_SACK_SERVER", .size = 1, .array_index = MPACK_ARRAY_INIT},
    {.id = 8, .type = TFE_CMSG_TCP_RESTORE_TS_CLIENT, .str_name = "TFE_CMSG_TCP_RESTORE_TS_CLIENT", .size = 1, .array_index = MPACK_ARRAY_INIT},
    {.id = 9, .type = TFE_CMSG_TCP_RESTORE_TS_SERVER, .str_name = "TFE_CMSG_TCP_RESTORE_TS_SERVER", .size = 1, .array_index = MPACK_ARRAY_INIT},
    {.id = 10, .type = TFE_CMSG_TCP_RESTORE_PROTOCOL, .str_name = "TFE_CMSG_TCP_RESTORE_PROTOCOL", .size = 1, .array_index = MPACK_ARRAY_INIT},
    {.id = 11, .type = TFE_CMSG_TCP_RESTORE_WINDOW_CLIENT, .str_name = "TFE_CMSG_TCP_RESTORE_WINDOW_CLIENT", .size = 2, .array_index = MPACK_ARRAY_INIT},
    {.id = 12, .type = TFE_CMSG_TCP_RESTORE_WINDOW_SERVER, .str_name = "TFE_CMSG_TCP_RESTORE_WINDOW_SERVER", .size = 2, .array_index = MPACK_ARRAY_INIT},
    {.id = 13, .type = TFE_CMSG_TCP_RESTORE_TS_CLIENT_VAL, .str_name = "TFE_CMSG_TCP_RESTORE_TS_CLIENT_VAL", .size = 4, .array_index = MPACK_ARRAY_INIT},
    {.id = 14, .type = TFE_CMSG_TCP_RESTORE_TS_SERVER_VAL, .str_name = "TFE_CMSG_TCP_RESTORE_TS_SERVER_VAL", .size = 4, .array_index = MPACK_ARRAY_INIT},
    {.id = 15, .type = TFE_CMSG_TCP_RESTORE_INFO_PACKET_CUR_DIR, .str_name = "TFE_CMSG_TCP_RESTORE_INFO_PACKET_CUR_DIR", .size = 1, .array_index = MPACK_ARRAY_INIT}, 
    {.id = 16, .type = TFE_CMSG_SRC_SUB_ID, .str_name = "TFE_CMSG_SRC_SUB_ID", .size = 256, .array_index = MPACK_ARRAY_INIT},              
    {.id = 17, .type = TFE_CMSG_DST_SUB_ID, .str_name = "TFE_CMSG_DST_SUB_ID", .size = 256, .array_index = MPACK_ARRAY_INIT},              
    {.id = 18, .type = TFE_CMSG_SRC_ASN, .str_name = "TFE_CMSG_SRC_ASN", .size = 64, .array_index = MPACK_ARRAY_INIT},                 
    {.id = 19, .type = TFE_CMSG_DST_ASN, .str_name = "TFE_CMSG_DST_ASN", .size = 64, .array_index = MPACK_ARRAY_INIT},                 
    {.id = 20, .type = TFE_CMSG_SRC_ORGANIZATION, .str_name = "TFE_CMSG_SRC_ORGANIZATION", .size = 256, .array_index = MPACK_ARRAY_INIT},        
    {.id = 21, .type = TFE_CMSG_DST_ORGANIZATION, .str_name = "TFE_CMSG_DST_ORGANIZATION", .size = 256, .array_index = MPACK_ARRAY_INIT},        
    {.id = 22, .type = TFE_CMSG_SRC_IP_LOCATION_COUNTRY, .str_name = "TFE_CMSG_SRC_IP_LOCATION_COUNTRY", .size = 256, .array_index = MPACK_ARRAY_INIT}, 
    {.id = 23, .type = TFE_CMSG_DST_IP_LOCATION_COUNTRY, .str_name = "TFE_CMSG_DST_IP_LOCATION_COUNTRY", .size = 256, .array_index = MPACK_ARRAY_INIT}, 
    {.id = 24, .type = TFE_CMSG_SRC_IP_LOCATION_PROVINE, .str_name = "TFE_CMSG_SRC_IP_LOCATION_PROVINE", .size = 256, .array_index = MPACK_ARRAY_INIT}, 
    {.id = 25, .type = TFE_CMSG_DST_IP_LOCATION_PROVINE, .str_name = "TFE_CMSG_DST_IP_LOCATION_PROVINE", .size = 256, .array_index = MPACK_ARRAY_INIT}, 
    {.id = 26, .type = TFE_CMSG_SRC_IP_LOCATION_CITY, .str_name = "TFE_CMSG_SRC_IP_LOCATION_CITY", .size = 256, .array_index = MPACK_ARRAY_INIT},    
    {.id = 27, .type = TFE_CMSG_DST_IP_LOCATION_CITY, .str_name = "TFE_CMSG_DST_IP_LOCATION_CITY", .size = 256, .array_index = MPACK_ARRAY_INIT},    
    {.id = 28, .type = TFE_CMSG_SRC_IP_LOCATION_SUBDIVISION, .str_name = "TFE_CMSG_SRC_IP_LOCATION_SUBDIVISION", .size = 256, .array_index = MPACK_ARRAY_INIT},
    {.id = 29, .type = TFE_CMSG_DST_IP_LOCATION_SUBDIVISION, .str_name = "TFE_CMSG_DST_IP_LOCATION_SUBDIVISION", .size = 256, .array_index = MPACK_ARRAY_INIT},
    {.id = 30, .type = TFE_CMSG_SSL_CLIENT_JA3_FINGERPRINT, .str_name = "TFE_CMSG_SSL_CLIENT_JA3_FINGERPRINT", .size = 32, .array_index = MPACK_ARRAY_INIT}, 
    {.id = 31, .type = TFE_CMSG_FQDN_CAT_ID_VAL, .str_name = "TFE_CMSG_FQDN_CAT_ID_VAL", .size = 4, .array_index = MPACK_ARRAY_FQDN_IDS},
    {.id = 32, .type = TFE_CMSG_TLV_NR_MAX, .str_name = "TFE_SEQ_SIDS", .size = 2, .array_index = MPACK_ARRAY_SEQ_SIDS},
    {.id = 33, .type = TFE_CMSG_TLV_NR_MAX, .str_name = "TFE_ACK_SIDS", .size = 2, .array_index = MPACK_ARRAY_ACK_SIDS},
    {.id = 34, .type = TFE_CMSG_TLV_NR_MAX, .str_name = "TFE_SEQ_ROUTE_CTX", .size = 1, .array_index = MPACK_ARRAY_SEQ_ROUTE_CTX},
    {.id = 35, .type = TFE_CMSG_TLV_NR_MAX, .str_name = "TFE_ACK_ROUTE_CTX", .size = 1, .array_index = MPACK_ARRAY_ACK_ROUTE_CTX}
};

static int fqdn_id_set_cmsg(struct ctrl_pkt_parser *handler, mpack_node_t node, int table_index)
{
    uint32_t fqdn_val[8] = {0};

    uint32_t array_cnt = mpack_node_array_length(node);
    tfe_cmsg_set(handler->cmsg, TFE_CMSG_FQDN_CAT_ID_NUM, (const unsigned char *)&array_cnt, sizeof(uint32_t));
    TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu array fqdn_id num: [%u]", LOG_TAG_CTRLPKT, handler->session_id, array_cnt);
    for (uint32_t i = 0; i < array_cnt; i++) {
        fqdn_val[i] = mpack_node_u32(mpack_node_array_at(node, i));
        TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu array fqdn_id msgpack cmsg: [%s] -> [%lu]", LOG_TAG_CTRLPKT, handler->session_id, mpack_table[table_index].str_name, fqdn_val[i]);
    }
    tfe_cmsg_set(handler->cmsg, TFE_CMSG_FQDN_CAT_ID_VAL, (const unsigned char*)fqdn_val, array_cnt * sizeof(uint32_t));
    return 0;
}

static int sids_array_parse_mpack(struct ctrl_pkt_parser *handler, mpack_node_t node, int table_index, int is_seq)
{
    struct sids *sid = is_seq ? &handler->seq_sids : &handler->ack_sids;
    if (mpack_node_array_length(node) > MR_SID_LIST_MAXLEN) {
        TFE_LOG_ERROR(g_default_logger, "%s: session: %lu sid[%ld] more than maxlen %d", LOG_TAG_CTRLPKT, handler->session_id, mpack_node_array_length(node), MR_SID_LIST_MAXLEN);
        return -1;
    }

    sid->num = mpack_node_array_length(node);
    TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu array sids num: [%d]", LOG_TAG_CTRLPKT, handler->session_id, sid->num);
    for (int i = 0; i < sid->num; i++)
    {
        sid->elems[i] = mpack_node_u16(mpack_node_array_at(node, i));
    }
    return 0;
}

static int route_ctx_parse_mpack(struct ctrl_pkt_parser *handler, mpack_node_t node, int table_index, int is_seq)
{
    struct route_ctx *ctx = is_seq ? &handler->seq_route_ctx : &handler->ack_route_ctx;
    if (mpack_node_array_length(node) > 64) {
        TFE_LOG_ERROR(g_default_logger, "%s: session: %lu route ctx[%ld] more than maxlen 64", LOG_TAG_CTRLPKT, handler->session_id, mpack_node_array_length(node));
        return -1;
    }

    ctx->len = mpack_node_array_length(node);
    TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu array route ctx num: [%u]", LOG_TAG_CTRLPKT, handler->session_id, ctx->len);
    for (int i = 0; i < ctx->len; i++)
    {
        ctx->data[i] = mpack_node_u8(mpack_node_array_at(node, i));
    }
    return 0;
}

static int proxy_parse_messagepack(mpack_node_t node, void *ctx)
{
    uint64_t value = 0;
    struct ctrl_pkt_parser *handler = (struct ctrl_pkt_parser *)ctx;

    if (mpack_node_is_nil(mpack_node_map_cstr(node, "rule_ids")))
    {
        TFE_LOG_ERROR(g_default_logger, "%s: unexpected control packet: (rule_ids no found)", LOG_TAG_CTRLPKT);
        return -1;
    }
    handler->tfe_policy_id_num = mpack_node_array_length(mpack_node_map_cstr(node, "rule_ids"));
    for (int i = 0; i < handler->tfe_policy_id_num; i++) {
        handler->tfe_policy_ids[i] = mpack_node_u64(mpack_node_array_at(mpack_node_map_cstr(node, "rule_ids"), i));
    }

    if (handler->tfe_policy_id_num) {
        tfe_cmsg_set(handler->cmsg, TFE_CMSG_POLICY_ID, (const unsigned char *)&handler->tfe_policy_ids[0], sizeof(uint64_t));
        TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu tfe policy id num: [%d]", LOG_TAG_CTRLPKT, handler->session_id, handler->tfe_policy_id_num);
        for (int i = 0; i < handler->tfe_policy_id_num; i++) {
            TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu policy id:%lu ", LOG_TAG_CTRLPKT, handler->session_id, handler->tfe_policy_ids[i]);
        }
    }

    mpack_node_t tcp_handshake = mpack_node_map_cstr(node, "tcp_handshake");
    int cmsg_array_cnt = mpack_node_array_length(tcp_handshake);
    for (int i = 0; i < cmsg_array_cnt; i++) {
        mpack_node_t ptr = mpack_node_array_at(tcp_handshake, i);
        
        switch (mpack_node_type(ptr)) {
            case mpack_type_uint:
                value = mpack_node_u64(ptr);
                tfe_cmsg_set(handler->cmsg, mpack_table[i].type, (const unsigned char *)&value, mpack_table[i].size);
                TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu interger msgpack cmsg: [%s] -> [%lu]", LOG_TAG_CTRLPKT, handler->session_id, mpack_table[i].str_name, value);
                break;
            case mpack_type_str:
                tfe_cmsg_set(handler->cmsg, mpack_table[i].type, (const unsigned char *)mpack_node_str(ptr), mpack_node_strlen(ptr));
                TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu string msgpack cmsg: [%s] -> [%s]", LOG_TAG_CTRLPKT, handler->session_id, mpack_table[i].str_name, mpack_node_str(ptr));
                break;
            case mpack_type_nil:
                TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu msgpack cmsg: [%s] -> [nil]", LOG_TAG_CTRLPKT, handler->session_id, mpack_table[i].str_name);
                break;
            case mpack_type_array:
                switch(mpack_table[i].array_index)
                {
                    case MPACK_ARRAY_FQDN_IDS:
                        fqdn_id_set_cmsg(handler, ptr, i);
                        break;
                    case MPACK_ARRAY_SEQ_SIDS:
                        sids_array_parse_mpack(handler, ptr, i, 1);
                        break;
                    case MPACK_ARRAY_ACK_SIDS:
                        sids_array_parse_mpack(handler, ptr, i, 0);
                        break;
                    case MPACK_ARRAY_SEQ_ROUTE_CTX:
                        route_ctx_parse_mpack(handler, ptr, i, 1);
                        break;
                    case MPACK_ARRAY_ACK_ROUTE_CTX:
                        route_ctx_parse_mpack(handler, ptr, i, 0);
                        break;
                    default:
                        break;
                }
                break;
            default:
                break;
        }
    }
    return 0;
}

int parse_messagepack(const char* data, size_t length, void *ctx)
{
    struct ctrl_pkt_parser *handler = (struct ctrl_pkt_parser *)ctx;
    char buff[16] = {0};
    mpack_node_t params;
    mpack_node_t sce_map;
    mpack_node_t proxy_map;
    mpack_tree_t tree;
    mpack_tree_init_data(&tree, data, length);
    mpack_tree_parse(&tree);
    mpack_node_t root = mpack_tree_root(&tree);
    if (mpack_node_is_nil(root))
    {
        TFE_LOG_ERROR(g_default_logger, "%s: unexpected control packet: (invalid mpack format)", LOG_TAG_CTRLPKT);
        goto error;
    }

    if (mpack_node_is_nil(mpack_node_map_cstr(root, "tsync")))
    {
        TFE_LOG_ERROR(g_default_logger, "%s: unexpected control packet: (tsync no found)", LOG_TAG_CTRLPKT);
        goto error;
    }
    mpack_node_copy_cstr(mpack_node_map_cstr(root, "tsync"), handler->tsync, sizeof(handler->tsync));
    if (strcmp(handler->tsync, "2.0") != 0)
    {
        TFE_LOG_ERROR(g_default_logger, "%s: unexpected control packet: (invalid tsync version) %s", LOG_TAG_CTRLPKT, handler->tsync);
        goto error;
    }

    if (mpack_node_is_nil(mpack_node_map_cstr(root, "session_id")))
    {
        TFE_LOG_ERROR(g_default_logger, "%s: unexpected control packet: (session_id no found)", LOG_TAG_CTRLPKT);
        goto error;
    }
    handler->session_id = mpack_node_u64(mpack_node_map_cstr(root, "session_id"));

    
    if (mpack_node_is_nil(mpack_node_map_cstr(root, "state")))
    {
        TFE_LOG_ERROR(g_default_logger, "%s: unexpected control packet: (state no found)", LOG_TAG_CTRLPKT);
        goto error;
    }
    mpack_node_copy_cstr(mpack_node_map_cstr(root, "state"), buff, sizeof(buff));
    if (strncasecmp(buff, "opening", sizeof(buff)) == 0)
    {
        handler->state = SESSION_STATE_OPENING;
        goto succ;
    }
    else if (strncasecmp(buff, "active", sizeof(buff)) == 0)
    {
        handler->state = SESSION_STATE_ACTIVE;
    }
    else if (strncasecmp(buff, "closing", sizeof(buff)) == 0)
    {
        handler->state = SESSION_STATE_CLOSING;
        goto succ;
    }
    else if (strncasecmp(buff, "resetall", sizeof(buff)) == 0)
    {
        handler->state = SESSION_STATE_RESETALL;
        goto succ;
    }
    else
    {
        TFE_LOG_DEBUG(g_default_logger, "%s: unexpected control packet: (invalid state value) %s", LOG_TAG_CTRLPKT, buff);
    }

    if (mpack_node_is_nil(mpack_node_map_cstr(root, "method")))
    {
        TFE_LOG_ERROR(g_default_logger, "%s: unexpected control packet: (method no found)", LOG_TAG_CTRLPKT);
        goto error;
    }
    mpack_node_copy_cstr(mpack_node_map_cstr(root, "method"), handler->method, sizeof(handler->method));

    if (mpack_node_is_nil(mpack_node_map_cstr(root, "params")))
    {
        TFE_LOG_ERROR(g_default_logger, "%s: unexpected control packet: (params no found)", LOG_TAG_CTRLPKT);
        goto error;
    }
    params = mpack_node_map_cstr(root, "params");

    if (!mpack_node_is_missing(mpack_node_map_str_optional(params, "sce", strlen("sce"))))
    {
        sce_map = mpack_node_map_cstr(params, "sce");
        if (mpack_node_is_nil(mpack_node_map_cstr(sce_map, "rule_ids")))
        {
            TFE_LOG_ERROR(g_default_logger, "%s: unexpected control packet: (sce rule_ids no found)", LOG_TAG_CTRLPKT);
            goto error;
        }
        handler->sce_policy_id_num = mpack_node_array_length(mpack_node_map_cstr(sce_map, "rule_ids"));
        for (int i = 0; i < handler->sce_policy_id_num; i++) {
            handler->sce_policy_ids[i] = mpack_node_u64(mpack_node_array_at(mpack_node_map_cstr(sce_map, "rule_ids"), i));
        }

        if (handler->sce_policy_id_num) {
            TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu sce policy id num: [%d]", LOG_TAG_CTRLPKT, handler->session_id, handler->sce_policy_id_num);
            for (int i = 0; i < handler->sce_policy_id_num; i++) {
                TFE_LOG_DEBUG(g_default_logger, "%s: session: %lu policy id:%lu ", LOG_TAG_CTRLPKT, handler->session_id, handler->sce_policy_ids[i]);
            }
        }
    }

    if (mpack_node_is_missing(mpack_node_map_str_optional(params, "proxy", strlen("proxy"))))
    {
        TFE_LOG_ERROR(g_default_logger, "%s: unexpected control packet: (proxy no found)", LOG_TAG_CTRLPKT);
        goto error;
    }

    proxy_map = mpack_node_map_cstr(params, "proxy");
    proxy_parse_messagepack(proxy_map, handler);

succ:
    mpack_tree_destroy(&tree);
    return 0;
error:
    mpack_tree_destroy(&tree);
    return -1;
}