#include <MESA_prof_load.h>
#include <link_db.h>
#include <rte_branch_prediction.h>
#include <sc_node.h>
#include <sc_node_common.h>
#include <sc_trace.h>

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

extern int ef_peer_index_get_by_tl_id_get(uint16_t traffic_link_id, uint16_t * out_ef_peer_index);

/* Link aware injector next node */
enum
{
    LAI_NEXT_PKT_DROP = 0,
    LAI_NEXT_EF_EGRESS,
    LAI_NEXT_VWIRE_EGRESS,
    LAI_NEXT_MAX
};

/* Link aware injector drop reason*/
enum lai_drop_reason
{
    LAI_DROP_RSN_TLID_MISS = 0,
    LAI_DROP_RSN_INVALID_ADAPTER_TYPE,
    LAI_DROP_RSN_EF_PEER_ID_MISS,
    LAI_DROP_RSN_EF_ADAPTER_ID_MISS,
    LAI_DROP_RSN_MAX
};

/* Link aware injector drop reason string */
char * lai_drop_rsn_str[LAI_DROP_RSN_MAX] = {"drop_rsn_traffic_link_id_miss", "drop_rsn_invalid_adapter_type",
                                             "drop_rsn_ef_peer_id_miss", "drop_rsn_ef_adapter_id_miss"};

/* Link aware injector stats struct */
struct lai_stats
{
    volatile uint64_t total_pkts;
    volatile uint64_t pkts_per_batch;
    volatile uint64_t to_ef_egress;
    volatile uint64_t to_vwire_egress;
    volatile uint64_t drop_reason[LAI_DROP_RSN_MAX];
} __rte_cache_aligned;

/* Link aware injector main struct */
struct node_lai_main
{
    uint32_t lai_sid;
    uint32_t ef_sid_start;
    uint32_t vwire_sid_start;
    uint16_t nr_ef_adapters;
    uint16_t * ef_adapter_ids;
    struct link_db_ctx * link_db_ctx;
};

/* Link aware injector node context */
struct lai_node_ctx
{
    /* Cached next index */
    uint16_t next_index;
};

#define LAI_NODE_LAST_NEXT(ctx) (((struct lai_node_ctx *)ctx)->next_index)

/* Global variable */
static struct node_lai_main g_lai_main = {};
static struct lai_stats lai_stats_per_graph[RTE_MAX_LCORE] = {};

static inline struct node_lai_main * node_lai_main_get()
{
    return &g_lai_main;
}

/* Init Link aware injector */
int lai_init(struct sc_main * sc)
{
    struct node_lai_main * lai_main = node_lai_main_get();
    memset(lai_main, 0, sizeof(struct node_lai_main));

    /* Get Link aware injector sid */
    MESA_load_profile_uint_def(sc->local_cfgfile, "feature_sids", "lai", &lai_main->lai_sid, 900);

    /* Link db ctx create */
    uint32_t max_entries;
    MESA_load_profile_uint_def(sc->local_cfgfile, "limits", "nr_max_link_dbs", &max_entries, 64);
    max_entries = max_entries * LINK_DB_TYPE_ALL;

    lai_main->link_db_ctx = link_db_create(LINK_DB_TYPE_ALL, max_entries);
    if (lai_main->link_db_ctx == NULL)
    {
        MR_ERROR("In link aware injector node, link db ctx create failed.");
        return RT_ERR;
    }

    if (link_db_config_parse(sc->local_cfgfile, lai_main->link_db_ctx) == RT_ERR)
    {
        MR_ERROR("In link aware injector node, link db config parse failed.");
        return RT_ERR;
    }

    /* Get ef and vwire sid start */
    lai_main->ef_sid_start = ef_sid_start_get();
    lai_main->vwire_sid_start = vwire_sid_start_get();

    if (lai_main->ef_sid_start == 0)
    {
        MR_ERROR("In link aware injector node, ef sid start get failed.");
        return RT_ERR;
    }
    else if (lai_main->vwire_sid_start == 0)
    {
        MR_ERROR("In link aware injector node, vwire sid start get failed.");
        return RT_ERR;
    }

    /* Inserter sid to forwarder table */
    forwarder_table_insert(lai_main->lai_sid, FORWARDER_TYPE_LAI);

    /* Get ef adapter ids */
    lai_main->nr_ef_adapters = ef_nr_adapters_get();
    if (lai_main->nr_ef_adapters != 0)
    {
        lai_main->ef_adapter_ids = ZMALLOC(sizeof(uint16_t) * lai_main->nr_ef_adapters);
        MR_VERIFY_MALLOC(lai_main->ef_adapter_ids);

        ef_adapter_ids_get(lai_main->ef_adapter_ids, lai_main->nr_ef_adapters);
        char str_ef_adapter_info[MR_STRING_MAX] = {0};
        int len = snprintf(str_ef_adapter_info, sizeof(str_ef_adapter_info),
                           "Link aware injector node, nr ef adapters:%u, ef adapter ids:", lai_main->nr_ef_adapters);
        for (int i = 0; i < lai_main->nr_ef_adapters; i++)
        {
            len += snprintf(str_ef_adapter_info + len, sizeof(str_ef_adapter_info) - len, " %u,",
                            lai_main->ef_adapter_ids[i]);
        }
        MR_INFO("%s", str_ef_adapter_info);
    }

    MR_INFO("Link aware injector node init success,link aware injector sid:%u, ef sid start: %u, vwire sid start: %u",
            lai_main->lai_sid, lai_main->ef_sid_start, lai_main->vwire_sid_start);

    return RT_SUCCESS;
}

