#include <assert.h>

#include "utils_internal.h"
#include "packet_internal.h"
#include "packet_manager.h"
#include "packet_builder.h"
#include "fieldstat/fieldstat_easy.h"

#define PACKET_MANAGER_LOG_ERROR(format, ...) STELLAR_LOG_ERROR(__thread_local_logger, "packet manager", format, ##__VA_ARGS__)
#define PACKET_MANAGER_LOG_FATAL(format, ...) STELLAR_LOG_FATAL(__thread_local_logger, "packet manager", format, ##__VA_ARGS__)
#define PACKET_MANAGER_LOG_INFO(format, ...) STELLAR_LOG_INFO(__thread_local_logger, "packet manager", format, ##__VA_ARGS__)

struct packet_manager_rte
{
    enum packet_stage curr_stage;
    struct packet_queue queue[PACKET_QUEUE_MAX];

    void *claim_arg;
    on_packet_claimed_callback *claim_cb;

    struct mq_runtime *mq_rte;
    struct packet_manager_stat stat;
};

struct packet_manager_sche
{
    struct exdata_schema *ex_sche;
    struct mq_schema *mq_sche;
    int pkt_msg_id[PACKET_STAGE_MAX];
};

struct packet_manager
{
    uint16_t thread_num;
    struct packet_manager_sche *sche;
    struct packet_manager_rte *rte[MAX_THREAD_NUM];

    struct fieldstat_easy *fs;
    int pkt_mgr_fs_idx[PKT_MGR_STAT_MAX];
};

/******************************************************************************
 * utils
 ******************************************************************************/

const char *packet_stage_to_str(enum packet_stage stage)
{
    switch (stage)
    {
    case PACKET_STAGE_PREROUTING:
        return "PACKET_STAGE_PREROUTING";
    case PACKET_STAGE_INPUT:
        return "PACKET_STAGE_INPUT";
    case PACKET_STAGE_FORWARD:
        return "PACKET_STAGE_FORWARD";
    case PACKET_STAGE_OUTPUT:
        return "PACKET_STAGE_OUTPUT";
    case PACKET_STAGE_POSTROUTING:
        return "PACKET_STAGE_POSTROUTING";
    default:
        return "PACKET_STAGE_UNKNOWN";
    }
}

uint64_t packet_manager_stat_get(struct packet_manager_stat *stat, enum pkt_mgr_stat_type type)
{
    switch (type)
    {
#define XX(_type, _name, _val) case _type: return stat->_val;
    PKT_MGR_STAT_MAP(XX)
#undef XX
    default:
        return 0;
    }
}

/******************************************************************************
 * packet manager rte
 ******************************************************************************/

static void packet_manager_rte_free(struct packet_manager_rte *pkt_mgr_rte)
{
    struct packet *pkt = NULL;

    if (pkt_mgr_rte)
    {
        for (int i = 0; i < PACKET_QUEUE_MAX; i++)
        {
            while ((pkt = TAILQ_FIRST(&pkt_mgr_rte->queue[i])))
            {
                TAILQ_REMOVE(&pkt_mgr_rte->queue[i], pkt, stage_tqe);
                packet_free(pkt);
            }
        }

        free(pkt_mgr_rte);
        pkt_mgr_rte = NULL;
    }
}

static struct packet_manager_rte *packet_manager_rte_new(struct mq_runtime *mq_rte)
{
    struct packet_manager_rte *pkt_mgr_rte = calloc(1, sizeof(struct packet_manager_rte));
    if (pkt_mgr_rte == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to allocate memory for packet_manager_rte");
        return NULL;
    }

    for (int i = 0; i < PACKET_QUEUE_MAX; i++)
    {
        TAILQ_INIT(&pkt_mgr_rte->queue[i]);
    }
    pkt_mgr_rte->mq_rte = mq_rte;

    return pkt_mgr_rte;
}

/******************************************************************************
 * packet manager sche
 ******************************************************************************/

static void on_packet_stage_dispatch(int pkt_msg_id, void *msg, on_msg_cb_func *cb, void *cb_arg, void *dispatch_arg)
{
    assert(msg);
    assert(dispatch_arg);

    enum packet_stage stage = PACKET_STAGE_MAX;
    struct packet *pkt = (struct packet *)msg;
    struct packet_manager_sche *pkt_mgr_sche = (struct packet_manager_sche *)dispatch_arg;

    for (int i = 0; i < PACKET_STAGE_MAX; i++)
    {
        if (pkt_mgr_sche->pkt_msg_id[i] == pkt_msg_id)
        {
            stage = i;
            break;
        }
    }

    ((on_packet_stage_callback *)(void *)cb)(pkt, stage, cb_arg);
}

