#include <assert.h>

#include "utils_internal.h"
#include "session_internal.h"
#include "session_manager_log.h"
#include "session_manager_cfg.h"
#include "session_manager_rte.h"
#include "session_manager_stat.h"
#include "fieldstat/fieldstat_easy.h"

#define CLEAN_SESSION_BURST 1024
#define SESSION_MANAGER_MODULE_NAME "session_manager_module"

#pragma GCC diagnostic ignored "-Wunused-parameter"

struct session_manager_sche
{
     int pkt_ex_id;
     int sess_msg_id_tcp;
     int sess_msg_id_udp;
     int sess_msg_id_ctrl;
     int sess_msg_id_stream;

     struct mq_schema *mq_sche;
     struct exdata_schema *ex_sche;
};

struct session_manager
{
     int stat_idx[SESS_MGR_STAT_MAX];
     struct fieldstat_easy *fs;
     struct session_manager_cfg *cfg;
     struct session_manager_sche *sche;
     struct session_manager_rte *rte[MAX_THREAD_NUM];
     struct mq_runtime *mq[MAX_THREAD_NUM];
     struct stellar_module_manager *mod_mgr;
};

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

static void clean_closed_session(struct session_manager *sess_mgr, uint16_t thread_id, uint64_t now_ms)
{
     struct session *sess = NULL;
     struct session *cleaned[CLEAN_SESSION_BURST] = {NULL};
     struct mq_runtime *mq_rte = sess_mgr->mq[thread_id];
     struct session_manager_rte *sess_mgr_rte = sess_mgr->rte[thread_id];

     uint64_t used = session_manager_rte_clean_session(sess_mgr_rte, now_ms, cleaned, CLEAN_SESSION_BURST);
     for (uint64_t i = 0; i < used; i++)
     {
          sess = cleaned[i];
          assert(session_get_current_state(sess) == SESSION_STATE_CLOSED);
          session_set_current_packet(sess, NULL);
          session_set_flow_type(sess, FLOW_TYPE_NONE);

          if (session_get_type(sess) == SESSION_TYPE_TCP)
          {
               mq_runtime_publish_message(mq_rte, sess_mgr->sche->sess_msg_id_stream, &sess->empty_seg);
               mq_runtime_publish_message(mq_rte, sess_mgr->sche->sess_msg_id_ctrl, sess);
               mq_runtime_publish_message(mq_rte, sess_mgr->sche->sess_msg_id_tcp, sess);
          }
          else
          {
               mq_runtime_publish_message(mq_rte, sess_mgr->sche->sess_msg_id_ctrl, sess);
               mq_runtime_publish_message(mq_rte, sess_mgr->sche->sess_msg_id_udp, sess);
          }
     }
}

static void on_sess_msg_dispatch(int sess_msg_id, void *msg, on_msg_cb_func *msg_cb, void *msg_cb_args, void *dispatch_args)
{
     struct session *sess = (struct session *)msg;
     struct packet *pkt = (struct packet *)session_get0_current_packet(sess);
     enum session_state state = session_get_current_state(sess);

     ((on_session_message_callback *)(void *)msg_cb)(sess, state, pkt, msg_cb_args);
}

static void on_tcp_payload_msg_dispatch(int sess_msg_id, void *msg, on_msg_cb_func *msg_cb, void *msg_cb_args, void *dispatch_args)
{
     struct tcp_segment *seg = (struct tcp_segment *)msg;
     struct session *sess = (struct session *)seg->user_data;
     enum session_state state = session_get_current_state(sess);

     ((on_tcp_payload_callback *)(void *)msg_cb)(sess, state, seg->data, seg->len, msg_cb_args);
}

