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

#include "stellar/exdata.h"
#include "stellar/packet_manager.h"
#include "stellar/session_manager.h"
#include "stellar/module_manager.h"

#include "utils_internal.h"
#include "session_internal.h"
#include "session_manager_runtime.h"

#define SESSION_MANAGER_MODULE_NAME "session_manager_module"

#pragma GCC diagnostic ignored "-Wunused-parameter"
#pragma GCC diagnostic ignored "-Wunused-function"

struct session_manager_schema
{
    struct exdata_schema *exdata;
    struct mq_schema *mq;

    int pkt_exdata_idx;

    int topic_id_tcp;
    int topic_id_udp;
    int topic_id_ctrl_pkt;
    int topic_id_tcp_stream;
};

struct session_manager
{
    uint16_t thread_num;
    struct session_manager_config *cfg;
    struct session_manager_schema *schema;
    struct session_manager_runtime *runtime[MAX_THREAD_NUM];
    struct stellar_module_manager *mod_mgr;
};

/******************************************************************************
 * callback
 ******************************************************************************/

static void on_session_dispatch(int topic_id, void *msg, on_msg_cb_func *cb, void *cb_args, void *dispatch_args)
{
    struct session *sess = (struct session *)msg;
    struct packet *pkt = (struct packet *)session_get0_current_packet(sess);

    ((on_session_callback *)(void *)cb)(sess, pkt, cb_args);
}

static void on_tcp_stream_dispatch(int topic_id, void *msg, on_msg_cb_func *cb, void *cb_args, void *dispatch_args)
{
    struct tcp_segment *seg = (struct tcp_segment *)msg;

    ((on_tcp_stream_callback *)(void *)cb)(seg->user_data, seg->data, seg->len, cb_args);
}

static void on_tcp_stream_free(void *msg, void *args)
{
    struct tcp_segment *seg = (struct tcp_segment *)msg;
    struct session *sess = (struct session *)seg->user_data;

    session_free_tcp_segment(sess, seg);
}

static void on_packet_forward(enum packet_stage stage, struct packet *pkt, void *args)
{
    struct session_manager *sess_mgr = (struct session_manager *)args;
    struct stellar_module_manager *mod_mgr = sess_mgr->mod_mgr;
    int thread_id = stellar_module_manager_get_thread_id(mod_mgr);
    struct mq_runtime *mq_rt = stellar_module_manager_get_mq_runtime(mod_mgr);
    struct session_manager_runtime *sess_mgr_rt = sess_mgr->runtime[thread_id];

    /*
     * We use the system's real time instead of monotonic time for the following reasons:
     *     -> Session creation/closure times require real time (e.g., for logging session activities).
     *     -> Session ID generation relies on real time        (e.g., for reverse calculating session creation time from the session ID).
     *
     * Note: Modifying the system time will affect the timing wheel, impacting session expiration, and TCP reassembly expiration.
     * Suggestion: After modifying the system time, restart the service to ensure consistent timing.
     */
    uint64_t now_ms = clock_get_real_time_ms();

    struct tcp_segment *seg = NULL;
    struct session *sess = session_manager_runtime_lookup_session_by_packet(sess_mgr_rt, pkt);
    if (sess == NULL)
    {
        sess = session_manager_runtime_new_session(sess_mgr_rt, pkt, now_ms);
        if (sess == NULL)
        {
            goto fast_path;
        }
        else
        {
            session_set_user_data(sess, exdata_runtime_new(sess_mgr->schema->exdata));
            goto slow_path;
        }
    }
    else
    {
        if (session_manager_runtime_update_session(sess_mgr_rt, sess, pkt, now_ms) == -1)
        {
            goto fast_path;
        }
        else
        {
            goto slow_path;
        }
    }

slow_path:
    if (session_get_type(sess) == SESSION_TYPE_TCP)
    {
        mq_runtime_publish_message(mq_rt, sess_mgr->schema->topic_id_tcp, sess);
        while ((seg = session_get_tcp_segment(sess)))
        {
            mq_runtime_publish_message(mq_rt, sess_mgr->schema->topic_id_tcp_stream, seg);
        }
    }
    else
    {
        mq_runtime_publish_message(mq_rt, sess_mgr->schema->topic_id_udp, sess);
    }

    packet_set_exdata(pkt, sess_mgr->schema->pkt_exdata_idx, sess);

fast_path:
    packet_set_exdata(pkt, sess_mgr->schema->pkt_exdata_idx, NULL);
}

