#include <assert.h>

#include "utils_internal.h"
#include "tcp_reassembly.h"
#include "session_internal.h"
#include "session_manager.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

#pragma GCC diagnostic ignored "-Wunused-parameter"

struct session_manager
{
     int pkt_ex_to_get_sess;
     int pkt_ex_to_free_sess;
     int pkt_ex_to_tcp_segment;
     struct exdata_schema *ex_sche;

     int stat_idx[SESS_MGR_STAT_MAX];
     struct fieldstat_easy *fs;
     struct session_manager_cfg *cfg;
     struct session_manager_rte *rte[MAX_THREAD_NUM];
     struct module_manager *mod_mgr;
     struct packet_manager *pkt_mgr;
};

static void free_session(int idx, void *ex_ptr, void *arg)
{
     struct session *sess = (struct session *)ex_ptr;
     struct session_manager *sess_mgr = (struct session_manager *)arg;
     assert(idx == sess_mgr->pkt_ex_to_free_sess);

     if (sess && session_get_current_state(sess) == SESSION_STATE_CLOSED)
     {
          int thread_id = module_manager_get_thread_id(sess_mgr->mod_mgr);
          struct session_manager_rte *sess_mgr_rte = session_manager_get_rte(sess_mgr, 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 free_tcp_segment(int idx, void *ex_ptr, void *arg)
{
     struct tcp_segment *segs = (struct tcp_segment *)ex_ptr;
     while (segs)
     {
          struct session *sess = (struct session *)tcp_segment_get_user_data(segs);
          struct tcp_segment *next = segs->next;
          session_free_tcp_segment(sess, segs);
          segs = next;
     }
}

static void notify_sess_closed_by_pseudo_pkt(struct session_manager *sess_mgr, int thread_id, struct session *sess)
{
     struct packet *pseudo = NULL;
     struct packet_manager *pkt_mgr = sess_mgr->pkt_mgr;
     assert(session_get_current_state(sess) == SESSION_STATE_CLOSED);

     if (session_get_first_packet(sess, FLOW_TYPE_C2S))
     {
          pseudo = packet_manager_dup_packet(pkt_mgr, thread_id, session_get_first_packet(sess, FLOW_TYPE_C2S));
     }
     else
     {
          pseudo = packet_manager_dup_packet(pkt_mgr, thread_id, session_get_first_packet(sess, FLOW_TYPE_S2C));
     }

     assert(pseudo);
     packet_tag_set(pseudo, PKT_TAG_KEY_SESS, PKT_TAG_VAL_SESS_FREE);
     packet_set_type(pseudo, PACKET_TYPE_PSEUDO);
     packet_set_action(pseudo, PACKET_ACTION_DROP);
     packet_set_exdata(pseudo, sess_mgr->pkt_ex_to_free_sess, sess);
     packet_manager_schedule_packet(pkt_mgr, thread_id, pseudo, PACKET_STAGE_FORWARD);
     SESSION_MANAGER_LOG_INFO("notify session %lu %s closed by pseudo packet: %p", session_get_id(sess), session_get_readable_addr(sess), pseudo);
}

static void on_polling(struct module_manager *mod_mgr, void *args)
{
     struct session *sess = NULL;
     uint64_t now_ms = clock_get_real_time_ms();
     struct session_manager *sess_mgr = (struct session_manager *)args;
     int thread_id = module_manager_get_thread_id(mod_mgr);
     struct session_manager_rte *sess_mgr_rte = session_manager_get_rte(sess_mgr, thread_id);

     static __thread uint64_t last_clean_expired_sess_ts = 0;
     if (now_ms - last_clean_expired_sess_ts > sess_mgr->cfg->expire_period_ms)
     {
          for (uint64_t i = 0; i < sess_mgr->cfg->expire_batch_max; i++)
          {
               sess = session_manager_rte_get_expired_session(sess_mgr_rte, now_ms);
               if (sess)
               {
                    notify_sess_closed_by_pseudo_pkt(sess_mgr, thread_id, sess);
               }
               else
               {
                    break;
               }
          }
          last_clean_expired_sess_ts = 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)
     {
          struct session_manager_stat *sess_mgr_stat = session_manager_rte_get_stat(sess_mgr_rte);
          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;
     }
}

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

void session_manager_on_packet_forward(struct packet *pkt, struct module *mod)
{
     struct session_manager *sess_mgr = (struct session_manager *)module_get_ctx(mod);
     int thread_id = module_manager_get_thread_id(sess_mgr->mod_mgr);
     struct session_manager_rte *sess_mgr_rte = session_manager_get_rte(sess_mgr, 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 session *sess = session_manager_rte_lookup_session_by_packet(sess_mgr_rte, pkt);
     if (sess == NULL)
     {
          if (packet_get_type(pkt) == PACKET_TYPE_RAW)
          {
               sess = session_manager_rte_new_session(sess_mgr_rte, pkt, now_ms);
               if (sess)
               {
                    session_set_user_data(sess, exdata_runtime_new(sess_mgr->ex_sche));
                    packet_tag_set(pkt, PKT_TAG_KEY_SESS, PKT_TAG_VAL_SESS_NEW);
               }
          }
          else
          {
               // TODO new session by pseudo packet
          }
     }
     else
     {
          if (packet_get_type(pkt) == PACKET_TYPE_PSEUDO)
          {
               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);
                    sess->stats[FLOW_TYPE_C2S][STAT_PSEUDO_PACKETS_RECEIVED]++;
                    sess->stats[FLOW_TYPE_C2S][STAT_PSEUDO_BYTES_RECEIVED] += packet_get_raw_len(pkt);
               }
               else
               {
                    session_set_flow_type(sess, FLOW_TYPE_S2C);
                    sess->stats[FLOW_TYPE_S2C][STAT_PSEUDO_PACKETS_RECEIVED]++;
                    sess->stats[FLOW_TYPE_S2C][STAT_PSEUDO_BYTES_RECEIVED] += packet_get_raw_len(pkt);
               }
               packet_tag_set(pkt, PKT_TAG_KEY_SESS, PKT_TAG_VAL_SESS_CTRL_MSG);
          }
          else
          {
               session_manager_rte_update_session(sess_mgr_rte, sess, pkt, now_ms);
               packet_tag_set(pkt, PKT_TAG_KEY_SESS, 0);
          }
     }

     packet_set_exdata(pkt, sess_mgr->pkt_ex_to_get_sess, sess);

     if (sess && session_get_type(sess) == SESSION_TYPE_TCP)
     {
          struct tcp_segment *prev = NULL;
          struct tcp_segment *next = NULL;
          while ((next = session_get_tcp_segment(sess)))
          {
               tcp_segment_set_user_data(next, sess);
               if (prev == NULL)
               {
                    prev = next;
                    prev->next = NULL;
                    packet_set_exdata(pkt, sess_mgr->pkt_ex_to_tcp_segment, prev);
                    packet_tag_set(pkt, PKT_TAG_KEY_SESS, PKT_TAG_VAL_SESS_TCP_STREAM);
               }
               else
               {
                    prev->next = next;
                    prev = next;
                    prev->next = NULL;
               }
          }
     }

     while ((sess = session_manager_rte_get_evicted_session(sess_mgr_rte)))
     {
          notify_sess_closed_by_pseudo_pkt(sess_mgr, thread_id, sess);
     }
}

void session_manager_on_packet_output(struct packet *pkt, struct module *mod)
{
     struct session_manager *sess_mgr = (struct session_manager *)module_get_ctx(mod);
     int thread_id = module_manager_get_thread_id(sess_mgr->mod_mgr);
     struct session_manager_rte *sess_mgr_rte = session_manager_get_rte(sess_mgr, thread_id);

     struct session *sess = (struct session *)packet_get_exdata(pkt, sess_mgr->pkt_ex_to_get_sess);
     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_pseudo = (packet_get_type(pkt) == PACKET_TYPE_PSEUDO);
          uint16_t len = packet_get_raw_len(pkt);
          switch (packet_get_action(pkt))
          {
          case PACKET_ACTION_DROP:
               session_inc_stat(sess, flow, (is_pseudo ? STAT_PSEUDO_PACKETS_DROPPED : STAT_RAW_PACKETS_DROPPED), 1);
               session_inc_stat(sess, flow, (is_pseudo ? STAT_PSEUDO_BYTES_DROPPED : STAT_RAW_BYTES_DROPPED), len);
               break;
          case PACKET_ACTION_FORWARD:
               session_inc_stat(sess, flow, (is_pseudo ? STAT_PSEUDO_PACKETS_TRANSMITTED : STAT_RAW_PACKETS_TRANSMITTED), 1);
               session_inc_stat(sess, flow, (is_pseudo ? STAT_PSEUDO_BYTES_TRANSMITTED : STAT_RAW_BYTES_TRANSMITTED), len);
               break;
          default:
               assert(0);
               break;
          }

          session_set_current_packet(sess, NULL);
     }

     if (packet_get_origin(pkt) == NULL)
     {
          session_manager_rte_record_duplicated_packet(sess_mgr_rte, pkt);
     }
}

struct session_manager *session_manager_new(struct packet_manager *pkt_mgr, 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, "metrics/session_manager.json", 2) != 0)
     {
          SESSION_MANAGER_LOG_ERROR("failed to enable auto output");
          goto error_out;
     }
