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

#include <MESA_prof_load.h>
#include <sc_common.h>
#include <sc_mrb.h>
#include <sc_node_common.h>
#include <sc_trace.h>

/* Bridge Next Node */
enum
{
    BRIDGE_NEXT_DROP = 0,
    BRIDGE_NEXT_ETH_EGRESS,
    BRIDGE_NEXT_MAX
};

/* Bridge pkt match results */
enum bridge_match_result
{
    BRIDGE_MATCH_MISS = 0,
    BRIDGE_MATCH_TO_MASTER,
    BRIDGE_MATCH_BROADCAST,
    BRIDGE_MATCH_DST_MAC,
    BRIDGE_MATCH_MAX
};

/* Bridge pkt match results string */
static const char * bridge_match_result_str[BRIDGE_MATCH_MAX] = {"miss", "to_master", "broadcast", "match_dst_mac"};

/* Bridge stats struct */
struct bridge_stats
{
    volatile uint64_t total_pkts;
    volatile uint64_t pkts_per_batch;
    volatile uint64_t pkt_clone_failed;
    volatile uint64_t match_result[BRIDGE_MATCH_MAX];
} __rte_cache_aligned;

struct bridge
{
    uint8_t nr_devices;
    port_id_t master_port;
    port_id_t ports[64];
    struct rte_ether_addr ether_addr[64];
} __rte_cache_aligned;

/* Bridge main struct */
struct node_bridge_main
{
    uint8_t nr_bridge;
    struct rte_mempool * direct_pool;
    struct bridge bridge_table[64];
} __rte_cache_aligned;

/* Global bridge main */
static struct node_bridge_main g_bridge_main = {};
static struct bridge_stats stats_per_graph[RTE_MAX_LCORE] = {};
static uint8_t last_bridge_index[RTE_MAX_LCORE][64] = {};

static inline struct node_bridge_main * node_bridge_main_get(void)
{
    return &g_bridge_main;
}

/************************************* Bridge config **************************************/
/* Parser the bridge config */
int parser_bridge_conf(struct sc_main * sc, struct node_bridge_main * bridge_main)
{
    uint32_t nr_bridge = 0;

    /* Parsing all config */
    for (int index = 0; index < RTE_DIM(bridge_main->bridge_table); index++)
    {
        char str_section[MR_STRING_MAX] = {};
        snprintf(str_section, sizeof(str_section), "bridge:%d", index);

        /* Get devices */
        char str_devices[MR_STRING_MAX] = {};
        if (MESA_load_profile_string_nodef(sc->local_cfgfile, str_section, "devices", str_devices,
                                           sizeof(str_devices)) < 0)
            continue;

        /* Split devices */
        char * str_tokens[MR_TOKENS_MAX] = {};
        int nr_str_tokens = rte_strsplit(str_devices, sizeof(str_devices), str_tokens, MR_TOKENS_MAX, ',');
        if (nr_str_tokens <= 0)
        {
            MR_ERROR("Bridge config :  %s ,devices '%s' Is invalid .", str_section, str_devices);
            return RT_ERR;
        }

        struct bridge * bridge = &bridge_main->bridge_table[nr_bridge];
        bridge->master_port = UINT32_MAX;
        for (int i = 0; i < nr_str_tokens; i++)
        {
            /* Check the dev name */
            struct mr_dev_desc * dev_desc = mr_dev_desc_lookup(sc->devmgr_main, str_tokens[i]);
            if (dev_desc == NULL)
            {
                MR_ERROR("Bridge config :  %s ,dev name '%s' Is invalid .", str_section, str_tokens[i]);
                return RT_ERR;
            }

            /* Save the device info */
            dev_desc->bridge_index = nr_bridge;
            bridge->nr_devices++;
            bridge->ports[i] = dev_desc->port_id;
            rte_ether_addr_copy(&dev_desc->eth_addr, &bridge->ether_addr[i]);
        }

        nr_bridge++;

        /* Parser master devices */
        char str_master_dev[MR_STRING_MAX] = {};
        if (MESA_load_profile_string_nodef(sc->local_cfgfile, str_section, "master_device", str_master_dev,
                                           sizeof(str_master_dev)) < 0)
            continue;

        /* Check the dev name */
        struct mr_dev_desc * dev_desc = mr_dev_desc_lookup(sc->devmgr_main, str_master_dev);
        if (dev_desc == NULL)
        {
            MR_ERROR("Bridge config :  %s ,master dev name '%s' Is invalid .", str_section, str_master_dev);
            return RT_ERR;
        }

        bridge->master_port = dev_desc->port_id;
    }

    bridge_main->nr_bridge = nr_bridge;
    return RT_SUCCESS;
}

