#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 *sess_queue; // ready session queue

     session_event_cb event_cb;
     void *arg;

     // timeout config
     uint64_t packet_timeout_ms;
     uint64_t closing_timeout_ms;

     // 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 event and timer
 ******************************************************************************/

void session_manager_trigger_session_event(struct session_manager *mgr, struct session *sess, uint32_t event)
{
     session_push_event(sess, event);
     session_queue_push(mgr->sess_queue, sess);
}

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 closing_expire_callback(struct session *sess, void *arg)
{
     SESSION_MANAGER_LOG_DEBUG("session %lu closing expire, free session", session_get_id(sess));
     struct session_manager *mgr = (struct session_manager *)arg;
     assert(mgr != NULL);

     uint32_t event;
     while (session_pop_event(sess, &event))
     {
     }

     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 packet_expire_callback(struct session *sess, void *arg)
{
     SESSION_MANAGER_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_trigger_session_event(mgr, sess, SESSION_EVENT_CLOSING);
     session_manager_update_session_timer(mgr, sess, closing_expire_callback, mgr->closing_timeout_ms);
}

/******************************************************************************
 * 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 tcphdr *hdr = (const struct tcphdr *)packet_get_tcp_hdr_ptr(pkt);
     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_trigger_session_event(mgr, sess, SESSION_EVENT_OPENING);
     session_manager_update_session_timer(mgr, sess, packet_expire_callback, mgr->packet_timeout_ms);

     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_trigger_session_event(mgr, sess, SESSION_EVENT_OPENING);
     session_manager_trigger_session_event(mgr, sess, SESSION_EVENT_PACKET);
     session_manager_update_session_timer(mgr, sess, packet_expire_callback, mgr->packet_timeout_ms);

     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_trigger_session_event(mgr, sess, SESSION_EVENT_CLOSING);
          session_manager_update_session_timer(mgr, sess, closing_expire_callback, mgr->closing_timeout_ms);
          return;
     }

     if (tcp_need_active(state))
     {
          update_counter_on_active(mgr, sess);
          session_set_state(sess, SESSION_STATE_ACTIVE);
          session_manager_trigger_session_event(mgr, sess, SESSION_EVENT_PACKET);
          session_manager_update_session_timer(mgr, sess, packet_expire_callback, mgr->packet_timeout_ms);
          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_trigger_session_event(mgr, sess, SESSION_EVENT_PACKET);
     session_manager_update_session_timer(mgr, sess, packet_expire_callback, mgr->packet_timeout_ms);
}

// 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->sess_queue = session_queue_create();
     if (mgr->sess_queue == NULL)
     {
          goto error;
     }

     mgr->closing_timeout_ms = 2 * 1000;
     mgr->packet_timeout_ms = 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->sess_queue);
          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_session_eventcb(struct session_manager *mgr, session_event_cb cb, void *arg)
{
     mgr->event_cb = cb;
     mgr->arg = arg;
}

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

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

// return NULL: discard
//      * tuple6 not find
//      * tcp first packet not syn
struct session *session_manager_find_session(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_manager_trigger_session_event(mgr, unused_sess, SESSION_EVENT_CLOSING);
               session_manager_update_session_timer(mgr, unused_sess, closing_expire_callback, mgr->closing_timeout_ms);
          }

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

void session_manager_dispatch(struct session_manager *mgr)
{
     uint32_t event;
     struct session *sess;

     void *cb_arg = mgr->arg;
     session_event_cb event_cb = mgr->event_cb;

     SESSION_MANAGER_LOG_DEBUG("current timestamp: %lu s", timestamp_get_sec());

     // limit expire session number
     for (int i = 0; i < 100; i++)
     {
          sess = session_timer_expire_session(mgr->sess_timer, timestamp_get_msec());
          if (sess == NULL)
          {
               break;
          }

          session_run_expirecb(sess);
     }

     while (1)
     {
          // get session from ready queue
          sess = session_queue_pop(mgr->sess_queue);
          if (sess == NULL)
          {
               break;
          }

          while (1)
          {
               // get event from session
               if (session_pop_event(sess, &event) == false)
               {
                    break;
               }

               SESSION_MANAGER_LOG_DEBUG("handle \"%s\" event on session %lu", session_event_tostring((enum session_event)event), session_get_id(sess));
               if (event_cb)
               {
                    event_cb(sess, event, cb_arg);
               }
          }
          session_set0_cur_pkt(sess, NULL);
          session_set_cur_dir(sess, SESSION_DIR_NONE);
     };
}

uint64_t session_manager_get_tcp_opening_sess_num(struct session_manager *mgr)
{
     return mgr->tcp_opening_sess_num;
}

uint64_t session_manager_get_tcp_closing_sess_num(struct session_manager *mgr)
{
     return mgr->tcp_closing_sess_num;
}

uint64_t session_manager_get_tcp_active_sess_num(struct session_manager *mgr)
{
     return mgr->tcp_active_sess_num;
}

uint64_t session_manager_get_udp_opening_sess_num(struct session_manager *mgr)
{
     return mgr->udp_opening_sess_num;
}

uint64_t session_manager_get_udp_closing_sess_num(struct session_manager *mgr)
{
     return mgr->udp_closing_sess_num;
}

uint64_t session_manager_get_udp_active_sess_num(struct session_manager *mgr)
{
     return mgr->udp_active_sess_num;
}