#if 0
     if (packet_manager_register_node(pkt_mgr, "session_manager", PACKET_STAGE_FORWARD, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP | PKT_TAG_VAL_IPPROTO_UDP, session_manager_on_packet_forward, sess_mgr))
     {
          SESSION_MANAGER_LOG_ERROR("failed to subscribe PACKET_STAGE_FORWARD");
          goto error_out;
     }
     if (packet_manager_register_node(pkt_mgr, "session_manager", PACKET_STAGE_OUTPUT, PKT_TAG_KEY_IPPROTO, PKT_TAG_VAL_IPPROTO_TCP | PKT_TAG_VAL_IPPROTO_UDP, session_manager_on_packet_output, sess_mgr))
     {
          SESSION_MANAGER_LOG_ERROR("failed to subscribe PACKET_STAGE_OUTPUT");
          goto error_out;
     }
#endif

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

     sess_mgr->pkt_ex_to_get_sess = packet_manager_new_packet_exdata_index(pkt_mgr, "pkt_ex_key_for_get_sess", NULL, NULL);
     if (sess_mgr->pkt_ex_to_get_sess == -1)
     {
          SESSION_MANAGER_LOG_ERROR("failed to create packet exdata index");
          goto error_out;
     }

     sess_mgr->pkt_ex_to_free_sess = packet_manager_new_packet_exdata_index(pkt_mgr, "pkt_ex_key_for_free_sess", free_session, sess_mgr);
     if (sess_mgr->pkt_ex_to_free_sess == -1)
     {
          SESSION_MANAGER_LOG_ERROR("failed to create packet exdata index");
          goto error_out;
     }

     sess_mgr->pkt_ex_to_tcp_segment = packet_manager_new_packet_exdata_index(pkt_mgr, "pkt_ex_key_for_tcp_stream", free_tcp_segment, sess_mgr);
     if (sess_mgr->pkt_ex_to_tcp_segment == -1)
     {
          SESSION_MANAGER_LOG_ERROR("failed to create packet exdata index");
          goto error_out;
     }

     sess_mgr->pkt_mgr = pkt_mgr;
     return sess_mgr;

