#include <rte_graph.h>
#include <rte_graph_worker.h>

#include <MESA_prof_load.h>
#include <cJSON.h>
#include <mrb_define.h>
#include <sc_common.h>
#include <sc_node_common.h>
#include <sc_trace.h>

/* Forwarder next node */
enum
{
    FORWARDER_PKT_DROP = 0,
    FORWARDER_NEXT_LB,
    FORWARDER_VWIRE_EGRESS,
    FORWARDER_EF_EGRESS,
    FORWARDER_TERA_EGRESS,
    FORWARDER_LAI,
    FORWARDER_NEXT_MAX
};

/* Forwarder drop reason */
enum forwarder_drop_reason
{
    FORWARDER_DROP_RSN_POP_SID_ERR = 0,
    FORWARDER_DROP_RSN_INVALID_SID,
    FORWARDER_DROP_RSN_MAX
};

/* Forwarder drop reason string */
static const char * forwarder_drop_reason_str[FORWARDER_DROP_RSN_MAX] = {"drop_rsn_pop_sid_err",
                                                                         "drop_rsn_invalid_sid"};

/* Forwarder stats */
struct forwarder_stats
{
    volatile uint64_t total_pkts;
    volatile uint64_t pkts_per_batch;
    volatile uint64_t to_load_balance;
    volatile uint64_t to_vwire_egress;
    volatile uint64_t to_ef_egress;
    volatile uint64_t to_tera_egress;
    volatile uint64_t to_lai;
    volatile uint64_t drop_reason[FORWARDER_DROP_RSN_MAX];
} __rte_cache_aligned;

struct forwarder_main
{
    uint16_t nr_rules;
    uint16_t forwarder_table[UINT16_MAX];
};

/* Global forwarder main */
static struct forwarder_main g_forwarder_main = {};
static struct forwarder_stats stats_per_graph[RTE_MAX_LCORE] = {};

static inline struct forwarder_main * forwarder_main_get(void)
{
    return &g_forwarder_main;
}

/* Sid check */
int sid_check(uint32_t sid)
{
    struct forwarder_main * forwarder_main = forwarder_main_get();

    if (sid >= RTE_DIM(forwarder_main->forwarder_table))
        return RT_ERR;

    if (forwarder_main->forwarder_table[sid] == 0)
        return RT_ERR;

    return RT_SUCCESS;
}

/* Forwarder table inserter */
void forwarder_table_insert(uint16_t sid, uint16_t type)
{
    char str_forwarder[1024] = {};
    struct forwarder_main * forwarder_main = forwarder_main_get();
    int len = snprintf(str_forwarder, sizeof(str_forwarder), "New entry in forwarder table created, sid=%u", sid);

    switch (type)
    {
    case FORWARDER_TYPE_LB:
        forwarder_main->forwarder_table[sid] = FORWARDER_NEXT_LB;
        len += snprintf(str_forwarder + len, sizeof(str_forwarder) - len, ", next_services=load balance");
        break;
    case FORWARDER_TYPE_VWIRE:
        forwarder_main->forwarder_table[sid] = FORWARDER_VWIRE_EGRESS;
        len += snprintf(str_forwarder + len, sizeof(str_forwarder) - len, ", next_services=vwire egress");
        break;
    case FORWARDER_TYPE_EF:
        forwarder_main->forwarder_table[sid] = FORWARDER_EF_EGRESS;
        len += snprintf(str_forwarder + len, sizeof(str_forwarder) - len, ", next_services=etherfabric egress");
        break;
    case FORWARDER_TYPE_TERA:
        forwarder_main->forwarder_table[sid] = FORWARDER_TERA_EGRESS;
        len += snprintf(str_forwarder + len, sizeof(str_forwarder) - len, ", next_services=tera egress");
        break;
    case FORWARDER_TYPE_LAI:
        forwarder_main->forwarder_table[sid] = FORWARDER_LAI;
        len += snprintf(str_forwarder + len, sizeof(str_forwarder) - len, ", next_services=link aware injector");
        break;
    }

    forwarder_main->nr_rules++;
    MR_INFO("%s", str_forwarder);
}

/* Generate and store the trace information */
static __rte_always_inline void gen_store_trace_info(struct rte_node * node, struct rte_mbuf * mbuf,
                                                     uint16_t next_node_index, struct forwarder_stats * stats,
                                                     enum forwarder_drop_reason drop_reason)
{
    /* Populate the next node infomation */
    char str_record[MR_STRING_MAX];
    int len = snprintf(str_record, sizeof(str_record), "next node:%s", node->nodes[next_node_index]->name);