/* 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, enum lai_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 drop reason */
    if (next_node_index == LAI_NEXT_PKT_DROP)
    {
        assert(drop_reason < LAI_DROP_RSN_MAX);
        len += snprintf(str_record + len, sizeof(str_record) - len, ", rsn:%s", lai_drop_rsn_str[drop_reason]);
    }

    struct mrb_metadata * mrb_meta = mrbuf_cz_data(mbuf, MR_NODE_CTRLZONE_ID);
    len += snprintf(str_record + len, sizeof(str_record) - len, ", trf lk id:%u, cur sid:%u", mrb_meta->traffic_link_id,
                    mrb_meta->cur_sid);

    /* 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);
}

/* Link aware injector node process function */
static __rte_always_inline uint16_t lai_node_process(struct rte_graph * graph, struct rte_node * node, void ** objs,
                                                     uint16_t nb_objs)
{
    void ** batch_pkts = objs;
    struct rte_mbuf ** mbufs = (struct rte_mbuf **)objs;

    /* prefetch the mbuf and it's content */
    rte_prefetch0(mbufs[0]);
    rte_prefetch0(mrbuf_cz_data(mbufs[0], MR_NODE_CTRLZONE_ID));

    uint16_t last_spec = 0;
    uint16_t n_left_from = nb_objs;
    uint16_t batch_next_node_index = LAI_NODE_LAST_NEXT(node->ctx);
    uint32_t last_traffic_link_id = UINT32_MAX;
    enum lai_drop_reason drop_reason = LAI_DROP_RSN_TLID_MISS;
    struct lai_stats stats = {};
    struct node_lai_main * lai_main = node_lai_main_get();
    struct link_db_reverse_result result = {.nr_adapters = 0, .match_result = RT_ERR};

    while (n_left_from > 0)
    {
        struct rte_mbuf * mbuf = mbufs[0];
        mbufs += 1;
        n_left_from -= 1;

        /* prefetch next mbuf */
        if (n_left_from > 0)
        {
            rte_prefetch0(mbufs[0]);
            rte_prefetch0(mrbuf_cz_data(mbufs[0], MR_NODE_CTRLZONE_ID));
        }

        /* Get metadata */
        uint16_t next_node_index;
        struct mrb_metadata * mrb_meta = mrbuf_cz_data(mbuf, MR_NODE_CTRLZONE_ID);

        /* Link db reverse match */
        if (last_traffic_link_id != mrb_meta->traffic_link_id)
        {
            last_traffic_link_id = mrb_meta->traffic_link_id;
            link_db_reverse_match(lai_main->link_db_ctx, last_traffic_link_id, &result);
        }

        /* Check the match result */
        if (unlikely(result.match_result != RT_SUCCESS))
        {
            drop_reason = LAI_DROP_RSN_TLID_MISS;
            stats.drop_reason[drop_reason]++;
            next_node_index = LAI_NEXT_PKT_DROP;
            goto node_enqueue;
        }

        /* Set the next node index */
        if (result.type == LINK_DB_TYPE_EF)
        {
            uint16_t ef_peer_index;
            int ret = ef_peer_index_get_by_tl_id_get(mrb_meta->traffic_link_id, &ef_peer_index);
            if (ret == RT_ERR)
            {
                drop_reason = LAI_DROP_RSN_EF_PEER_ID_MISS;
                stats.drop_reason[drop_reason]++;
                next_node_index = LAI_NEXT_PKT_DROP;
                goto node_enqueue;
            }

            uint16_t adapter_id;
            if (result.nr_adapters == 0)
            {
                if (lai_main->nr_ef_adapters == 0)
                {
                    drop_reason = LAI_DROP_RSN_EF_ADAPTER_ID_MISS;
                    stats.drop_reason[drop_reason]++;
                    next_node_index = LAI_NEXT_PKT_DROP;
                    goto node_enqueue;
                }
                else
                {
                    adapter_id = lai_main->ef_adapter_ids[mbuf->hash.usr % lai_main->nr_ef_adapters];
                }
            }
            else
            {
                adapter_id = result.adapter_id[mbuf->hash.usr % result.nr_adapters];
            }

            mrb_meta->cur_sid = lai_main->ef_sid_start + adapter_id;
            mrb_meta->ef_peer_index = ef_peer_index;
            mrb_meta->ef_link_id = result.ef_link_id;
            mrb_meta->packet_create_from_nf = 1;

            next_node_index = LAI_NEXT_EF_EGRESS;
            stats.to_ef_egress++;
        }
        else if (result.type == LINK_DB_TYPE_VWIRE)
        {
            mrb_meta->cur_sid = lai_main->vwire_sid_start + result.adapter_id[0];
            next_node_index = LAI_NEXT_VWIRE_EGRESS;
            stats.to_vwire_egress++;
        }
        else
        {
            drop_reason = LAI_DROP_RSN_INVALID_ADAPTER_TYPE;
            stats.drop_reason[drop_reason]++;
            next_node_index = LAI_NEXT_PKT_DROP;
        }

    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, drop_reason);
        }

        /* Check if the next index needs to be changed */
        if (unlikely(batch_next_node_index != next_node_index))
        {
            /* Enqueue the last packets if the next index has changed */
            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 hasn't changed, update the last packets */
            last_spec++;
        }
    }

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

    /* Update last next index */
    LAI_NODE_LAST_NEXT(node->ctx) = batch_next_node_index;

    /* Update stats */
    struct lai_stats * graph_stats = &lai_stats_per_graph[graph->id];
    graph_stats->total_pkts += nb_objs;
    graph_stats->pkts_per_batch = nb_objs;
    graph_stats->to_ef_egress += stats.to_ef_egress;
    graph_stats->to_vwire_egress += stats.to_vwire_egress;
    for (int i = 0; i < LAI_DROP_RSN_MAX; i++)
    {
        graph_stats->drop_reason[i] += stats.drop_reason[i];
    }

    return nb_objs;
}

