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

#include "timestamp.h"
#include "session_manager.h"
#include "session_pool.h"
#include "session_table.h"
#include "session_timer.h"
#include "session_queue.h"
#include "session_private.h"
#include "packet_helpers.h"
#include "tcp_helpers.h"
#include "udp_helpers.h"

struct session_manager
{
     struct session_pool *sess_pool;
     struct session_table *sess_table;
     struct session_timer *sess_timer;
     struct session_queue *evicted_sess;

     // timeout config
     uint64_t timeout_toclosing;
     uint64_t timeout_toclosed;

     // session number
     uint64_t tcp_opening_sess_num;
     uint64_t tcp_closing_sess_num;
     uint64_t tcp_active_sess_num;

     uint64_t udp_opening_sess_num;
     uint64_t udp_closing_sess_num;
     uint64_t udp_active_sess_num;
};

/******************************************************************************
 * utils
 ******************************************************************************/

// TODO
static uint64_t alloc_session_id(void)
{
     return 0;
}

static void metadata_ex_free_cb(struct session *sess, uint8_t idx, void *ex_ptr, void *arg)
{
     if (ex_ptr)
     {
          metadata_free((struct metadata *)ex_ptr);
     }
}

static void packet_ex_free_cb(struct session *sess, uint8_t idx, void *ex_ptr, void *arg)
{
     if (ex_ptr)
     {
          packet_free((struct packet *)ex_ptr);
     }
}

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

static void update_counter_on_opening(struct session_manager *mgr, struct session *sess)
{
     if (session_get_state(sess) == SESSION_STATE_INIT)
     {
          if (session_get_type(sess) == SESSION_TYPE_TCP)
          {
               mgr->tcp_opening_sess_num++;
          }
          else
          {
               mgr->udp_opening_sess_num++;
          }
     }
}

static void update_counter_on_active(struct session_manager *mgr, struct session *sess)
{
     if (session_get_state(sess) == SESSION_STATE_OPENING)
     {
          if (session_get_type(sess) == SESSION_TYPE_TCP)
          {
               mgr->tcp_opening_sess_num--;
               mgr->tcp_active_sess_num++;
          }
          else
          {
               mgr->udp_opening_sess_num--;
               mgr->udp_active_sess_num++;
          }
     }
}

static void update_counter_on_closing(struct session_manager *mgr, struct session *sess)
{
     if (session_get_state(sess) == SESSION_STATE_OPENING)
     {
          if (session_get_type(sess) == SESSION_TYPE_TCP)
          {
               mgr->tcp_opening_sess_num--;
               mgr->tcp_closing_sess_num++;
          }
          else
          {
               mgr->udp_opening_sess_num--;
               mgr->udp_closing_sess_num++;
          }

          return;
     }

     if (session_get_state(sess) == SESSION_STATE_ACTIVE)
     {
          if (session_get_type(sess) == SESSION_TYPE_TCP)
          {
               mgr->tcp_active_sess_num--;
               mgr->tcp_closing_sess_num++;
          }
          else
          {
               mgr->udp_active_sess_num--;
               mgr->udp_closing_sess_num++;
          }

          return;
     }
}

static void update_counter_on_closed(struct session_manager *mgr, struct session *sess)
{
     if (session_get_state(sess) == SESSION_STATE_CLOSING)
     {
          if (session_get_type(sess) == SESSION_TYPE_TCP)
          {
               mgr->tcp_closing_sess_num--;
          }
          else
          {
               mgr->udp_closing_sess_num--;
          }
     }
}

/******************************************************************************
 * judge session direction
 ******************************************************************************/

static enum session_dir judge_direction_by_tuple6(const struct tuple6 *key)
{
     // big port is client
     if (ntohs(key->src_port) > ntohs(key->dst_port))
     {
          return SESSION_DIR_C2S;
     }
     else
     {
          return SESSION_DIR_S2C;
     }
}

static enum session_dir judge_direction_by_session(const struct session *sess, const struct tuple6 *key)
{
     if (tuple6_cmp(session_get0_tuple6(sess), key) == 0)
     {
          return session_get_tuple6_dir(sess);
     }
     else
     {
          if (session_get_tuple6_dir(sess) == SESSION_DIR_C2S)
          {
               return SESSION_DIR_S2C;
          }
          else
          {
               return SESSION_DIR_C2S;
          }
     }
}

/******************************************************************************
 * update session timer
 ******************************************************************************/