static void packet_manager_sche_free(struct packet_manager_sche *pkt_mgr_sche)
{
    if (pkt_mgr_sche)
    {
        if (pkt_mgr_sche->mq_sche)
        {
            for (int i = 0; i < PACKET_STAGE_MAX; i++)
            {
                if (pkt_mgr_sche->pkt_msg_id[i] >= 0)
                {
                    mq_schema_destroy_topic(pkt_mgr_sche->mq_sche, pkt_mgr_sche->pkt_msg_id[i]);
                }
            }
        }

        if (pkt_mgr_sche->ex_sche)
        {
            exdata_schema_free(pkt_mgr_sche->ex_sche);
        }

        free(pkt_mgr_sche);
        pkt_mgr_sche = NULL;
    }
}

static struct packet_manager_sche *packet_manager_sche_new(struct mq_schema *mq_sche)
{
    struct packet_manager_sche *pkt_mgr_sche = calloc(1, sizeof(struct packet_manager_sche));
    if (pkt_mgr_sche == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to allocate memory for packet_schema");
        return NULL;
    }

    pkt_mgr_sche->mq_sche = mq_sche;
    pkt_mgr_sche->ex_sche = exdata_schema_new();
    if (pkt_mgr_sche->ex_sche == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to create exdata_schema");
        goto error_out;
    }

    for (int i = 0; i < PACKET_STAGE_MAX; i++)
    {
        pkt_mgr_sche->pkt_msg_id[i] = mq_schema_create_topic(pkt_mgr_sche->mq_sche, packet_stage_to_str(i), &on_packet_stage_dispatch, pkt_mgr_sche, NULL, NULL);
        if (pkt_mgr_sche->pkt_msg_id[i] < 0)
        {
            PACKET_MANAGER_LOG_ERROR("failed to create topic %s", packet_stage_to_str(i));
            goto error_out;
        }
    }

    return pkt_mgr_sche;

error_out:
    packet_manager_sche_free(pkt_mgr_sche);
    return NULL;
}

/******************************************************************************
 *  packet manager
 ******************************************************************************/

struct packet_manager *packet_manager_new(struct mq_schema *mq_sche, uint16_t thread_num)
{
    struct packet_manager *pkt_mgr = calloc(1, sizeof(struct packet_manager));
    if (pkt_mgr == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to allocate memory for packet_manager");
        return NULL;
    }

    pkt_mgr->thread_num = thread_num;
    pkt_mgr->sche = packet_manager_sche_new(mq_sche);
    if (pkt_mgr->sche == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to create packet_schema");
        goto error_out;
    }

    pkt_mgr->fs = fieldstat_easy_new(pkt_mgr->thread_num, "packet_manager", NULL, 0);
    if (pkt_mgr->fs == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to create fieldstat_easy");
        goto error_out;
    }
    if (fieldstat_easy_enable_auto_output(pkt_mgr->fs, "metrics/packet_manager.json", 2) != 0)
    {
        PACKET_MANAGER_LOG_ERROR("failed to enable auto output for fieldstat_easy");
        goto error_out;
    }
    for (int i = 0; i < PKT_MGR_STAT_MAX; i++)
    {
        pkt_mgr->pkt_mgr_fs_idx[i] = fieldstat_easy_register_counter(pkt_mgr->fs, pkt_mgr_stat_str[i]);
    }

    return pkt_mgr;

error_out:
    packet_manager_free(pkt_mgr);
    return NULL;
}

void packet_manager_free(struct packet_manager *pkt_mgr)
{
    if (pkt_mgr)
    {
        if (pkt_mgr->fs)
        {
            fieldstat_easy_free(pkt_mgr->fs);
        }

        packet_manager_sche_free(pkt_mgr->sche);

        free(pkt_mgr);
        pkt_mgr = NULL;
    }
}

int packet_manager_new_packet_exdata_index(struct packet_manager *pkt_mgr, const char *name, exdata_free *func, void *arg)
{
    assert(pkt_mgr);
    return exdata_schema_new_index(pkt_mgr->sche->ex_sche, name, func, arg);
}

int packet_manager_subscribe(struct packet_manager *pkt_mgr, enum packet_stage stage, on_packet_stage_callback *cb, void *arg)
{
    assert(pkt_mgr);
    return mq_schema_subscribe(pkt_mgr->sche->mq_sche, pkt_mgr->sche->pkt_msg_id[stage], (on_msg_cb_func *)(void *)cb, arg);
}