error_out:
     session_manager_free(sess_mgr);
     return NULL;
}

void session_manager_free(struct session_manager *sess_mgr)
{
     if (sess_mgr)
     {
          if (sess_mgr->ex_sche)
          {
               exdata_schema_free(sess_mgr->ex_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);
     }
}

struct session_manager_cfg *session_manager_get_cfg(const struct session_manager *sess_mgr)
{
     assert(sess_mgr);
     return sess_mgr->cfg;
}

struct session_manager_rte *session_manager_get_rte(const struct session_manager *sess_mgr, uint16_t thread_id)
{
     assert(sess_mgr);
     if (thread_id < sess_mgr->cfg->thread_num)
     {
          return sess_mgr->rte[thread_id];
     }
     else
     {
          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->ex_sche, name, func, arg);
}

int session_manager_init(struct session_manager *sess_mgr, uint16_t thread_id)
{
     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->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 session_manager_rte *rte = session_manager_get_rte(sess_mgr, thread_id);

     if (rte == NULL)
     {
          return;
     }

     struct session *sess = NULL;
     struct session_manager_stat *stat = session_manager_rte_get_stat(rte);

     while (stat->tcp_sess_used || stat->udp_sess_used)
     {
          while ((sess = session_manager_rte_get_expired_session(rte, UINT64_MAX)))
          {
               exdata_runtime_free((struct exdata_runtime *)session_get_user_data(sess));
               session_manager_rte_free_session(rte, sess);
          }
     }

     session_manager_rte_free(rte);
}

struct session *session_manager_lookup_session_by_packet(struct session_manager *sess_mgr, uint16_t thread_id, const struct packet *pkt)
{
     struct session_manager_rte *rte = session_manager_get_rte(sess_mgr, thread_id);
     if (rte == NULL)
     {
          return NULL;
     }
     return session_manager_rte_lookup_session_by_packet(rte, pkt);
}

struct session *session_manager_lookup_session_by_id(struct session_manager *sess_mgr, uint16_t thread_id, uint64_t sess_id)
{
     struct session_manager_rte *rte = session_manager_get_rte(sess_mgr, thread_id);
     if (rte == NULL)
     {
          return NULL;
     }
     return session_manager_rte_lookup_session_by_id(rte, sess_id);
}

void session_manager_discard_session(struct session_manager *sess_mgr, uint16_t thread_id, struct session *sess)
{
     struct session_manager_rte *rte = session_manager_get_rte(sess_mgr, thread_id);
     if (rte == NULL)
     {
          return;
     }
     session_manager_rte_discard_session(rte, sess);
}

struct session *packet_exdata_to_session(const struct session_manager *sess_mgr, const struct packet *pkt)
{
     return (struct session *)packet_get_exdata(pkt, sess_mgr->pkt_ex_to_get_sess);
}

struct tcp_segment *packet_exdata_to_tcp_segment(const struct session_manager *sess_mgr, const struct packet *pkt)
{
     return (struct tcp_segment *)packet_get_exdata(pkt, sess_mgr->pkt_ex_to_tcp_segment);
}

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

struct session_manager *module_to_session_manager(struct module *mod)
{
     assert(mod);
     assert(strcmp(module_get_name(mod), SESSION_MANAGER_MODULE_NAME) == 0);
     return module_get_ctx(mod);
}

struct module *session_manager_on_init(struct module_manager *mod_mgr)
{
     assert(mod_mgr);
     struct module *pkt_mgr_mod = module_manager_get_module(mod_mgr, PACKET_MANAGER_MODULE_NAME);
     struct packet_manager *pkt_mgr = module_to_packet_manager(pkt_mgr_mod);
     assert(pkt_mgr);
     const char *toml_file = module_manager_get_toml_path(mod_mgr);
     assert(toml_file);

     struct session_manager *sess_mgr = session_manager_new(pkt_mgr, toml_file);
     if (sess_mgr == NULL)
     {
          return NULL;
     }
     sess_mgr->mod_mgr = mod_mgr;
     module_manager_register_polling_node(mod_mgr, on_polling, sess_mgr);

     struct module *sess_mgr_mod = 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;
     }
     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 module_manager *mod_mgr, struct module *mod)
{
     if (mod)
     {
          struct session_manager *sess_mgr = module_get_ctx(mod);
          session_manager_free(sess_mgr);
          module_free(mod);
          SESSION_MANAGER_LOG_FATAL("session_manager exit");
     }
}

struct module *session_manager_on_thread_init(struct module_manager *mod_mgr, int thread_id, struct module *mod)
{
     struct session_manager *sess_mgr = module_get_ctx(mod);
     assert(sess_mgr);

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

void session_manager_on_thread_exit(struct module_manager *mod_mgr, int thread_id, struct module *mod)
{
     struct session_manager *sess_mgr = module_get_ctx(mod);
     if (sess_mgr)
     {
          assert(thread_id < (int)sess_mgr->cfg->thread_num);
          session_manager_clean(sess_mgr, thread_id);
     }
}