    /* Populate the reason for next node */
    struct mrb_metadata * mrb_meta = mrbuf_cz_data(mbuf, MR_NODE_CTRLZONE_ID);
    uint16_t cur_sid = mrb_meta->cur_sid;

    if (unlikely(next_node_index == FORWARDER_PKT_DROP))
    {
        assert(drop_reason < FORWARDER_DROP_RSN_MAX);
        len += snprintf(str_record + len, sizeof(str_record) - len, ", rsn:%s", forwarder_drop_reason_str[drop_reason]);
    }
    else
    {
        len += snprintf(str_record + len, sizeof(str_record) - len, ", pop sid:%u", cur_sid);
    }

    /* Populate the sids information */
    len += embed_sid_info(mbuf, str_record + len, sizeof(str_record) - len);

    /* Emit the trace record */
    struct dp_trace_record_meta meta = {
        .measurement_type = DP_TRACE_MEASUREMENT_TYPE_TRACE, .appsym = MR_TRACE_APPSYM, .module = node->name};
    dp_trace_record_emit_str(sc_main_get()->trace, mbuf, rte_lcore_id(), &meta, str_record);
}

/* Forwarder node process function */
static __rte_always_inline uint16_t forwarder_node_process(struct rte_graph * graph, struct rte_node * node,
                                                           void ** objs, uint16_t cnt)
{
    uint16_t last_spec = 0, next_node_index = 0;
    struct rte_mbuf * mbuf;

    /* Get pkts num and pkts buffer */
    struct forwarder_main * forwarder_main = forwarder_main_get();
    uint16_t batch_next_node_index = FORWARDER_NEXT_LB;
    uint16_t n_left_from = cnt;
    struct rte_mbuf ** pkts = (struct rte_mbuf **)objs;
    void ** batch_pkts = objs;
    struct forwarder_stats stats = {};
    enum forwarder_drop_reason drop_reason = FORWARDER_DROP_RSN_MAX;

    unsigned int next_node_index_counters[FORWARDER_NEXT_MAX] = {};

    /* Single packet processing */
    while (n_left_from > 0)
    {
        mbuf = pkts[0];
        pkts += 1;
        n_left_from -= 1;

        /* Get ctrlzone */
        struct mrb_metadata * mrb_metadata = mrbuf_cz_data(mbuf, MR_NODE_CTRLZONE_ID);

        /* Get sid */
        uint16_t sid;
        int ret = sid_list_pop(&mrb_metadata->sid_list, &sid);
        if (unlikely(ret == RT_ERR))
        {
            /* dump the packet when pop sid failed */
            /*
            char str_metadata[256] = {};
            sprintf(str_metadata, METADATA_INFO, mrb_metadata->dir, mrb_metadata->packet_create_from_nf,
                    mrb_metadata->link_id, mrb_metadata->is_ctrlbuf, mrb_metadata->payload_offset,
                    mrb_metadata->session_id, mrb_metadata->start_sid, mrb_metadata->nr_sid, mrb_metadata->cur_sid,
                    mrb_metadata->sids[0], mrb_metadata->sids[1], mrb_metadata->sids[2], mrb_metadata->sids[3],
                    mrb_metadata->sids[4], mrb_metadata->sids[5], mrb_metadata->sids[6], mrb_metadata->sids[7],
                    mrb_metadata->port_ingress, mrb_metadata->port_egress, mrb_metadata->link_db_index);
            MR_ERROR("The pkt metadata : %s  .", str_metadata);
            rte_pktmbuf_dump(stderr, mbuf, rte_pktmbuf_data_len(mbuf));
            */

            drop_reason = FORWARDER_DROP_RSN_POP_SID_ERR;
            stats.drop_reason[drop_reason]++;
            next_node_index = FORWARDER_PKT_DROP;
            goto node_enqueue;
        }

        /* According sid get next node index */
        assert(sid < RTE_DIM(forwarder_main->forwarder_table));
        next_node_index = forwarder_main->forwarder_table[sid];
        assert(next_node_index < FORWARDER_NEXT_MAX);

        next_node_index_counters[next_node_index] += 1;

        /* Deal invalid sid */
        if (unlikely(next_node_index == 0))
        {
            drop_reason = FORWARDER_DROP_RSN_INVALID_SID;
            stats.drop_reason[drop_reason]++;
            next_node_index = FORWARDER_PKT_DROP;
            goto node_enqueue;
        }

        /* Save current sid */
        mrb_metadata->cur_sid = sid;

    node_enqueue:
        /* Check if tracing is enabled for the current Mbuf */
        if (unlikely(dp_trace_record_can_emit(mbuf, DP_TRACE_MEASUREMENT_TYPE_TRACE)))
        {
            gen_store_trace_info(node, mbuf, next_node_index, &stats, drop_reason);
            // gen_store_trace_info_sid_list(node, mbuf);
        }

        /* Judge the next index whether to change */
        if (unlikely(batch_next_node_index != next_node_index))
        {
            /* If the next index has been changed,enqueue last pkts */
            rte_node_enqueue(graph, node, batch_next_node_index, batch_pkts, last_spec);
            batch_pkts += last_spec;
            last_spec = 1;
            batch_next_node_index = next_node_index;
        }
        else
        {
            /* If the next index not change, update the lasts */
            last_spec++;
        }
    }

    /* Update graph stats */
    struct forwarder_stats * graph_stats = &stats_per_graph[graph->id];
    graph_stats->total_pkts += cnt;
    graph_stats->pkts_per_batch = cnt;
    graph_stats->to_load_balance += next_node_index_counters[FORWARDER_NEXT_LB];
    graph_stats->to_vwire_egress += next_node_index_counters[FORWARDER_VWIRE_EGRESS];
    graph_stats->to_ef_egress += next_node_index_counters[FORWARDER_EF_EGRESS];
    graph_stats->to_tera_egress += next_node_index_counters[FORWARDER_TERA_EGRESS];
    graph_stats->to_lai += next_node_index_counters[FORWARDER_LAI];

    for (uint16_t i = 0; i < FORWARDER_DROP_RSN_MAX; i++)
    {
        graph_stats->drop_reason[i] += stats.drop_reason[i];
    }

    /* Process the remaining packets */
    if (likely(last_spec > 0))
    {
        rte_node_enqueue(graph, node, batch_next_node_index, batch_pkts, last_spec);
    }

    return cnt;
}