/* Link aware injector node init function */
static int lai_node_init(const struct rte_graph * graph, struct rte_node * node)
{
    LAI_NODE_LAST_NEXT(node->ctx) = LAI_NEXT_EF_EGRESS;
    return 0;
}

/* Link aware injector node base */
static struct rte_node_register lai_node_base = {
    .process = lai_node_process,
    .name = "link_aware_injector",
    .init = lai_node_init,
    .nb_edges = LAI_NEXT_MAX,
    .next_nodes =
        {
            [LAI_NEXT_PKT_DROP] = "pkt_drop_trap",
            [LAI_NEXT_EF_EGRESS] = "ef_egress",
            [LAI_NEXT_VWIRE_EGRESS] = "vwire_egress",
        },
};

RTE_NODE_REGISTER(lai_node_base);

cJSON * lai_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_eth_egress[nr_graphs];
    uint64_t to_vwire_egress[nr_graphs];
    uint64_t drop_reason[LAI_DROP_RSN_MAX][nr_graphs];

    for (uint32_t graph_id = 0; graph_id < nr_graphs; graph_id++)
    {
        struct lai_stats * stats = &lai_stats_per_graph[graph_id];
        if (stats->total_pkts == 0)
        {
            total_pkts[graph_id] = 0;
            pkts_per_batch[graph_id] = 0;
            to_eth_egress[graph_id] = 0;
            to_vwire_egress[graph_id] = 0;
            for (int i = 0; i < LAI_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_eth_egress[graph_id] = stats->to_ef_egress;
        to_vwire_egress[graph_id] = stats->to_vwire_egress;
        for (int i = 0; i < LAI_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, "lai, total_pkts", json_total_pkts);

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

    cJSON * json_to_eth_egress = create_uint64_array(to_eth_egress, nr_graphs);
    cJSON_AddItemToObject(json_root, "lai, to_eth_egress", json_to_eth_egress);

    cJSON * json_to_vwire_egress = create_uint64_array(to_vwire_egress, nr_graphs);
    cJSON_AddItemToObject(json_root, "lai, to_vwire_egress", json_to_vwire_egress);
    for (int i = 0; i < LAI_DROP_RSN_MAX; i++)
    {
        char str_title[MR_STRING_MAX];
        snprintf(str_title, sizeof(str_title), "lai, %s", lai_drop_rsn_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;
}