void session_manager_update_session_timer(struct session_manager *mgr, struct session *sess, session_expire_cb cb, uint64_t timeout_ms)
{
     session_timer_del_session(mgr->sess_timer, sess);
     session_set_expirecb(sess, cb, mgr, timestamp_get_msec() + timeout_ms);
     session_timer_add_session(mgr->sess_timer, sess);
}

/******************************************************************************
 * expire callback
 ******************************************************************************/

static void session_to_closed(struct session *sess, void *arg)
{
     SESSION_LOG_DEBUG("session %lu closing expire, free session", session_get_id(sess));
     struct session_manager *mgr = (struct session_manager *)arg;
     assert(mgr != NULL);

     update_counter_on_closed(mgr, sess);
     session_set_state(sess, SESSION_STATE_CLOSED);
     session_set0_cur_pkt(sess, NULL);
     session_set_cur_dir(sess, SESSION_DIR_NONE);
     for (uint8_t i = 0; i < EX_DATA_MAX_COUNT; i++)
     {
          session_free_ex_data(sess, i);
     }
     session_table_delete_session(mgr->sess_table, session_get0_tuple6(sess));
     session_timer_del_session(mgr->sess_timer, sess);
     session_pool_free(mgr->sess_pool, sess);
}

static void session_to_closing(struct session *sess, void *arg)
{
     SESSION_LOG_DEBUG("session %lu packet expire, trigger closing event", session_get_id(sess));
     struct session_manager *mgr = (struct session_manager *)arg;
     assert(mgr != NULL);

     update_counter_on_closing(mgr, sess);
     session_set_state(sess, SESSION_STATE_CLOSING);
     session_manager_update_session_timer(mgr, sess, session_to_closed, mgr->timeout_toclosed);
}

/******************************************************************************
 * session ex data
 ******************************************************************************/

static int tcp_need_closing(uint64_t state)
{
     if ((state & TCP_C2S_FIN_RECVED) && (state & TCP_S2C_FIN_RECVED))
     {
          return 1;
     }

     if (state & TCP_C2S_RST_RECVED)
     {
          return 1;
     }

     if (state & TCP_S2C_RST_RECVED)
     {
          return 1;
     }

     return 0;
}

static int tcp_need_active(uint64_t state)
{
     if ((state & TCP_C2S_PAYLOAD_RECVED) || (state & TCP_S2C_PAYLOAD_RECVED))
     {
          return 1;
     }

     return 0;
}

static void update_session_base(struct session *sess, const struct packet *pkt, enum session_dir curr_dir)
{
     uint64_t len = packet_get_raw_len(pkt);
     const struct metadata *md = packet_get0_metadata(pkt);
     if (curr_dir == SESSION_DIR_C2S)
     {
          session_inc_c2s_metrics(sess, 1, len);
          if (session_get0_ex_data(sess, c2s_1st_md_ex) == NULL)
          {
               session_set_ex_data(sess, c2s_1st_md_ex, metadata_dup(md));
          }
          if (session_get0_ex_data(sess, c2s_1st_pkt_ex) == NULL)
          {
               session_set_ex_data(sess, c2s_1st_pkt_ex, packet_dup(pkt));
          }
     }
     else
     {
          session_inc_s2c_metrics(sess, 1, len);
          if (session_get0_ex_data(sess, s2c_1st_md_ex) == NULL)
          {
               session_set_ex_data(sess, s2c_1st_md_ex, metadata_dup(md));
          }
          if (session_get0_ex_data(sess, s2c_1st_pkt_ex) == NULL)
          {
               session_set_ex_data(sess, s2c_1st_pkt_ex, packet_dup(pkt));
          }
     }
     session_set_last_time(sess, timestamp_get_msec());
     session_set0_cur_pkt(sess, pkt);
     session_set_cur_dir(sess, curr_dir);
}