/* Dump bridge config */
void dump_bridge_config(struct node_bridge_main * bridge_main)
{
    if (bridge_main->nr_bridge == 0)
        return;

    MR_INFO(" ");
    MR_INFO("Bridge config total num: %u", bridge_main->nr_bridge);

    for (int index = 0; index <= bridge_main->nr_bridge; index++)
    {
        struct bridge * bridge_item = &bridge_main->bridge_table[index];
        char str_section[MR_STRING_MAX];

        int len = snprintf(str_section, sizeof(str_section), " Port IDs: [");

        for (int i = 0; i < bridge_item->nr_devices; i++)
        {
            len += snprintf(str_section + len, sizeof(str_section), "%u", bridge_item->ports[i]);

            if (i < bridge_item->nr_devices - 1)
            {
                len += snprintf(str_section + len, sizeof(str_section), ", ");
            }
        }

        len += snprintf(str_section + len, sizeof(str_section), "]");

        MR_INFO("Bridge index : %d, %s", index, str_section);
    }
}

/* Bridge init */
int bridge_init(struct sc_main * sc)
{
    /* Get indirect pool */
    struct node_bridge_main * bridge_main = node_bridge_main_get();
    bridge_main->direct_pool = mrb_direct_mempool_locate(sc->mrb_pool_main, "single-d", 0, 0);

    if (bridge_main->direct_pool == NULL)
    {
        MR_ERROR("Bridge node, get 'single-d' failed .");
        return RT_ERR;
    }

    /* Init the bridge index */
    unsigned int dev_desc_iterator = 0;
    struct mr_dev_desc * dev_desc = NULL;

    while ((dev_desc = mr_dev_desc_iterate(sc_main_get()->devmgr_main, &dev_desc_iterator)) != NULL)
    {
        dev_desc->bridge_index = UINT8_MAX;
    }

    /* Parser the local bridge config */
    int ret = parser_bridge_conf(sc, bridge_main);

    /* Dump the config */
    if (ret != RT_ERR)
        dump_bridge_config(bridge_main);

    return ret;
}

/* Bridge look up */
int bridge_lookup(struct bridge * bridge, struct rte_ether_hdr * ether_hdr, uint8_t * last_index, port_id_t * out_port)
{
    if (rte_is_same_ether_addr(&bridge->ether_addr[*last_index], &ether_hdr->dst_addr))
    {
        *out_port = bridge->ports[*last_index];
        return RT_SUCCESS;
    }

    for (int index = 0; index < bridge->nr_devices; index++)
    {
        if (index == *last_index)
            continue;

        if (rte_is_same_ether_addr(&bridge->ether_addr[index], &ether_hdr->dst_addr))
        {
            *out_port = bridge->ports[index];
            *last_index = index;
            return RT_SUCCESS;
        }
    }
    return RT_ERR;
}

/* Generate and store the trace information for original pkt */
static __rte_always_inline void gen_store_trace_info_original(struct rte_node * node, struct rte_mbuf * mbuf,
                                                              uint16_t next_node_index, struct bridge_stats * stats,
                                                              enum bridge_match_result match_result)
{
    /* Populate the next node infomation */
    char str_record[MR_STRING_MAX];
    assert(match_result < BRIDGE_MATCH_MAX);
    int len = snprintf(str_record, sizeof(str_record), ", rsn:%s", bridge_match_result_str[match_result]);