/* Forwarder node base */
static struct rte_node_register forwarder_node_base = {
    .process = forwarder_node_process,
    .name = "forwarder",
    .init = NULL,
    .nb_edges = FORWARDER_NEXT_MAX,
    .next_nodes =
        {
            [FORWARDER_PKT_DROP] = "pkt_drop_trap",
            [FORWARDER_NEXT_LB] = "lb",
            [FORWARDER_VWIRE_EGRESS] = "vwire_egress",
            [FORWARDER_EF_EGRESS] = "ef_egress",
            [FORWARDER_TERA_EGRESS] = "tera_egress",
            [FORWARDER_LAI] = "link_aware_injector",
        },
};

RTE_NODE_REGISTER(forwarder_node_base);

/************************************** Forwarder statistics **************************************/
cJSON * forwarder_node_monit_loop(struct sc_main * sc)
{
    cJSON * json_root = cJSON_CreateObject();
    unsigned int nr_graphs = sc->nr_io_thread;

    uint64_t total_pkts[nr_graphs];
    uint64_t pkts_per_batch[nr_graphs];
    uint64_t to_load_balance[nr_graphs];
    uint64_t to_vwire_egress[nr_graphs];
    uint64_t to_ef_egress[nr_graphs];
    uint64_t to_tera_egress[nr_graphs];
    uint64_t to_lai[nr_graphs];
    uint64_t drop_reason[FORWARDER_DROP_RSN_MAX][nr_graphs];

    for (uint32_t graph_id = 0; graph_id < nr_graphs; graph_id++)
    {
        struct forwarder_stats * stats = &stats_per_graph[graph_id];
        if (stats->total_pkts == 0)
        {
            total_pkts[graph_id] = 0;
            pkts_per_batch[graph_id] = 0;
            to_load_balance[graph_id] = 0;
            to_vwire_egress[graph_id] = 0;
            to_ef_egress[graph_id] = 0;
            to_tera_egress[graph_id] = 0;
            to_lai[graph_id] = 0;

            for (int i = 0; i < FORWARDER_DROP_RSN_MAX; i++)
            {
                drop_reason[i][graph_id] = 0;
            }

            continue;
        }

        total_pkts[graph_id] = stats->total_pkts;
        pkts_per_batch[graph_id] = stats->pkts_per_batch;
        to_load_balance[graph_id] = stats->to_load_balance;
        to_vwire_egress[graph_id] = stats->to_vwire_egress;
        to_ef_egress[graph_id] = stats->to_ef_egress;
        to_tera_egress[graph_id] = stats->to_tera_egress;
        to_lai[graph_id] = stats->to_lai;

        for (int i = 0; i < FORWARDER_DROP_RSN_MAX; i++)
        {
            drop_reason[i][graph_id] = stats->drop_reason[i];
        }
    }

    cJSON * json_total_pkts = create_uint64_array(total_pkts, nr_graphs);
    cJSON_AddItemToObject(json_root, "forwarder, total_pkts", json_total_pkts);

    cJSON * json_pkts_per_batch = create_uint64_array(pkts_per_batch, nr_graphs);
    cJSON_AddItemToObject(json_root, "forwarder, pkts_per_batch", json_pkts_per_batch);

    cJSON * json_to_load_balance = create_uint64_array(to_load_balance, nr_graphs);
    cJSON_AddItemToObject(json_root, "forwarder, to_load_balance", json_to_load_balance);

    cJSON * json_to_vwire_egress = create_uint64_array(to_vwire_egress, nr_graphs);
    cJSON_AddItemToObject(json_root, "forwarder, to_vwire_egress", json_to_vwire_egress);

    cJSON * json_to_ef_egress = create_uint64_array(to_ef_egress, nr_graphs);
    cJSON_AddItemToObject(json_root, "forwarder, to_etherfabric_egress", json_to_ef_egress);

    cJSON * json_to_tera_egress = create_uint64_array(to_tera_egress, nr_graphs);
    cJSON_AddItemToObject(json_root, "forwarder, to_tera_egress", json_to_tera_egress);

    cJSON * json_to_lai = create_uint64_array(to_lai, nr_graphs);
    cJSON_AddItemToObject(json_root, "forwarder, to_link_aware_injector", json_to_lai);

    for (int i = 0; i < FORWARDER_DROP_RSN_MAX; i++)
    {
        char str_title[MR_STRING_MAX];
        snprintf(str_title, sizeof(str_title), "forwarder, %s", forwarder_drop_reason_str[i]);

        cJSON * json_drop_reason = create_uint64_array(drop_reason[i], nr_graphs);
        cJSON_AddItemToObject(json_root, str_title, json_drop_reason);
    }

    return json_root;
}

