#include "utils.h"
#include "stellar/mq.h"
#include "packet_private.h"
#include "packet_manager_private.h"

#define PACKET_MANAGER_LOG_ERROR(format, ...) STELLAR_LOG_ERROR(__thread_local_logger, "packet manager", format, ##__VA_ARGS__)
#define PACKET_MANAGER_LOG_INFO(format, ...) STELLAR_LOG_WARN(__thread_local_logger, "packet manager", format, ##__VA_ARGS__)

TAILQ_HEAD(packet_queue, packet);

struct packet_manager_config
{
    uint16_t nr_worker_thread;
};

struct packet_manager_stat
{
    uint64_t input_pkts;
    uint64_t output_pkts;

    uint64_t take_pkts;
    uint64_t schedule_pkts;

    uint64_t curr_queue_len[PACKET_STAGE_MAX];
};

struct packet_manager_schema
{
    struct mq_schema *mq;
    int topic_id[PACKET_STAGE_MAX];
};

struct packet_manager_runtime
{
    struct mq_runtime *mq;
    struct packet_queue queue[PACKET_STAGE_MAX];
    struct packet_manager_stat stat;
};

struct packet_manager
{
    struct packet_manager_config *cfg;
    struct packet_manager_schema *schema;
    struct packet_manager_runtime *runtime[MAX_THREAD_NUM];
};

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

static void packet_manager_config_free(struct packet_manager_config *cfg)
{
    if (cfg)
    {
        free(cfg);
        cfg = NULL;
    }
}

static struct packet_manager_config *packet_manager_config_new(const char *toml_file)
{
    struct packet_manager_config *cfg = calloc(1, sizeof(struct packet_manager_config));
    if (cfg == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to allocate memory for packet_manager_config");
        return NULL;
    }

    if (load_and_validate_toml_integer_config(toml_file, "packet_io.nr_worker_thread", (uint64_t *)&cfg->nr_worker_thread, 1, MAX_THREAD_NUM) != 0)
    {
        PACKET_MANAGER_LOG_ERROR("failed to load packet_io.nr_worker_thread from %s", toml_file);
        free(cfg);
        return NULL;
    }

    return cfg;
}

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