    /* Populate the egress port information */
    if (next_node_index != BRIDGE_NEXT_DROP)
    {
        struct mrb_metadata * mrb_meta = mrbuf_cz_data(mbuf, MR_NODE_CTRLZONE_ID);
        len += snprintf(str_record + len, sizeof(str_record) - len, ", tx:%u", mrb_meta->port_egress);
    }

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

/* Generate and store the trace information for clone pkt */
#if 0
static __rte_always_inline void gen_store_trace_info_copy(struct rte_node * node, struct rte_mbuf * mbuf,
                                                          uint16_t next_node_index, uint8_t measurement_type)
{

    /* Populate the next node infomation */
    char str_record[MR_STRING_MAX];
    struct mrb_metadata * mrb_meta = mrbuf_cz_data(mbuf, MR_NODE_CTRLZONE_ID);
    snprintf(str_record, sizeof(str_record), "cloned packet, rsn: broadcast");

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

static __rte_always_inline uint16_t bridge_node_process(struct rte_graph * graph, struct rte_node * node, void ** objs,
                                                        uint16_t cnt)
{
    /* Get pkts num and pkts buffer */
    unsigned int last_spec = 0;
    unsigned int n_left_from = cnt;
    struct rte_mbuf ** pkts = (struct rte_mbuf **)objs;
    void ** batch_pkts = objs;
    uint16_t batch_next_node_index = BRIDGE_NEXT_ETH_EGRESS;

    rte_graph_t graph_id = graph->id;
    struct bridge_stats stats = {};
    enum bridge_match_result match_result = BRIDGE_MATCH_MAX;
    struct sc_main * sc = sc_main_get();
    struct node_bridge_main * bridge_main = node_bridge_main_get();

    while (n_left_from > 0)
    {
        uint16_t next_node_index = BRIDGE_NEXT_ETH_EGRESS;

        struct rte_mbuf * mbuf = pkts[0];
        assert(mbuf != NULL);

        struct mrb_metadata * mrb_meta = mrbuf_cz_data(mbuf, MR_NODE_CTRLZONE_ID);
        assert(mrb_meta != NULL);

        port_id_t port_ingress = mrb_meta->port_ingress;
        struct mr_dev_desc * dev_desc = mr_dev_desc_lookup_by_port_id(sc->devmgr_main, port_ingress);
        assert(dev_desc != NULL);

        uint8_t bridge_index = dev_desc->bridge_index;
        assert((bridge_index < RTE_DIM(bridge_main->bridge_table)) && (bridge_index != UINT8_MAX));
        //__rte_mbuf_sanity_check(mbuf, 1);

        pkts += 1;
        n_left_from -= 1;

        /* Bridge lookup */
        port_id_t egress_port;
        struct bridge * bridge = &bridge_main->bridge_table[bridge_index];
        struct rte_ether_hdr * ether_hdr = rte_pktmbuf_mtod(mbuf, struct rte_ether_hdr *);

        int ret = bridge_lookup(bridge, ether_hdr, &last_bridge_index[graph_id][bridge_index], &egress_port);
        if (likely(ret == RT_SUCCESS))
        {
            match_result = BRIDGE_MATCH_DST_MAC;
            mrb_meta->port_egress = egress_port;
            goto node_enqueue;
        }

        /* Check broadcast pkt */
        if (likely(!rte_is_broadcast_ether_addr(&ether_hdr->dst_addr)))
        {
            /* Send to master pkts */
            egress_port = bridge->master_port;
            if (likely((egress_port != port_ingress) && (egress_port != UINT32_MAX)))
            {

                match_result = BRIDGE_MATCH_TO_MASTER;
                stats.match_result[match_result]++;
                mrb_meta->port_egress = egress_port;
                goto node_enqueue;
            }

            /* Drop pkts */
            match_result = BRIDGE_MATCH_MISS;
            stats.match_result[match_result]++;
            next_node_index = BRIDGE_NEXT_DROP;
            goto node_enqueue;
        }

        /* Deal broadcast pkt */
        match_result = BRIDGE_MATCH_BROADCAST;
        stats.match_result[match_result]++;
        uint8_t original_pkt = 1;
        for (int index = 0; index < bridge->nr_devices; index++)
        {
            egress_port = bridge->ports[index];

            if (egress_port == port_ingress)
                continue;

            if (original_pkt)
            {
                mrb_meta->port_egress = egress_port;
                next_node_index = BRIDGE_NEXT_ETH_EGRESS;
                original_pkt = 0;
                continue;
            }

            /* copy pkt */
            struct rte_mbuf * mbuf_copy = rte_pktmbuf_copy(mbuf, bridge_main->direct_pool, 0, UINT32_MAX);
            if (unlikely(mbuf_copy == NULL))
            {
                stats.pkt_clone_failed++;
                continue;
            }

            /* Populate the egress port information */
            struct mrb_metadata * mbuf_copy_mrb_meta = mrbuf_cz_data(mbuf_copy, MR_NODE_CTRLZONE_ID);
            mrb_metadata_clear(mbuf_copy_mrb_meta);
            mbuf_copy_mrb_meta->port_egress = egress_port;
            mbuf_copy_mrb_meta->traffic_link_id = mrb_meta->traffic_link_id;

#if 0
            /* Mark packets for trace*/
            dp_trace_filter_exec(sc_main_get()->trace, mbuf_copy, 0, rte_lcore_id());

            /* Check if tracing is enabled for the current Mbuf */
            if (unlikely(dp_trace_record_can_emit(mbuf_copy, DP_TRACE_MEASUREMENT_TYPE_TELEMETRY)))
            {
                gen_store_trace_info_copy(node, mbuf_copy, BRIDGE_NEXT_ETH_EGRESS, DP_TRACE_MEASUREMENT_TYPE_TELEMETRY);
            }

            if (unlikely(dp_trace_record_can_emit(mbuf_copy, DP_TRACE_MEASUREMENT_TYPE_TRACE)))
            {
                gen_store_trace_info_copy(node, mbuf_copy, BRIDGE_NEXT_ETH_EGRESS, DP_TRACE_MEASUREMENT_TYPE_TRACE);
            }
#endif

            /* Send clone pkt */
            rte_node_enqueue(graph, node, BRIDGE_NEXT_ETH_EGRESS, (void **)&mbuf_copy, 1);
        }

    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_original(node, mbuf, next_node_index, &stats, match_result);
        }

        /* Batch processing */
        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++;
        }
    }

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