cJSON * forwarder_table_monit_loop(struct sc_main * sc)
{
    uint32_t nr_items = 0;
    struct forwarder_main * forwarder_main = forwarder_main_get();
    cJSON * json_root = cJSON_CreateObject();

    for (int index = 0; index < RTE_DIM(forwarder_main->forwarder_table); index++)
    {
        if (forwarder_main->forwarder_table[index] == 0)
            continue;

        char str_info[MR_STRING_MAX] = {};
        int len = snprintf(str_info, sizeof(str_info), "Sid:%u, ", index);

        switch (forwarder_main->forwarder_table[index])
        {
        case FORWARDER_NEXT_LB:
            len += snprintf(str_info + len, sizeof(str_info) - len, "Next Services: LB");
            break;
        case FORWARDER_VWIRE_EGRESS:
            len += snprintf(str_info + len, sizeof(str_info) - len, "Next Services: Vwire");
            break;
        case FORWARDER_EF_EGRESS:
            len += snprintf(str_info + len, sizeof(str_info) - len, "Next Services: Etherfabric");
            break;
        case FORWARDER_TERA_EGRESS:
            len += snprintf(str_info + len, sizeof(str_info) - len, "Next Services: Tera");
            break;
        case FORWARDER_LAI:
            len += snprintf(str_info + len, sizeof(str_info) - len, "Next Services: Link Aware Injector");
            break;
        default:
            len += snprintf(str_info + len, sizeof(str_info) - len, "Next Services: Unknown");
            break;
        }

        char str_idx[MR_STRING_MAX] = {};
        snprintf(str_idx, sizeof(str_idx) - 1, "index:%u", nr_items);
        cJSON_AddStringToObject(json_root, str_idx, str_info);
        nr_items++;
    }

    cJSON_AddNumberToObject(json_root, "total_num", nr_items);
    return json_root;
}