static void on_packet_stage_dispatch(int topic_id, void *msg, on_msg_cb_func *cb, void *cb_arg, void *dispatch_arg)
{
    struct packet_manager_schema *schema = (struct packet_manager_schema *)dispatch_arg;
    struct packet *pkt = (struct packet *)msg;

    enum packet_stage stage = PACKET_STAGE_MAX;
    for (int i = 0; i < PACKET_STAGE_MAX; i++)
    {
        if (schema->topic_id[i] == topic_id)
        {
            stage = i;
            break;
        }
    }

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

static void packet_manager_schema_free(struct packet_manager_schema *schema)
{
    if (schema)
    {
        if (schema->mq)
        {
            for (int i = 0; i < PACKET_STAGE_MAX; i++)
            {
                if (schema->topic_id[i] != -1)
                {
                    mq_schema_destroy_topic(schema->mq, schema->topic_id[i]);
                }
            }

            mq_schema_free(schema->mq);
        }

        free(schema);
        schema = NULL;
    }
}

static struct packet_manager_schema *packet_manager_schema_new()
{
    struct packet_manager_schema *schema = calloc(1, sizeof(struct packet_manager_schema));
    if (schema == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to allocate memory for packet_manager_schema");
        return NULL;
    }

    schema->mq = mq_schema_new();
    if (schema->mq == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to create mq_schema");
        goto error_out;
    }

    for (int i = 0; i < PACKET_STAGE_MAX; i++)
    {
        schema->topic_id[i] = -1;
    }
    schema->topic_id[PACKET_STAGE_PREROUTING] = mq_schema_create_topic(schema->mq, "packet_stage_prerouting", on_packet_stage_dispatch, schema, NULL, NULL);
    schema->topic_id[PACKET_STAGE_INPUT] = mq_schema_create_topic(schema->mq, "packet_stage_input", on_packet_stage_dispatch, schema, NULL, NULL);
    schema->topic_id[PACKET_STAGE_FORWARD] = mq_schema_create_topic(schema->mq, "packet_stage_forward", on_packet_stage_dispatch, schema, NULL, NULL);
    schema->topic_id[PACKET_STAGE_OUTPUT] = mq_schema_create_topic(schema->mq, "packet_stage_output", on_packet_stage_dispatch, schema, NULL, NULL);
    schema->topic_id[PACKET_STAGE_POSTROUTING] = mq_schema_create_topic(schema->mq, "packet_stage_postrouting", on_packet_stage_dispatch, schema, NULL, NULL);
    for (int i = 0; i < PACKET_STAGE_MAX; i++)
    {
        if (schema->topic_id[i] < 0)
        {
            PACKET_MANAGER_LOG_ERROR("failed to create topic");
            goto error_out;
        }
    }

    return schema;

error_out:
    packet_manager_schema_free(schema);
    return NULL;
}

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

static void packet_manager_runtime_free(struct packet_manager_runtime *runtime)
{
    if (runtime)
    {
        if (runtime->mq)
        {
            mq_runtime_free(runtime->mq);
        }
        for (int i = 0; i < PACKET_STAGE_MAX; i++)
        {
            struct packet *pkt = NULL;
            while ((pkt = TAILQ_FIRST(&runtime->queue[i])))
            {
                TAILQ_REMOVE(&runtime->queue[i], pkt, stage_tqe);

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

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

    runtime->mq = mq_runtime_new(schema->mq);
    if (runtime->mq == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to create mq_runtime");
        free(runtime);
        return NULL;
    }

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

    return runtime;
}

static void packet_manager_runtime_stat(struct packet_manager_runtime *runtime)
{
    PACKET_MANAGER_LOG_INFO("input_pkts: %lu, output_pkts: %lu, take_pkts: %lu, schedule_pkts: %lu",
                            runtime->stat.input_pkts, runtime->stat.output_pkts, runtime->stat.take_pkts, runtime->stat.schedule_pkts);
    for (int i = 0; i < PACKET_STAGE_MAX; i++)
    {
        PACKET_MANAGER_LOG_INFO("curr_queue_len[%d]: %lu", i, runtime->stat.curr_queue_len[i]);
    }
}

/******************************************************************************
 * Public API
 ******************************************************************************/

struct packet_manager *packet_manager_new(const char *toml_file)
{
    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->cfg = packet_manager_config_new(toml_file);
    if (pkt_mgr->cfg == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to create packet_manager_config");
        goto error_out;
    }

    pkt_mgr->schema = packet_manager_schema_new();
    if (pkt_mgr->schema == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to create packet_manager_schema");
        goto error_out;
    }

    for (uint16_t i = 0; i < pkt_mgr->cfg->nr_worker_thread; i++)
    {
        pkt_mgr->runtime[i] = packet_manager_runtime_new(pkt_mgr->schema);
        if (pkt_mgr->runtime[i] == NULL)
        {
            PACKET_MANAGER_LOG_ERROR("failed to create packet_manager_runtime");
            goto error_out;
        }
    }

    return pkt_mgr;

error_out:
    packet_manager_free(pkt_mgr);
    return NULL;
}

void packet_manager_free(struct packet_manager *pkt_mgr)
{
    if (pkt_mgr)
    {
        packet_manager_config_free(pkt_mgr->cfg);
        packet_manager_schema_free(pkt_mgr->schema);

        for (uint16_t i = 0; i < pkt_mgr->cfg->nr_worker_thread; i++)
        {
            if (pkt_mgr->runtime[i])
            {
                PACKET_MANAGER_LOG_INFO("worker thread %d packet manager runtime stat", i);
                packet_manager_runtime_stat(pkt_mgr->runtime[i]);
                packet_manager_runtime_free(pkt_mgr->runtime[i]);
            }
        }

        free(pkt_mgr);
        pkt_mgr = NULL;
    }
}

void packet_manager_runtime_input(struct packet_manager_runtime *pkt_mgr_rt, struct packet *pkt)
{
    pkt_mgr_rt->stat.input_pkts++;
    pkt_mgr_rt->stat.curr_queue_len[PACKET_STAGE_PREROUTING]++;
    TAILQ_INSERT_TAIL(&pkt_mgr_rt->queue[PACKET_STAGE_PREROUTING], pkt, stage_tqe);
}

struct packet *packet_manager_runtime_output(struct packet_manager_runtime *pkt_mgr_rt)
{
    struct packet *pkt = TAILQ_FIRST(&pkt_mgr_rt->queue[PACKET_STAGE_POSTROUTING]);
    if (pkt)
    {
        pkt_mgr_rt->stat.output_pkts++;
        pkt_mgr_rt->stat.curr_queue_len[PACKET_STAGE_POSTROUTING]--;
        TAILQ_REMOVE(&pkt_mgr_rt->queue[PACKET_STAGE_POSTROUTING], 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++)
    {
        struct packet *pkt = NULL;
        while ((pkt = TAILQ_FIRST(&pkt_mgr_rt->queue[i])))
        {
            pkt_mgr_rt->stat.curr_queue_len[i]--;
            TAILQ_REMOVE(&pkt_mgr_rt->queue[i], pkt, stage_tqe);
            mq_runtime_publish_message(pkt_mgr_rt->mq, i, pkt);
            mq_runtime_dispatch(pkt_mgr_rt->mq);

            if (packet_is_stolen(pkt))
            {
                continue;
            }

            if (i + 1 == PACKET_STAGE_MAX)
            {
                break;
            }
            else
            {
                pkt_mgr_rt->stat.curr_queue_len[i + 1]++;
                TAILQ_INSERT_TAIL(&pkt_mgr_rt->queue[i + 1], pkt, stage_tqe);
            }
        }
    }
}

struct packet_manager_schema *packet_manager_get_schema(struct packet_manager *pkt_mgr)
{
    if (pkt_mgr)
    {
        return pkt_mgr->schema;
    }
    else
    {
        return NULL;
    }
}

struct packet_manager_runtime *packet_manager_get_runtime(struct packet_manager *pkt_mgr, uint16_t thr_idx)
{
    if (pkt_mgr && thr_idx < pkt_mgr->cfg->nr_worker_thread)
    {
        return pkt_mgr->runtime[thr_idx];
    }
    else
    {
        return NULL;
    }
}

int packet_manager_schema_add_subscriber(struct packet_manager_schema *schema, enum packet_stage stage, on_packet_stage_callback cb, void *args)
{
    return mq_schema_subscribe(schema->mq, schema->topic_id[stage], (on_msg_cb_func *)cb, args);
}

void packet_manager_runtime_take_packet(struct packet_manager_runtime *pkt_mgr_rt, struct packet *pkt)
{
    pkt_mgr_rt->stat.take_pkts++;
    packet_set_stolen(pkt, true);
}

void packet_manager_runtime_schedule_packet(struct packet_manager_runtime *pkt_mgr_rt, struct packet *pkt, enum packet_stage stage)
{
    pkt_mgr_rt->stat.schedule_pkts++;
    packet_set_stolen(pkt, false);
    TAILQ_INSERT_TAIL(&pkt_mgr_rt->queue[stage], pkt, stage_tqe);
}