int packet_manager_init(struct packet_manager *pkt_mgr, uint16_t thread_id, struct mq_runtime *mq_rte)
{
    assert(pkt_mgr);
    assert(thread_id < pkt_mgr->thread_num);
    assert(mq_rte);

    pkt_mgr->rte[thread_id] = packet_manager_rte_new(mq_rte);
    if (pkt_mgr->rte[thread_id] == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to create packet_manager_rte");
        return -1;
    }
    else
    {
        return 0;
    }
}

void packet_manager_clean(struct packet_manager *pkt_mgr, uint16_t thread_id)
{
    assert(pkt_mgr);
    assert(thread_id < pkt_mgr->thread_num);

    struct packet_manager_rte *pkt_mgr_rte = pkt_mgr->rte[thread_id];
    packet_manager_rte_free(pkt_mgr_rte);
    pkt_mgr_rte = NULL;
}

void packet_manager_ingress(struct packet_manager *pkt_mgr, uint16_t thread_id, struct packet *pkt)
{
    struct packet_manager_rte *pkt_mgr_rte = pkt_mgr->rte[thread_id];
    struct exdata_runtime *ex_rte = exdata_runtime_new(pkt_mgr->sche->ex_sche);
    packet_set_user_data(pkt, ex_rte);

    pkt_mgr_rte->stat.pkts_ingress++;
    pkt_mgr_rte->stat.queue[PACKET_STAGE_PREROUTING].pkts_in++;
    TAILQ_INSERT_TAIL(&pkt_mgr_rte->queue[PACKET_STAGE_PREROUTING], pkt, stage_tqe);
}

struct packet *packet_manager_egress(struct packet_manager *pkt_mgr, uint16_t thread_id)
{
    struct packet_manager_rte *pkt_mgr_rte = pkt_mgr->rte[thread_id];
    struct packet *pkt = TAILQ_FIRST(&pkt_mgr_rte->queue[PACKET_STAGE_MAX]);

    if (pkt)
    {
        pkt_mgr_rte->stat.pkts_egress++;
        pkt_mgr_rte->stat.queue[PACKET_STAGE_MAX].pkts_out++;
        TAILQ_REMOVE(&pkt_mgr_rte->queue[PACKET_STAGE_MAX], pkt, stage_tqe);

        struct exdata_runtime *ex_rte = packet_get_user_data(pkt);
        exdata_runtime_free(ex_rte);
        return pkt;
    }
    else
    {
        return NULL;
    }
}

void packet_manager_dispatch(struct packet_manager *pkt_mgr, uint16_t thread_id)
{
    struct packet_manager_rte *pkt_mgr_rte = pkt_mgr->rte[thread_id];

    for (int i = 0; i < PACKET_STAGE_MAX; i++)
    {
        pkt_mgr_rte->curr_stage = i;

        struct packet *pkt = NULL;
        while ((pkt = TAILQ_FIRST(&pkt_mgr_rte->queue[pkt_mgr_rte->curr_stage])))
        {
            packet_set_claim(pkt, false);
            pkt_mgr_rte->claim_cb = NULL;
            pkt_mgr_rte->claim_arg = NULL;

            TAILQ_REMOVE(&pkt_mgr_rte->queue[pkt_mgr_rte->curr_stage], pkt, stage_tqe);
            pkt_mgr_rte->stat.queue[pkt_mgr_rte->curr_stage].pkts_out++;

            mq_runtime_publish_message(pkt_mgr_rte->mq_rte, pkt_mgr_rte->curr_stage, pkt);
            mq_runtime_dispatch(pkt_mgr_rte->mq_rte);

            if (packet_is_claim(pkt))
            {
                if (pkt_mgr_rte->claim_cb)
                {
                    struct exdata_runtime *ex_rte = packet_get_user_data(pkt);
                    exdata_runtime_reset(ex_rte);

                    pkt_mgr_rte->claim_cb(pkt, pkt_mgr_rte->claim_arg);
                }
                continue;
            }

            TAILQ_INSERT_TAIL(&pkt_mgr_rte->queue[pkt_mgr_rte->curr_stage + 1], pkt, stage_tqe);
            pkt_mgr_rte->stat.queue[pkt_mgr_rte->curr_stage + 1].pkts_in++;
        }
    }
    pkt_mgr_rte->curr_stage = -1;
}