static void update_tcp_ex_data(struct session *sess, const struct packet *pkt, enum session_dir curr_dir)
{
     const struct layer_record *tcp_layer = packet_get_innermost_layer(pkt, LAYER_TYPE_TCP);
     const struct tcphdr *hdr = (const struct tcphdr *)tcp_layer->hdr_ptr;

     uint64_t state = (uint64_t)session_get0_ex_data(sess, tcp_builtin_ex);
     if (tcp_hdr_has_flag_rst(hdr))
     {
          if (curr_dir == SESSION_DIR_C2S)
          {
               state |= TCP_C2S_RST_RECVED;
               session_set_ex_data(sess, tcp_builtin_ex, (void *)(state));
          }
          else
          {
               state |= TCP_S2C_RST_RECVED;
               session_set_ex_data(sess, tcp_builtin_ex, (void *)(state));
          }
     }

     if (tcp_hdr_has_flag_fin(hdr))
     {
          if (curr_dir == SESSION_DIR_C2S)
          {
               state |= TCP_C2S_FIN_RECVED;
               session_set_ex_data(sess, tcp_builtin_ex, (void *)(state));
          }
          else
          {
               state |= TCP_S2C_FIN_RECVED;
               session_set_ex_data(sess, tcp_builtin_ex, (void *)(state));
          }
     }

     if (tcp_hdr_has_flag_syn(hdr))
     {
          if (tcp_hdr_has_flag_ack(hdr))
          {
               state |= TCP_SYNACK_RECVED;
               session_set_ex_data(sess, tcp_builtin_ex, (void *)(state));
          }
          else
          {
               state |= TCP_SYN_RECVED;
               session_set_ex_data(sess, tcp_builtin_ex, (void *)(state));
          }
     }

     if (tcp_layer->pld_len > 0)
     {
          if (curr_dir == SESSION_DIR_C2S)
          {
               state |= TCP_C2S_PAYLOAD_RECVED;
               session_set_ex_data(sess, tcp_builtin_ex, (void *)(state));
          }
          else
          {
               state |= TCP_S2C_PAYLOAD_RECVED;
               session_set_ex_data(sess, tcp_builtin_ex, (void *)(state));
          }
     }
}

static void update_udp_ex_data(struct session *sess, const struct packet *pkt, enum session_dir curr_dir)
{
     uint64_t state = (uint64_t)session_get0_ex_data(sess, udp_builtin_ex);
     if (curr_dir == SESSION_DIR_C2S)
     {
          session_set_ex_data(sess, udp_builtin_ex, (void *)(state | UDP_C2S_RECVED));
     }
     else
     {
          session_set_ex_data(sess, udp_builtin_ex, (void *)(state | UDP_S2C_RECVED));
     }
}

/******************************************************************************
 * handle session
 ******************************************************************************/

// return  0: success
// return -1: tcp not syn packet, discard
static int handle_tcp_new_session(struct session_manager *mgr, struct tuple6 *key, struct session *sess, const struct packet *pkt)
{
     const struct layer_record *tcp_layer = packet_get_innermost_layer(pkt, LAYER_TYPE_TCP);
     if (tcp_layer == NULL)
     {
          // not tcp packet, discard
          return -1;
     }
     const struct tcphdr *hdr = (const struct tcphdr *)tcp_layer->hdr_ptr;
     if (!tcp_hdr_has_flag_syn(hdr))
     {
          // not syn packet, discard
          return -1;
     }

     enum session_dir curr_dir = SESSION_DIR_NONE;

     session_init(sess);
     // syn packet
     if (!tcp_hdr_has_flag_ack(hdr))
     {
          curr_dir = SESSION_DIR_C2S;
          session_set_ex_data(sess, tcp_builtin_ex, (void *)TCP_SYN_RECVED);
     }
     // syn ack packet
     else
     {
          curr_dir = SESSION_DIR_S2C;
          session_set_ex_data(sess, tcp_builtin_ex, (void *)TCP_SYNACK_RECVED);
     }

     session_set_id(sess, alloc_session_id());
     session_set_tuple6(sess, key);
     session_set_tuple6_dir(sess, curr_dir);

     session_set_type(sess, SESSION_TYPE_TCP);
     update_counter_on_opening(mgr, sess);
     session_set_state(sess, SESSION_STATE_OPENING);

     session_set_create_time(sess, timestamp_get_msec());
     update_session_base(sess, pkt, curr_dir);

     session_manager_update_session_timer(mgr, sess, session_to_closing, mgr->timeout_toclosing);

     return 0;
}

// always return 0
static int handle_udp_new_session(struct session_manager *mgr, struct tuple6 *key, struct session *sess, const struct packet *pkt)
{
     enum session_dir curr_dir = judge_direction_by_tuple6(key);

     session_init(sess);
     update_udp_ex_data(sess, pkt, curr_dir);
     session_set_id(sess, alloc_session_id());
     session_set_tuple6(sess, key);
     session_set_tuple6_dir(sess, curr_dir);

     /*
      * when a UDP Session is created, the Opening and active events are triggered,
      * (the plugin is called twice by the opening/active events in turn),
      * and the state of the UDP session is directly switched to the active state.
      */
     session_set_type(sess, SESSION_TYPE_UDP);
     update_counter_on_opening(mgr, sess);
     session_set_state(sess, SESSION_STATE_OPENING);
     update_counter_on_active(mgr, sess);
     session_set_state(sess, SESSION_STATE_ACTIVE);

     session_set_create_time(sess, timestamp_get_msec());
     update_session_base(sess, pkt, curr_dir);

     session_manager_update_session_timer(mgr, sess, session_to_closing, mgr->timeout_toclosing);

     return 0;
}