static void on_packet_output(enum packet_stage stage, struct packet *pkt, void *args)
{
    struct session_manager *sess_mgr = (struct session_manager *)args;
    struct stellar_module_manager *mod_mgr = sess_mgr->mod_mgr;
    int thread_id = stellar_module_manager_get_thread_id(mod_mgr);
    struct session_manager_runtime *sess_mgr_rt = sess_mgr->runtime[thread_id];

    struct session *sess = (struct session *)packet_get_exdata(pkt, sess_mgr->schema->pkt_exdata_idx);
    if (sess)
    {
        enum flow_type type = session_get_flow_type(sess);
        int is_ctrl = packet_is_ctrl(pkt);
        uint16_t len = packet_get_raw_len(pkt);
        switch (packet_get_action(pkt))
        {
        case PACKET_ACTION_DROP:
            session_inc_stat(sess, type, (is_ctrl ? STAT_CONTROL_PACKETS_DROPPED : STAT_RAW_PACKETS_DROPPED), 1);
            session_inc_stat(sess, type, (is_ctrl ? STAT_CONTROL_BYTES_DROPPED : STAT_RAW_BYTES_DROPPED), len);
            break;
        case PACKET_ACTION_FORWARD:
            session_inc_stat(sess, type, (is_ctrl ? STAT_CONTROL_PACKETS_TRANSMITTED : STAT_RAW_PACKETS_TRANSMITTED), 1);
            session_inc_stat(sess, type, (is_ctrl ? STAT_CONTROL_BYTES_TRANSMITTED : STAT_RAW_BYTES_TRANSMITTED), len);
            break;
        default:
            assert(0);
            break;
        }

        session_set_current_packet(sess, NULL);
        session_set_flow_type(sess, FLOW_TYPE_NONE);
    }

    if (packet_get_origin(pkt) == NULL)
    {
        session_manager_runtime_record_duplicated_packet(sess_mgr_rt, pkt);
    }
}

static void clean_session(struct session_manager_runtime *sess_mgr_rt, uint64_t now_ms)
{
#define MAX_CLEANED_SESS 1024
    char buffer[4096] = {0};
    struct session *sess = NULL;
    struct session *cleaned_sess[MAX_CLEANED_SESS] = {NULL};

    uint64_t used = session_manager_runtime_clean_session(sess_mgr_rt, now_ms, cleaned_sess, MAX_CLEANED_SESS);
    for (uint64_t j = 0; j < used; j++)
    {
        sess = cleaned_sess[j];

        session_to_str(sess, 0, buffer, sizeof(buffer));
        SESSION_MANAGER_LOG_INFO("session free: %s", buffer);

        // TODO publish session free msg
        // TODO mq_runtime_dispatch_immediate()

        struct exdata_runtime *exdata_rt = (struct exdata_runtime *)session_get_user_data(sess);
        exdata_runtime_free(exdata_rt);

        session_manager_runtime_free_session(sess_mgr_rt, sess);
    }
}

static void on_polling(struct stellar_module_manager *mod_mgr, void *args)
{
    struct session_manager *sess_mgr = (struct session_manager *)args;
    int thread_id = stellar_module_manager_get_thread_id(mod_mgr);
    struct session_manager_runtime *sess_mgr_rt = sess_mgr->runtime[thread_id];
    uint64_t now_ms = clock_get_real_time_ms();

    clean_session(sess_mgr_rt, now_ms);

    // TODO
    // ouput stat to fs4
    // session_manager_runtime_print_stat(sess_mgr_rt);
}

/******************************************************************************
 * session manager schema
 ******************************************************************************/

void session_manager_schema_free(struct session_manager_schema *sess_mgr_schema)
{
    if (sess_mgr_schema)
    {
        if (sess_mgr_schema->mq)
        {
            mq_schema_destroy_topic(sess_mgr_schema->mq, sess_mgr_schema->topic_id_tcp);
            mq_schema_destroy_topic(sess_mgr_schema->mq, sess_mgr_schema->topic_id_udp);
            mq_schema_destroy_topic(sess_mgr_schema->mq, sess_mgr_schema->topic_id_ctrl_pkt);
            mq_schema_destroy_topic(sess_mgr_schema->mq, sess_mgr_schema->topic_id_tcp_stream);
        }
        exdata_schema_free(sess_mgr_schema->exdata);

        free(sess_mgr_schema);
        sess_mgr_schema = NULL;
    }
}

struct session_manager_schema *session_manager_schema_new(struct packet_manager *pkt_mgr, struct mq_schema *mq, void *subscribe_args)
{
    if (packet_manager_subscribe(pkt_mgr, PACKET_STAGE_FORWARD, on_packet_forward, subscribe_args))
    {
        SESSION_MANAGER_LOG_ERROR("failed to subscribe PACKET_STAGE_FORWARD");
        return NULL;
    }
    if (packet_manager_subscribe(pkt_mgr, PACKET_STAGE_OUTPUT, on_packet_output, subscribe_args))
    {
        SESSION_MANAGER_LOG_ERROR("failed to subscribe PACKET_STAGE_OUTPUT");
        return NULL;
    }

