#include <assert.h>
#include <stdlib.h>

#include "packet_internal.h"
#include "packet_manager_runtime.h"

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";
    }
}

struct packet_manager_runtime *packet_manager_runtime_new()
{
    struct packet_manager_runtime *pkt_mgr_rt = calloc(1, sizeof(struct packet_manager_runtime));
    if (pkt_mgr_rt == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to allocate memory for packet_manager_runtime");
        return NULL;
    }

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

    return pkt_mgr_rt;
}

void packet_manager_runtime_free(struct packet_manager_runtime *pkt_mgr_rt)
{
    if (pkt_mgr_rt)
    {
        for (int i = 0; i < PACKET_QUEUE_MAX; i++)
        {
            struct packet *pkt = NULL;
            while ((pkt = TAILQ_FIRST(&pkt_mgr_rt->queue[i])))
            {
                TAILQ_REMOVE(&pkt_mgr_rt->queue[i], pkt, stage_tqe);

                // TODO: free packet and free mbuff
                packet_free(pkt);
            }
        }
    }
    free(pkt_mgr_rt);
    pkt_mgr_rt = NULL;
}

void packet_manager_runtime_init(struct packet_manager_runtime *pkt_mgr_rt, struct mq_runtime *mq_rt)
{
    pkt_mgr_rt->mq = mq_rt;
}

void packet_manager_runtime_ingress(struct packet_manager_runtime *pkt_mgr_rt, struct packet *pkt)
{
    pkt_mgr_rt->stat.total.pkts_ingress++;
    pkt_mgr_rt->stat.queue[PACKET_STAGE_PREROUTING].pkts_in++;
    TAILQ_INSERT_TAIL(&pkt_mgr_rt->queue[PACKET_STAGE_PREROUTING], pkt, stage_tqe);
}

struct packet *packet_manager_runtime_egress(struct packet_manager_runtime *pkt_mgr_rt)
{
    struct packet *pkt = TAILQ_FIRST(&pkt_mgr_rt->queue[PACKET_STAGE_MAX]);
    if (pkt)
    {
        pkt_mgr_rt->stat.total.pkts_egress++;
        pkt_mgr_rt->stat.queue[PACKET_STAGE_MAX].pkts_out++;
        TAILQ_REMOVE(&pkt_mgr_rt->queue[PACKET_STAGE_MAX], pkt, stage_tqe);
    }
    return pkt;
}

void packet_manager_runtime_dispatch(struct packet_manager_runtime *pkt_mgr_rt)
{
    for (int i = 0; i < PACKET_STAGE_MAX; i++)
    {
        pkt_mgr_rt->stage = i;

        struct packet *pkt = NULL;
        while ((pkt = TAILQ_FIRST(&pkt_mgr_rt->queue[pkt_mgr_rt->stage])))
        {
            packet_set_claim(pkt, false);
            pkt_mgr_rt->claimed_cb = NULL;
            pkt_mgr_rt->cb_args = NULL;

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

            mq_runtime_publish_message(pkt_mgr_rt->mq, pkt_mgr_rt->stage, pkt);
            mq_runtime_dispatch(pkt_mgr_rt->mq);

            if (packet_is_claim(pkt))
            {
                if (pkt_mgr_rt->claimed_cb)
                {
                    pkt_mgr_rt->claimed_cb(pkt, pkt_mgr_rt->cb_args);
                }
                packet_set_claim(pkt, false);
                continue;
            }

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

int packet_manager_runtime_claim_packet(struct packet_manager_runtime *pkt_mgr_rt, struct packet *pkt, on_packet_claimed_callback cb, void *args)
{
    if (packet_is_claim(pkt))
    {
        PACKET_MANAGER_LOG_ERROR("packet is already claimed, cannot claim again");
        return -1;
    }
    else
    {
        pkt_mgr_rt->claimed_cb = cb;
        pkt_mgr_rt->cb_args = args;
        packet_set_claim(pkt, true);
        pkt_mgr_rt->stat.queue[pkt_mgr_rt->stage].pkts_claim++;
        return 0;
    }
}

void packet_manager_runtime_schedule_packet(struct packet_manager_runtime *pkt_mgr_rt, struct packet *pkt, enum packet_stage stage)
{
    if (stage >= PACKET_STAGE_MAX)
    {
        PACKET_MANAGER_LOG_ERROR("invalid stage %d", stage);
        assert(0);
        return;
    }

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

struct packet_manager_stat *packet_manager_runtime_get_stat(struct packet_manager_runtime *pkt_mgr_rt)
{
    return &pkt_mgr_rt->stat;
}

void packet_manager_runtime_print_stat(struct packet_manager_runtime *pkt_mgr_rt)
{
    PACKET_MANAGER_LOG_INFO("runtime: %p, pkts_ingress: %lu, pkts_egress: %lu",
                            pkt_mgr_rt, pkt_mgr_rt->stat.total.pkts_ingress, pkt_mgr_rt->stat.total.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_rt,
                                packet_stage_to_str(i),
                                pkt_mgr_rt->stat.queue[i].pkts_in,
                                pkt_mgr_rt->stat.queue[i].pkts_out,
                                pkt_mgr_rt->stat.queue[i].pkts_claim,
                                pkt_mgr_rt->stat.queue[i].pkts_schedule);
    }
}