static void handle_tcp_old_session(struct session_manager *mgr, struct tuple6 *key, struct session *sess, const struct packet *pkt)
{
     enum session_dir curr_dir = judge_direction_by_session(sess, key);
     update_tcp_ex_data(sess, pkt, curr_dir);
     update_session_base(sess, pkt, curr_dir);

     uint64_t state = (uint64_t)session_get0_ex_data(sess, tcp_builtin_ex);
     if (tcp_need_closing(state))
     {
          update_counter_on_closing(mgr, sess);
          session_set_state(sess, SESSION_STATE_CLOSING);
          session_manager_update_session_timer(mgr, sess, session_to_closed, mgr->timeout_toclosed);
          return;
     }

     if (tcp_need_active(state))
     {
          update_counter_on_active(mgr, sess);
          session_set_state(sess, SESSION_STATE_ACTIVE);
          session_manager_update_session_timer(mgr, sess, session_to_closing, mgr->timeout_toclosing);
          return;
     }
}

static void handle_udp_old_session(struct session_manager *mgr, struct tuple6 *key, struct session *sess, const struct packet *pkt)
{
     enum session_dir curr_dir = judge_direction_by_session(sess, key);
     update_udp_ex_data(sess, pkt, curr_dir);
     update_session_base(sess, pkt, curr_dir);

     update_counter_on_active(mgr, sess);
     session_set_state(sess, SESSION_STATE_ACTIVE);
     session_manager_update_session_timer(mgr, sess, session_to_closing, mgr->timeout_toclosing);
}

// return  0: success
// return -1: tcp not syn packet, discard
static int handle_new_session(struct session_manager *mgr, struct tuple6 *key, struct session *sess, const struct packet *pkt)
{
     if (key->ip_proto == IPPROTO_TCP)
     {
          return handle_tcp_new_session(mgr, key, sess, pkt);
     }
     else
     {
          return handle_udp_new_session(mgr, key, sess, pkt);
     }
}

static void handle_old_session(struct session_manager *mgr, struct tuple6 *key, struct session *sess, const struct packet *pkt)
{
     if (session_get_state(sess) == SESSION_STATE_CLOSING)
     {
          return;
     }

     if (key->ip_proto == IPPROTO_TCP)
     {
          handle_tcp_old_session(mgr, key, sess, pkt);
     }
     else
     {
          handle_udp_old_session(mgr, key, sess, pkt);
     }
}

/******************************************************************************
 * public API
 ******************************************************************************/

struct session_manager *session_manager_create(uint64_t max_session_num)
{
     struct session_manager *mgr = (struct session_manager *)calloc(1, sizeof(struct session_manager));
     if (mgr == NULL)
     {
          return NULL;
     }

     mgr->sess_pool = session_pool_create(max_session_num);
     if (mgr->sess_pool == NULL)
     {
          goto error;
     }

     mgr->sess_table = session_table_create();
     if (mgr->sess_table == NULL)
     {
          goto error;
     }

     mgr->sess_timer = session_timer_create();
     if (mgr->sess_timer == NULL)
     {
          goto error;
     }

     mgr->evicted_sess = session_queue_create();
     if (mgr->evicted_sess == NULL)
     {
          goto error;
     }

     mgr->timeout_toclosed = 2 * 1000;
     mgr->timeout_toclosing = 5 * 1000;

     mgr->tcp_opening_sess_num = 0;
     mgr->tcp_closing_sess_num = 0;
     mgr->tcp_active_sess_num = 0;

     mgr->udp_opening_sess_num = 0;
     mgr->udp_closing_sess_num = 0;
     mgr->udp_active_sess_num = 0;