    /* Update graph stats */
    struct bridge_stats * graph_stats = &stats_per_graph[graph_id];
    graph_stats->total_pkts += cnt;
    graph_stats->pkts_per_batch = cnt;
    graph_stats->pkt_clone_failed += stats.pkt_clone_failed;

    for (int i = 0; i < RTE_DIM(stats.match_result); i++)
    {
        graph_stats->match_result[i] += stats.match_result[i];
    }

    return 0;
}

/* Bridge Node Base */
static struct rte_node_register bridge_node_base = {
    .process = bridge_node_process,
    .name = "bridge",
    .init = NULL,
    .nb_edges = BRIDGE_NEXT_MAX,
    .next_nodes =
        {
            [BRIDGE_NEXT_DROP] = "pkt_drop_trap",
            [BRIDGE_NEXT_ETH_EGRESS] = "eth_egress",
        },
};
RTE_NODE_REGISTER(bridge_node_base);

/* Bridge Stat */
cJSON * bridge_node_monit_loop(struct sc_main * sc)
{
    unsigned int nr_graphs = sc->nr_io_thread;
    cJSON * json_root = cJSON_CreateObject();

    uint64_t total_pkts[nr_graphs];
    uint64_t pkts_per_batch[nr_graphs];
    uint64_t total_pkt_clone_failed[nr_graphs];
    uint64_t total_match_result[nr_graphs][BRIDGE_MATCH_MAX];

    for (uint32_t graph_id = 0; graph_id < nr_graphs; graph_id++)
    {
        struct bridge_stats * stats = &stats_per_graph[graph_id];
        if (stats->total_pkts == 0)
        {
            total_pkts[graph_id] = 0;
            pkts_per_batch[graph_id] = 0;
            total_pkt_clone_failed[graph_id] = 0;

            memset(total_match_result[graph_id], 0, sizeof(total_match_result[graph_id]));
            continue;
        }

        total_pkts[graph_id] = stats->total_pkts;
        pkts_per_batch[graph_id] = stats->pkts_per_batch;
        total_pkt_clone_failed[graph_id] = stats->pkt_clone_failed;

        for (int i = 0; i < BRIDGE_MATCH_MAX; i++)
        {
            total_match_result[graph_id][i] = stats->match_result[i];
        }
    }

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

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

    cJSON * json_clone_failed = create_uint64_array(total_pkt_clone_failed, nr_graphs);
    cJSON_AddItemToObject(json_root, "bridge, clone_failed", json_clone_failed);

    for (int i = 0; i < BRIDGE_MATCH_MAX; i++)
    {
        char str_title[MR_STRING_MAX];
        snprintf(str_title, sizeof(str_title), "bridge, %s", bridge_match_result_str[i]);

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

    return json_root;
}