    struct session_manager_schema *sess_mgr_schema = calloc(1, sizeof(struct session_manager_schema));
    if (sess_mgr_schema == NULL)
    {
        SESSION_MANAGER_LOG_ERROR("failed to allocate memory for session_manager_schema");
        return NULL;
    }

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

    sess_mgr_schema->mq = mq;
    sess_mgr_schema->pkt_exdata_idx = packet_manager_new_packet_exdata_index(pkt_mgr, "session_manager", NULL, NULL);
    if (sess_mgr_schema->pkt_exdata_idx == -1)
    {
        SESSION_MANAGER_LOG_ERROR("failed to create packet exdata index");
        goto error_out;
    }

    sess_mgr_schema->topic_id_tcp = mq_schema_create_topic(sess_mgr_schema->mq, "TCP", &on_session_dispatch, NULL, NULL, NULL);
    if (sess_mgr_schema->topic_id_tcp == -1)
    {
        SESSION_MANAGER_LOG_ERROR("failed to create topic TCP");
        goto error_out;
    }
    sess_mgr_schema->topic_id_udp = mq_schema_create_topic(sess_mgr_schema->mq, "UDP", &on_session_dispatch, NULL, NULL, NULL);
    if (sess_mgr_schema->topic_id_udp == -1)
    {
        SESSION_MANAGER_LOG_ERROR("failed to create topic UDP");
        goto error_out;
    }
    sess_mgr_schema->topic_id_ctrl_pkt = mq_schema_create_topic(sess_mgr_schema->mq, "CTRL_PKT", &on_session_dispatch, NULL, NULL, NULL);
    if (sess_mgr_schema->topic_id_ctrl_pkt == -1)
    {
        SESSION_MANAGER_LOG_ERROR("failed to create topic CTRL_PKT");
        goto error_out;
    }
    sess_mgr_schema->topic_id_tcp_stream = mq_schema_create_topic(sess_mgr_schema->mq, "TCP_STREAM", &on_tcp_stream_dispatch, NULL, &on_tcp_stream_free, NULL);
    if (sess_mgr_schema->topic_id_tcp_stream == -1)
    {
        SESSION_MANAGER_LOG_ERROR("failed to create topic TCP_STREAM");
        goto error_out;
    }

    return sess_mgr_schema;

error_out:
    session_manager_schema_free(sess_mgr_schema);
    return NULL;
}

/******************************************************************************
 * session manager
 ******************************************************************************/

void session_manager_free(struct session_manager *sess_mgr)
{
    struct session_manager_stat *stat = NULL;
    struct session_manager_runtime *sess_mgr_rt = NULL;

    if (sess_mgr)
    {
        for (int i = 0; i < sess_mgr->thread_num; i++)
        {
            sess_mgr_rt = sess_mgr->runtime[i];
            if (sess_mgr_rt == NULL)
            {
                continue;
            }

            stat = session_manager_runtime_get_stat(sess_mgr_rt);
            while (stat->tcp_sess_used || stat->udp_sess_used)
            {
                clean_session(sess_mgr_rt, UINT64_MAX);
            }

            SESSION_MANAGER_LOG_INFO("runtime: %p, idx: %d, will be cleaned", sess_mgr_rt, i);
            session_manager_runtime_print_stat(sess_mgr_rt);
            session_manager_runtime_free(sess_mgr->runtime[i]);
        }

        session_manager_schema_free(sess_mgr->schema);
        session_manager_config_free(sess_mgr->cfg);
        free(sess_mgr);
    }
}

struct session_manager *session_manager_new(struct stellar_module_manager *mod_mgr, struct packet_manager *pkt_mgr, struct mq_schema *mq_schema, const char *toml_file)
{

    uint64_t thread_num;
    uint64_t instance_id;
    uint64_t now_ms = clock_get_real_time_ms();

    if (load_toml_integer_config(toml_file, "instance.id", (uint64_t *)&instance_id, 0, 4095))
    {
        return NULL;
    }
    if (load_toml_integer_config(toml_file, "packet_io.thread_num", (uint64_t *)&thread_num, 0, MAX_THREAD_NUM))
    {
        return NULL;
    }

    struct session_manager *sess_mgr = calloc(1, sizeof(struct session_manager));
    if (sess_mgr == NULL)
    {
        SESSION_MANAGER_LOG_ERROR("failed to allocate memory for session_manager");
        return NULL;
    }