int packet_manager_claim_packet(struct packet_manager *pkt_mgr, uint16_t thread_id, struct packet *pkt, on_packet_claimed_callback cb, void *arg)
{
    assert(pkt_mgr);
    struct packet_manager_rte *pkt_mgr_rte = pkt_mgr->rte[thread_id];

    if (packet_is_claim(pkt))
    {
        PACKET_MANAGER_LOG_ERROR("packet is already claimed, cannot claim again");
        return -1;
    }
    else
    {
        pkt_mgr_rte->claim_cb = cb;
        pkt_mgr_rte->claim_arg = arg;
        packet_set_claim(pkt, true);
        pkt_mgr_rte->stat.queue[pkt_mgr_rte->curr_stage].pkts_claim++;
        return 0;
    }
}

void packet_manager_schedule_packet(struct packet_manager *pkt_mgr, uint16_t thread_id, struct packet *pkt, enum packet_stage stage)
{
    assert(pkt_mgr);
    struct packet_manager_rte *pkt_mgr_rte = pkt_mgr->rte[thread_id];

    if (stage >= PACKET_STAGE_MAX)
    {
        PACKET_MANAGER_LOG_ERROR("invalid stage %d", stage);
        assert(0);
        return;
    }

    pkt_mgr_rte->stat.queue[stage].pkts_schedule++;
    pkt_mgr_rte->stat.queue[stage].pkts_in++;
    TAILQ_INSERT_TAIL(&pkt_mgr_rte->queue[stage], pkt, stage_tqe);
}

struct packet_manager_stat *packet_manager_get_stat(struct packet_manager *pkt_mgr, uint16_t thread_id)
{
    struct packet_manager_rte *pkt_mgr_rte = pkt_mgr->rte[thread_id];
    return &pkt_mgr_rte->stat;
}

void packet_manager_print_stat(struct packet_manager *pkt_mgr, uint16_t thread_id)
{
    struct packet_manager_rte *pkt_mgr_rte = pkt_mgr->rte[thread_id];

    PACKET_MANAGER_LOG_INFO("runtime: %p, pkts_ingress: %lu, pkts_egress: %lu",
                            pkt_mgr_rte, pkt_mgr_rte->stat.pkts_ingress,
                            pkt_mgr_rte->stat.pkts_egress);
    for (int i = 0; i < PACKET_QUEUE_MAX; i++)
    {
        PACKET_MANAGER_LOG_INFO("runtime: %p, %-24s stat => pkts_in: %lu, pkts_out: %lu, pkts_claim: %lu, pkts_schedule: %lu",
                                pkt_mgr_rte,
                                packet_stage_to_str(i),
                                pkt_mgr_rte->stat.queue[i].pkts_in,
                                pkt_mgr_rte->stat.queue[i].pkts_out,
                                pkt_mgr_rte->stat.queue[i].pkts_claim,
                                pkt_mgr_rte->stat.queue[i].pkts_schedule);
    }
}

struct packet *packet_manager_build_tcp_packet(struct packet_manager *pkt_mgr, const struct packet *origin_pkt, uint32_t tcp_seq, uint32_t tcp_ack, uint8_t tcp_flags,
                                               const char *tcp_options, uint16_t tcp_options_len, const char *tcp_payload, uint16_t tcp_payload_len)
{
    struct packet *pkt = packet_build_tcp(origin_pkt, tcp_seq, tcp_ack, tcp_flags, tcp_options, tcp_options_len, tcp_payload, tcp_payload_len);
    if (pkt == NULL)
    {
        return NULL;
    }

    struct exdata_runtime *ex_rte = exdata_runtime_new(pkt_mgr->sche->ex_sche);
    packet_set_user_data(pkt, ex_rte);

    return pkt;
}

struct packet *packet_manager_build_udp_packet(struct packet_manager *pkt_mgr, const struct packet *origin_pkt, const char *udp_payload, uint16_t udp_payload_len)
{
    struct packet *pkt = packet_build_udp(origin_pkt, udp_payload, udp_payload_len);
    if (pkt == NULL)
    {
        return NULL;
    }

    struct exdata_runtime *ex_rte = exdata_runtime_new(pkt_mgr->sche->ex_sche);
    packet_set_user_data(pkt, ex_rte);

    return pkt;
}

struct packet *packet_manager_build_l3_packet(struct packet_manager *pkt_mgr, const struct packet *origin_pkt, uint8_t ip_proto, const char *l3_payload, uint16_t l3_payload_len)
{
    struct packet *pkt = packet_build_l3(origin_pkt, ip_proto, l3_payload, l3_payload_len);
    if (pkt == NULL)
    {
        return NULL;
    }

    struct exdata_runtime *ex_rte = exdata_runtime_new(pkt_mgr->sche->ex_sche);
    packet_set_user_data(pkt, ex_rte);

    return pkt;
}