static void on_sess_msg_free(void *msg, void *args)
{
     struct session *sess = (struct session *)msg;

     if (session_get_current_state(sess) == SESSION_STATE_CLOSED)
     {
          struct session_manager *sess_mgr = (struct session_manager *)args;
          int thread_id = stellar_module_manager_get_thread_id(sess_mgr->mod_mgr);
          struct session_manager_rte *sess_mgr_rte = sess_mgr->rte[thread_id];

          char buffer[4096] = {0};
          session_to_str(sess, 0, buffer, sizeof(buffer));
          SESSION_MANAGER_LOG_INFO("session free: %s", buffer);

          struct exdata_runtime *exdata_rt = (struct exdata_runtime *)session_get_user_data(sess);
          exdata_runtime_free(exdata_rt);
          session_manager_rte_free_session(sess_mgr_rte, sess);
     }
}

static void on_tcp_payload_msg_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(struct packet *pkt, enum packet_stage stage, void *args)
{
     struct session_manager *sess_mgr = (struct session_manager *)args;
     int thread_id = stellar_module_manager_get_thread_id(sess_mgr->mod_mgr);
     struct mq_runtime *mq_rte = sess_mgr->mq[thread_id];
     struct session_manager_rte *sess_mgr_rte = sess_mgr->rte[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 tuple6 key;
     struct tcp_segment *seg = NULL;
     struct session *sess = session_manager_rte_lookup_session_by_packet(sess_mgr_rte, pkt);
     if (sess == NULL)
     {
          if (packet_is_ctrl(pkt))
          {
               goto fast_path;
          }

          sess = session_manager_rte_new_session(sess_mgr_rte, pkt, now_ms);
          if (sess == NULL)
          {
               goto fast_path;
          }
          else
          {
               session_set_user_data(sess, exdata_runtime_new(sess_mgr->sche->ex_sche));
               goto slow_path;
          }
     }
     else
     {
          if (packet_is_ctrl(pkt))
          {
               goto ctrl_path;
          }

          if (session_manager_rte_update_session(sess_mgr_rte, sess, pkt, now_ms) == -1)
          {
               goto fast_path;
          }
          else
          {
               goto slow_path;
          }
     }

ctrl_path:
     session_set_current_packet(sess, pkt);
     packet_get_innermost_tuple6(pkt, &key);
     if (tuple6_cmp(session_get_tuple6(sess), &key) == 0)
     {
          session_set_flow_type(sess, FLOW_TYPE_C2S);
     }
     else
     {
          session_set_flow_type(sess, FLOW_TYPE_S2C);
     }
     packet_set_exdata(pkt, sess_mgr->sche->pkt_ex_id, sess);
     mq_runtime_publish_message(mq_rte, sess_mgr->sche->sess_msg_id_ctrl, sess);
     return;

slow_path:
     if (session_get_type(sess) == SESSION_TYPE_TCP)
     {
          mq_runtime_publish_message(mq_rte, sess_mgr->sche->sess_msg_id_tcp, sess);
          while ((seg = session_get_tcp_segment(sess)))
          {
               mq_runtime_publish_message(mq_rte, sess_mgr->sche->sess_msg_id_stream, seg);
          }
     }
     else
     {
          mq_runtime_publish_message(mq_rte, sess_mgr->sche->sess_msg_id_udp, sess);
     }
     packet_set_exdata(pkt, sess_mgr->sche->pkt_ex_id, sess);
     return;

fast_path:
     packet_set_exdata(pkt, sess_mgr->sche->pkt_ex_id, NULL);
     return;
}

static void on_packet_output(struct packet *pkt, enum packet_stage stage, void *args)
{
     struct session_manager *sess_mgr = (struct session_manager *)args;
     int thread_id = stellar_module_manager_get_thread_id(sess_mgr->mod_mgr);
     struct session_manager_rte *sess_mgr_rte = sess_mgr->rte[thread_id];

     struct session *sess = (struct session *)packet_get_exdata(pkt, sess_mgr->sche->pkt_ex_id);
     if (sess)
     {
          struct tuple6 key;
          enum flow_type flow = FLOW_TYPE_NONE;
          packet_get_innermost_tuple6(pkt, &key);
          if (tuple6_cmp(session_get_tuple6(sess), &key) == 0)
          {
               flow = FLOW_TYPE_C2S;
          }
          else
          {
               flow = FLOW_TYPE_S2C;
          }

          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, flow, (is_ctrl ? STAT_CONTROL_PACKETS_DROPPED : STAT_RAW_PACKETS_DROPPED), 1);
               session_inc_stat(sess, flow, (is_ctrl ? STAT_CONTROL_BYTES_DROPPED : STAT_RAW_BYTES_DROPPED), len);
               break;
          case PACKET_ACTION_FORWARD:
               session_inc_stat(sess, flow, (is_ctrl ? STAT_CONTROL_PACKETS_TRANSMITTED : STAT_RAW_PACKETS_TRANSMITTED), 1);
               session_inc_stat(sess, flow, (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_rte_record_duplicated_packet(sess_mgr_rte, pkt);
     }
}

static void on_polling(struct stellar_module_manager *mod_mgr, void *args)
{
     uint64_t now_ms = clock_get_real_time_ms();
     struct session_manager *sess_mgr = (struct session_manager *)args;
     int thread_id = stellar_module_manager_get_thread_id(mod_mgr);
     struct session_manager_rte *sess_mgr_rte = sess_mgr->rte[thread_id];
     struct session_manager_stat *sess_mgr_stat = session_manager_rte_get_stat(sess_mgr_rte);

     clean_closed_session(sess_mgr, thread_id, now_ms);

     static __thread uint64_t last_sync_stat_ms = 0;
     static __thread struct session_manager_stat last_stat = {0};
     if (now_ms - last_sync_stat_ms >= SYNC_STAT_INTERVAL_MS)
     {
          for (int i = 0; i < SESS_MGR_STAT_MAX; i++)
          {
               uint64_t val = session_manager_stat_get(sess_mgr_stat, i) - session_manager_stat_get(&last_stat, i);
               fieldstat_easy_counter_incrby(sess_mgr->fs, thread_id, sess_mgr->stat_idx[i], NULL, 0, val);
          }
          last_stat = *sess_mgr_stat;
          last_sync_stat_ms = now_ms;
     }
}

static void session_manager_sche_free(struct session_manager_sche *sess_mgr_schema)
{
     if (sess_mgr_schema)
     {
          if (sess_mgr_schema->mq_sche)
          {
               mq_schema_destroy_topic(sess_mgr_schema->mq_sche, sess_mgr_schema->sess_msg_id_tcp);
               mq_schema_destroy_topic(sess_mgr_schema->mq_sche, sess_mgr_schema->sess_msg_id_udp);
               mq_schema_destroy_topic(sess_mgr_schema->mq_sche, sess_mgr_schema->sess_msg_id_ctrl);
               mq_schema_destroy_topic(sess_mgr_schema->mq_sche, sess_mgr_schema->sess_msg_id_stream);
          }
          exdata_schema_free(sess_mgr_schema->ex_sche);

          free(sess_mgr_schema);
          sess_mgr_schema = NULL;
     }
}

static struct session_manager_sche *session_manager_sche_new(struct packet_manager *pkt_mgr, struct mq_schema *mq_sche, void *sess_mgr)
{
     if (packet_manager_subscribe(pkt_mgr, PACKET_STAGE_FORWARD, on_packet_forward, sess_mgr))
     {
          SESSION_MANAGER_LOG_ERROR("failed to subscribe PACKET_STAGE_FORWARD");
          return NULL;
     }
     if (packet_manager_subscribe(pkt_mgr, PACKET_STAGE_OUTPUT, on_packet_output, sess_mgr))
     {
          SESSION_MANAGER_LOG_ERROR("failed to subscribe PACKET_STAGE_OUTPUT");
          return NULL;
     }

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

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

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

     /*
      * Publish session closed messages to multiple topics.
      * Each topic has its own session message free callback.
      * To prevent the same session from being freeed multiple times,
      * only TCP/UDP topics register session message free callbacks,
      * and other topics do not register session message callbacks;
      *
      * Restriction: MQ ensures that the session message free order is consistent with the publishing order
      */
     sess_mgr_schema->sess_msg_id_tcp = mq_schema_create_topic(sess_mgr_schema->mq_sche, "SESSION_MESSAGE_TCP", &on_sess_msg_dispatch, NULL, &on_sess_msg_free, sess_mgr);
     if (sess_mgr_schema->sess_msg_id_tcp == -1)
     {
          SESSION_MANAGER_LOG_ERROR("failed to create topic SESSION_MESSAGE_FREE");
          goto error_out;
     }
     sess_mgr_schema->sess_msg_id_udp = mq_schema_create_topic(sess_mgr_schema->mq_sche, "SESSION_MESSAGE_UDP", &on_sess_msg_dispatch, NULL, &on_sess_msg_free, sess_mgr);
     if (sess_mgr_schema->sess_msg_id_udp == -1)
     {
          SESSION_MANAGER_LOG_ERROR("failed to create topic SESSION_MESSAGE_UDP");
          goto error_out;
     }
     sess_mgr_schema->sess_msg_id_ctrl = mq_schema_create_topic(sess_mgr_schema->mq_sche, "SESSION_MESSAGE_CTRL_PKT", &on_sess_msg_dispatch, NULL, NULL, NULL);
     if (sess_mgr_schema->sess_msg_id_ctrl == -1)
     {
          SESSION_MANAGER_LOG_ERROR("failed to create topic SESSION_MESSAGE_CTRL_PKT");
          goto error_out;
     }
     sess_mgr_schema->sess_msg_id_stream = mq_schema_create_topic(sess_mgr_schema->mq_sche, "SESSION_MESSAGE_TCP_STREAM", &on_tcp_payload_msg_dispatch, NULL, &on_tcp_payload_msg_free, sess_mgr);
     if (sess_mgr_schema->sess_msg_id_stream == -1)
     {
          SESSION_MANAGER_LOG_ERROR("failed to create topic SESSION_MESSAGE_TCP_STREAM");
          goto error_out;
     }

     return sess_mgr_schema;

error_out:
     session_manager_sche_free(sess_mgr_schema);
     return NULL;
}

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

void session_manager_free(struct session_manager *sess_mgr)
{
     if (sess_mgr)
     {
          if (sess_mgr->sche)
          {
               session_manager_sche_free(sess_mgr->sche);
          }
          if (sess_mgr->fs)
          {
               fieldstat_easy_free(sess_mgr->fs);
          }
          if (sess_mgr->cfg)
          {
               session_manager_cfg_free(sess_mgr->cfg);
          }
          free(sess_mgr);
     }
}

static struct session_manager *session_manager_new(struct packet_manager *pkt_mgr, struct mq_schema *mq_schema, const char *toml_file)
{
     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_cfg_new(toml_file);
     if (sess_mgr->cfg == NULL)
     {
          SESSION_MANAGER_LOG_ERROR("failed to create session_manager_cfg");
          goto error_out;
     }
     session_manager_cfg_print(sess_mgr->cfg);

     sess_mgr->fs = fieldstat_easy_new(sess_mgr->cfg->thread_num, "session_manager", NULL, 0);
     if (sess_mgr->fs == NULL)
     {
          SESSION_MANAGER_LOG_ERROR("failed to create fieldstat_easy");
          goto error_out;
     }
     for (int i = 0; i < SESS_MGR_STAT_MAX; i++)
     {
          sess_mgr->stat_idx[i] = fieldstat_easy_register_counter(sess_mgr->fs, sess_mgr_stat_str[i]);
     }
     if (fieldstat_easy_enable_auto_output(sess_mgr->fs, "session_manager.fs4", 2) != 0)
     {
          SESSION_MANAGER_LOG_ERROR("failed to enable auto output");
          goto error_out;
     }

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

     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);
     assert(func);
     return exdata_schema_new_index(sess_mgr->sche->ex_sche, name, func, arg);
}

int session_manager_subscribe_tcp(struct session_manager *sess_mgr, on_session_message_callback *cb, void *args)
{
     assert(sess_mgr);
     assert(cb);
     return mq_schema_subscribe(sess_mgr->sche->mq_sche, sess_mgr->sche->sess_msg_id_tcp, (on_msg_cb_func *)(void *)cb, args);
}

int session_manager_subscribe_udp(struct session_manager *sess_mgr, on_session_message_callback *cb, void *args)
{
     assert(sess_mgr);
     assert(cb);
     return mq_schema_subscribe(sess_mgr->sche->mq_sche, sess_mgr->sche->sess_msg_id_udp, (on_msg_cb_func *)(void *)cb, args);
}

int session_manager_subscribe_control_packet(struct session_manager *sess_mgr, on_session_message_callback *cb, void *args)
{
     assert(sess_mgr);
     assert(cb);
     return mq_schema_subscribe(sess_mgr->sche->mq_sche, sess_mgr->sche->sess_msg_id_ctrl, (on_msg_cb_func *)(void *)cb, args);
}

int session_manager_subscribe_tcp_stream(struct session_manager *sess_mgr, on_tcp_payload_callback *cb, void *args)
{
     assert(sess_mgr);
     assert(cb);
     return mq_schema_subscribe(sess_mgr->sche->mq_sche, sess_mgr->sche->sess_msg_id_stream, (on_msg_cb_func *)(void *)cb, args);
}

int session_manager_init(struct session_manager *sess_mgr, uint16_t thread_id, struct mq_runtime *mq_rte)
{
     assert(sess_mgr);
     uint64_t now_ms = clock_get_real_time_ms();

     sess_mgr->cfg->session_id_seed = sess_mgr->cfg->instance_id << 8 | thread_id;
     sess_mgr->mq[thread_id] = mq_rte;
     sess_mgr->rte[thread_id] = session_manager_rte_new(sess_mgr->cfg, now_ms);
     if (sess_mgr->rte[thread_id] == NULL)
     {
          SESSION_MANAGER_LOG_ERROR("failed to create session_manager_rte");
          return -1;
     }
     else
     {
          return 0;
     }
}

void session_manager_clean(struct session_manager *sess_mgr, uint16_t thread_id)
{
     assert(sess_mgr);

     struct mq_runtime *mq_rte = sess_mgr->mq[thread_id];
     if (sess_mgr->rte[thread_id] == NULL)
     {
          return;
     }

     struct session_manager_stat *stat = session_manager_rte_get_stat(sess_mgr->rte[thread_id]);
     while (stat->tcp_sess_used || stat->udp_sess_used)
     {
          clean_closed_session(sess_mgr, thread_id, UINT64_MAX);
          // here we need to dispatch the message to ensure that the session is cleaned up
          mq_runtime_dispatch(mq_rte);
     }

     session_manager_rte_free(sess_mgr->rte[thread_id]);
     sess_mgr->rte[thread_id] = NULL;
}

/******************************************************************************
 * 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_sche = stellar_module_manager_get_mq_schema(mod_mgr);
     assert(mq_sche);
     const char *toml_file = stellar_module_manager_get_toml_path(mod_mgr);
     assert(toml_file);

     struct session_manager *sess_mgr = session_manager_new(pkt_mgr, mq_sche, toml_file);
     if (sess_mgr == NULL)
     {
          return NULL;
     }
     stellar_module_manager_polling_subscribe(mod_mgr, on_polling, sess_mgr);

     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 init");
     return sess_mgr_mod;
}

void session_manager_on_exit(struct stellar_module_manager *mod_mgr, 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 exit");
     }
}

struct stellar_module *session_manager_on_thread_init(struct stellar_module_manager *mod_mgr, int thread_id, struct stellar_module *mod)
{
     struct session_manager *sess_mgr = stellar_module_get_ctx(mod);
     assert(sess_mgr);
     struct mq_runtime *mq_rte = stellar_module_manager_get_mq_runtime(mod_mgr);
     assert(mq_rte);

     if (session_manager_init(sess_mgr, thread_id, mq_rte) != 0)
     {
          SESSION_MANAGER_LOG_ERROR("failed to int session_manager_init");
          return NULL;
     }
     else
     {
          return mod;
     }
}

void session_manager_on_thread_exit(struct stellar_module_manager *mod_mgr, int thread_id, struct stellar_module *mod)
{
     struct session_manager *sess_mgr = stellar_module_get_ctx(mod);
     assert(sess_mgr);
     session_manager_clean(sess_mgr, thread_id);
}