    sess_mgr->cfg = session_manager_config_new(toml_file);
    if (sess_mgr->cfg == NULL)
    {
        SESSION_MANAGER_LOG_ERROR("failed to create session_manager_config");
        goto error_out;
    }
    session_manager_config_print(sess_mgr->cfg);

    sess_mgr->schema = session_manager_schema_new(pkt_mgr, mq_schema, sess_mgr);
    if (sess_mgr->schema == NULL)
    {
        goto error_out;
    }

    sess_mgr->thread_num = (uint16_t)thread_num;
    for (int i = 0; i < sess_mgr->thread_num; i++)
    {
        sess_mgr->cfg->session_id_seed = instance_id << 8 | i;
        sess_mgr->runtime[i] = session_manager_runtime_new(sess_mgr->cfg, now_ms);
        if (sess_mgr->runtime[i] == NULL)
        {
            SESSION_MANAGER_LOG_ERROR("failed to create session_manager_runtime");
            goto error_out;
        }
    }

    stellar_module_manager_polling_subscribe(mod_mgr, on_polling, sess_mgr);

    return sess_mgr;

error_out:
    session_manager_free(sess_mgr);
    return NULL;
}

int session_manager_new_session_exdata_index(struct session_manager *sess_mgr, const char *name, exdata_free *func, void *arg)
{
    assert(sess_mgr);
    assert(name);

    return exdata_schema_new_index(sess_mgr->schema->exdata, name, func, arg);
}

int session_manager_subscribe_tcp(struct session_manager *sess_mgr, on_session_callback *cb, void *args)
{
    assert(sess_mgr);
    assert(cb);

    return mq_schema_subscribe(sess_mgr->schema->mq, sess_mgr->schema->topic_id_tcp, (on_msg_cb_func *)(void *)cb, args);
}

int session_manager_subscribe_udp(struct session_manager *sess_mgr, on_session_callback *cb, void *args)
{
    assert(sess_mgr);
    assert(cb);

    return mq_schema_subscribe(sess_mgr->schema->mq, sess_mgr->schema->topic_id_udp, (on_msg_cb_func *)(void *)cb, args);
}

int session_manager_subscribe_control_packet(struct session_manager *sess_mgr, on_session_callback *cb, void *args)
{
    assert(sess_mgr);
    assert(cb);

    return mq_schema_subscribe(sess_mgr->schema->mq, sess_mgr->schema->topic_id_ctrl_pkt, (on_msg_cb_func *)(void *)cb, args);
}

int session_manager_subscribe_tcp_stream(struct session_manager *sess_mgr, on_tcp_stream_callback *cb, void *args)
{
    assert(sess_mgr);
    assert(cb);

    return mq_schema_subscribe(sess_mgr->schema->mq, sess_mgr->schema->topic_id_tcp_stream, (on_msg_cb_func *)(void *)cb, args);
}

/******************************************************************************
 * session manager module
 ******************************************************************************/

struct session_manager *stellar_module_get_session_manager(struct stellar_module_manager *mod_mgr)
{
    assert(mod_mgr);
    struct stellar_module *sess_mgr_mod = stellar_module_manager_get_module(mod_mgr, SESSION_MANAGER_MODULE_NAME);
    if (sess_mgr_mod == NULL)
    {
        return NULL;
    }
    return stellar_module_get_ctx(sess_mgr_mod);
}

struct stellar_module *session_manager_on_init(struct stellar_module_manager *mod_mgr)
{
    assert(mod_mgr);
    struct packet_manager *pkt_mgr = stellar_module_get_packet_manager(mod_mgr);
    assert(pkt_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 session_manager *sess_mgr = session_manager_new(mod_mgr, pkt_mgr, mq_schema, toml_file);
    if (sess_mgr == NULL)
    {
        return NULL;
    }

    struct stellar_module *sess_mgr_mod = stellar_module_new(SESSION_MANAGER_MODULE_NAME, NULL);
    if (sess_mgr_mod == NULL)
    {
        SESSION_MANAGER_LOG_ERROR("failed to create session_manager");
        session_manager_free(sess_mgr);
        return NULL;
    }
    stellar_module_set_ctx(sess_mgr_mod, sess_mgr);

    SESSION_MANAGER_LOG_FATAL("session_manager initialized");
    return sess_mgr_mod;
}

void session_manager_on_exit(struct stellar_module_manager *mod_mgr __attribute__((unused)), struct stellar_module *mod)
{
    if (mod)
    {
        struct session_manager *sess_mgr = stellar_module_get_ctx(mod);

        session_manager_free(sess_mgr);
        stellar_module_free(mod);
        SESSION_MANAGER_LOG_FATAL("session_manager exited");
    }
}