#include <assert.h>

#include "utils.h"
#include "packet_internal.h"
#include "packet_manager_runtime.h"
#include "packet_manager_internal.h"

#define PACKET_MANAGER_MODULE_NAME "packet_manager_module"

struct packet_manager_config
{
    uint16_t nr_worker_thread;
};

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

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

    uint64_t val = 0;
    if (load_and_validate_toml_integer_config(toml_file, "packet_io.nr_worker_thread", &val, 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;
    }
    cfg->nr_worker_thread = val;

    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)
{
    assert(msg);
    assert(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 *)(void *)cb)(stage, pkt, cb_arg);
}

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

        if (pkt_mgr_schema->exdata)
        {
            exdata_schema_free(pkt_mgr_schema->exdata);
        }

        free(pkt_mgr_schema);
        pkt_mgr_schema = NULL;
    }
}

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

    pkt_mgr_schema->exdata = exdata_schema_new();
    if (pkt_mgr_schema->exdata == NULL)
    {
        PACKET_MANAGER_LOG_ERROR("failed to create exdata_schema");
        goto error_out;
    }

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

    return pkt_mgr_schema;

error_out:
    packet_manager_schema_free(pkt_mgr_schema);
    return NULL;
}

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

struct packet_manager *packet_manager_new(struct mq_schema *mq, 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(mq);
    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();
        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)
{
    struct packet_manager_runtime *pkt_mgr_rt = NULL;

    if (pkt_mgr)
    {
        if (pkt_mgr->cfg)
        {
            for (uint16_t i = 0; i < pkt_mgr->cfg->nr_worker_thread; i++)
            {
                pkt_mgr_rt = pkt_mgr->runtime[i];
                if (pkt_mgr_rt)
                {
                    PACKET_MANAGER_LOG_INFO("runtime: %p, idx: %d, will be cleaned", pkt_mgr_rt, i);
                    packet_manager_runtime_print_stat(pkt_mgr_rt);
                    packet_manager_runtime_free(pkt_mgr_rt);
                }
            }
        }

        packet_manager_schema_free(pkt_mgr->schema);
        packet_manager_config_free(pkt_mgr->cfg);

        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)
{
    return exdata_schema_new_index(pkt_mgr->schema->exdata, name, func, arg);
}

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

void packet_manager_init(struct packet_manager *pkt_mgr, uint16_t thread_id, struct mq_runtime *mq_rt)
{
    assert(pkt_mgr);
    assert(thread_id < pkt_mgr->cfg->nr_worker_thread);
    assert(mq_rt);
    struct packet_manager_runtime *runtime = pkt_mgr->runtime[thread_id];

    packet_manager_runtime_init(runtime, mq_rt);
}

void packet_manager_ingress(struct packet_manager *pkt_mgr, uint16_t thread_id, struct packet *pkt)
{
    struct packet_manager_runtime *runtime = pkt_mgr->runtime[thread_id];
    struct exdata_runtime *exdata_rt = exdata_runtime_new(pkt_mgr->schema->exdata);
    packet_set_user_data(pkt, exdata_rt);

    packet_manager_runtime_ingress(runtime, pkt);
}

struct packet *packet_manager_egress(struct packet_manager *pkt_mgr, uint16_t thread_id)
{
    struct packet_manager_runtime *runtime = pkt_mgr->runtime[thread_id];
    struct packet *pkt = packet_manager_runtime_egress(runtime);
    if (pkt)
    {
        struct exdata_runtime *exdata_rt = packet_get_user_data(pkt);
        exdata_runtime_free(exdata_rt);
    }
    return pkt;
}

void packet_manager_dispatch(struct packet_manager *pkt_mgr, uint16_t thread_id)
{
    struct packet_manager_runtime *runtime = pkt_mgr->runtime[thread_id];
    packet_manager_runtime_dispatch(runtime);
}

int packet_manager_claim_packet(struct packet_manager *pkt_mgr, uint16_t thread_id, struct packet *pkt, on_packet_claimed_callback cb, void *args)
{
    struct packet_manager_runtime *runtime = pkt_mgr->runtime[thread_id];
    return packet_manager_runtime_claim_packet(runtime, pkt, cb, args);
}

void packet_manager_schedule_packet(struct packet_manager *pkt_mgr, uint16_t thread_id, struct packet *pkt, enum packet_stage stage)
{
    struct packet_manager_runtime *runtime = pkt_mgr->runtime[thread_id];
    packet_manager_runtime_schedule_packet(runtime, pkt, stage);
}

struct packet_manager_stat *packet_manager_get_stat(struct packet_manager *pkt_mgr, uint16_t thread_id)
{
    struct packet_manager_runtime *runtime = pkt_mgr->runtime[thread_id];
    return packet_manager_runtime_get_stat(runtime);
}

void packet_manager_print_stat(struct packet_manager *pkt_mgr, uint16_t thread_id)
{
    struct packet_manager_runtime *runtime = pkt_mgr->runtime[thread_id];
    packet_manager_runtime_print_stat(runtime);
}

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

struct packet_manager *stellar_module_get_packet_manager(struct stellar_module_manager *mod_mgr)
{
    assert(mod_mgr);
    struct stellar_module *pkt_mgr_mod = stellar_module_manager_get_module(mod_mgr, PACKET_MANAGER_MODULE_NAME);
    if (pkt_mgr_mod == NULL)
    {
        return NULL;
    }
    return (struct packet_manager*)stellar_module_get_ctx(pkt_mgr_mod);
}

struct stellar_module *packet_manager_on_init(struct stellar_module_manager *mod_mgr)
{
    assert(mod_mgr);
    struct mq_schema *mq_schema = stellar_module_manager_get_mq_schema(mod_mgr);
    assert(mq_schema);
    const char *toml_file = stellar_module_manager_get_toml_path(mod_mgr);
    assert(toml_file);

    struct packet_manager *pkt_mgr = packet_manager_new(mq_schema, toml_file);
    if (pkt_mgr == NULL)
    {
        return NULL;
    }

    struct stellar_module *pkt_mgr_mod = stellar_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;
    }
    stellar_module_set_ctx(pkt_mgr_mod, pkt_mgr);

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

void packet_manager_on_exit(struct stellar_module_manager *mod_mgr __attribute__((unused)), struct stellar_module *mod)
{
    if (mod)
    {
        struct packet_manager *pkt_mgr = stellar_module_get_ctx(mod);

        packet_manager_free(pkt_mgr);
        stellar_module_free(mod);
        PACKET_MANAGER_LOG_FATAL("packet_manager exited");
    }
}