struct packet *packet_manager_dup_packet(struct packet_manager *pkt_mgr, const struct packet *origin_pkt)
{
    struct packet *pkt = packet_dup(origin_pkt);
    if (pkt == NULL)
    {
        return NULL;
    }

    struct exdata_runtime *ex_rte = exdata_runtime_new(pkt_mgr->sche->ex_sche);
    packet_set_user_data(pkt, ex_rte);

    return pkt;
}

void packet_manager_free_packet(struct packet_manager *pkt_mgr __attribute__((unused)), struct packet *pkt)
{
    if (pkt)
    {
        struct exdata_runtime *ex_rte = packet_get_user_data(pkt);
        exdata_runtime_free(ex_rte);

        packet_free(pkt);
    }
}

/******************************************************************************
 * packet manager module
 ******************************************************************************/

static void on_polling(struct module_manager *mod_mgr, void *args)
{
    uint64_t now_ms = clock_get_real_time_ms();
    int thread_id = module_manager_get_thread_id(mod_mgr);
    struct packet_manager *pkt_mgr = (struct packet_manager *)args;

    static __thread uint64_t last_sync_stat_ms = 0;
    static __thread struct packet_manager_stat pkt_mgr_last_stat = {0};
    if (now_ms - last_sync_stat_ms >= SYNC_STAT_INTERVAL_MS)
    {
        struct packet_manager_stat *pkt_mgr_curr_stat = &pkt_mgr->rte[thread_id]->stat;
        for (int i = 0; i < PKT_MGR_STAT_MAX; i++)
        {
            uint64_t val = packet_manager_stat_get(pkt_mgr_curr_stat, i) - packet_manager_stat_get(&pkt_mgr_last_stat, i);
            fieldstat_easy_counter_incrby(pkt_mgr->fs, thread_id, pkt_mgr->pkt_mgr_fs_idx[i], NULL, 0, val);
        }
        pkt_mgr_last_stat = *pkt_mgr_curr_stat;
        last_sync_stat_ms = now_ms;
    }
}

struct packet_manager *module_to_packet_manager(struct module *mod)
{
    assert(mod);
    assert(strcmp(module_get_name(mod), PACKET_MANAGER_MODULE_NAME) == 0);
    return (struct packet_manager *)module_get_ctx(mod);
}

struct module *packet_manager_on_init(struct module_manager *mod_mgr)
{
    assert(mod_mgr);
    struct mq_schema *mq_sche = module_manager_get_mq_schema(mod_mgr);
    assert(mq_sche);
    uint16_t thread_num = module_manager_get_max_thread_num(mod_mgr);

    struct packet_manager *pkt_mgr = packet_manager_new(mq_sche, thread_num);
    if (pkt_mgr == NULL)
    {
        return NULL;
    }
    module_manager_polling_subscribe(mod_mgr, on_polling, pkt_mgr);

    struct module *pkt_mgr_mod = module_new(PACKET_MANAGER_MODULE_NAME, NULL);
    if (pkt_mgr_mod == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to create packet_manager");
        packet_manager_free(pkt_mgr);
        return NULL;
    }
    module_set_ctx(pkt_mgr_mod, pkt_mgr);

    PACKET_MANAGER_LOG_FATAL("packet_manager init");
    return pkt_mgr_mod;
}

void packet_manager_on_exit(struct module_manager *mod_mgr __attribute__((unused)), struct module *mod)
{
    if (mod)
    {
        struct packet_manager *pkt_mgr = module_get_ctx(mod);

        packet_manager_free(pkt_mgr);
        module_free(mod);
        PACKET_MANAGER_LOG_FATAL("packet_manager exit");
    }
}

struct module *packet_manager_on_thread_init(struct module_manager *mod_mgr, int thread_id, struct module *mod)
{
    struct packet_manager *pkt_mgr = module_get_ctx(mod);
    assert(pkt_mgr);
    struct mq_runtime *mq_rte = module_manager_get_mq_runtime(mod_mgr);
    assert(mq_rte);
    assert(thread_id < pkt_mgr->thread_num);

    if (packet_manager_init(pkt_mgr, thread_id, mq_rte) != 0)
    {
        PACKET_MANAGER_LOG_ERROR("failed to init packet_manager_init");
        return NULL;
    }
    else
    {
        return mod;
    }
}

void packet_manager_on_thread_exit(struct module_manager *mod_mgr __attribute__((unused)), int thread_id, struct module *mod)
{
    struct packet_manager *pkt_mgr = module_get_ctx(mod);
    if (pkt_mgr)
    {
        assert(thread_id < pkt_mgr->thread_num);
        packet_manager_clean(pkt_mgr, thread_id);
    }
}