     tcp_builtin_ex = session_get_ex_new_index("tcp_builtin_ex", NULL, NULL);
     udp_builtin_ex = session_get_ex_new_index("udp_builtin_ex", NULL, NULL);
     c2s_1st_md_ex = session_get_ex_new_index("c2s_1st_md_ex", metadata_ex_free_cb, NULL);
     s2c_1st_md_ex = session_get_ex_new_index("s2c_1st_md_ex", metadata_ex_free_cb, NULL);
     c2s_1st_pkt_ex = session_get_ex_new_index("c2s_1st_pkt_ex", packet_ex_free_cb, NULL);
     s2c_1st_pkt_ex = session_get_ex_new_index("s2c_1st_pkt_ex", packet_ex_free_cb, NULL);

     return mgr;

error:
     session_manager_destroy(mgr);
     return NULL;
}

void session_manager_destroy(struct session_manager *mgr)
{
     if (mgr)
     {
          session_queue_destroy(mgr->evicted_sess);
          session_timer_destroy(mgr->sess_timer);
          session_table_destroy(mgr->sess_table);
          session_pool_destroy(mgr->sess_pool);
          free(mgr);
          mgr = NULL;
     }
}

void session_manager_set_timeout_toclosing(struct session_manager *mgr, uint64_t timeout_ms)
{
     mgr->timeout_toclosing = timeout_ms;
}

void session_manager_set_timeout_toclosed(struct session_manager *mgr, uint64_t timeout_ms)
{
     mgr->timeout_toclosed = timeout_ms;
}

struct session *session_manager_lookup(struct session_manager *mgr, const struct packet *pkt)
{
     struct tuple6 key;
     if (packet_get_innermost_tuple6(pkt, &key) == -1)
     {
          return NULL;
     }

     return session_table_find_session(mgr->sess_table, &key);
}

// return null: invalid tuple6 or tcp first packet is not syn
struct session *session_manager_update(struct session_manager *mgr, const struct packet *pkt)
{
     struct tuple6 key;
     if (packet_get_innermost_tuple6(pkt, &key) == -1)
     {
          return NULL;
     }

     struct session *sess = session_table_find_session(mgr->sess_table, &key);
     if (sess == NULL)
     {
          if (session_pool_get_count(mgr->sess_pool) == 1)
          {
               struct session *unused_sess = session_table_find_least_recently_unused_session(mgr->sess_table);
               assert(unused_sess);

               update_counter_on_closing(mgr, unused_sess);
               session_set_state(unused_sess, SESSION_STATE_CLOSING);
               session_queue_push(mgr->evicted_sess, unused_sess);
               session_manager_update_session_timer(mgr, unused_sess, session_to_closed, mgr->timeout_toclosed);
          }

          sess = session_pool_alloc(mgr->sess_pool);
          assert(sess != NULL);

          // return  0: success
          // return -1: tcp not syn packet, discard
          if (handle_new_session(mgr, &key, sess, pkt) == 0)
          {
               session_table_add_session(mgr->sess_table, &key, sess);
          }
          else
          {
               session_pool_free(mgr->sess_pool, sess);
               return NULL;
          }
     }
     else
     {
          handle_old_session(mgr, &key, sess, pkt);
     }

     return sess;
}

struct session *session_manager_expire(struct session_manager *mgr)
{
     SESSION_LOG_DEBUG("current timestamp: %lu s", timestamp_get_sec());
     struct session *sess = session_timer_expire_session(mgr->sess_timer, timestamp_get_msec());
     if (sess == NULL)
     {
          return NULL;
     }

     session_run_expirecb(sess);

     if (session_get_state(sess) == SESSION_STATE_CLOSED)
     {
          return NULL;
     }

     return sess;
}

struct session *session_manager_evicte(struct session_manager *mgr)
{
     return session_queue_pop(mgr->evicted_sess);
}

uint64_t session_manager_get_sessions(struct session_manager *mgr, enum session_type type, enum session_state state)
{
     if (type == SESSION_TYPE_TCP)
     {
          switch (state)
          {
          case SESSION_STATE_OPENING:
               return mgr->tcp_opening_sess_num;
          case SESSION_STATE_ACTIVE:
               return mgr->tcp_active_sess_num;
          case SESSION_STATE_CLOSING:
               return mgr->tcp_closing_sess_num;
          default:
               return 0;
          }
     }

     if (type == SESSION_TYPE_UDP)
     {
          switch (state)
          {
          case SESSION_STATE_OPENING:
               return mgr->udp_opening_sess_num;
          case SESSION_STATE_ACTIVE:
               return mgr->udp_active_sess_num;
          case SESSION_STATE_CLOSING:
               return mgr->udp_closing_sess_num;
          default:
               return 0;
          }
     }

     return 0;
}