#include <fcntl.h>
#include <errno.h>
#include <unistd.h>
#include <stdarg.h>
#include <assert.h>
#include <limits.h>
#include <pthread.h>

#include "packet_dump.h"
#include "utils_internal.h"
#include "session_internal.h"

#include "stellar/log.h"
#include "stellar/module.h"
#include "stellar/session.h"

#pragma GCC diagnostic ignored "-Wunused-parameter"

struct session_debugger
{
    struct logger *logger;
    pthread_spinlock_t lock; // for hexdump thread safe
    struct session_manager *sess_mgr;
    int fd;
    int sess_exdata_idx;
};

struct session_debugger_exdata
{
    struct session_debugger *dbg;
    struct session *sess;

    // data packet
    uint64_t c2s_rx_data_pkts;
    uint64_t s2c_rx_data_pkts;

    uint64_t c2s_rx_data_bytes;
    uint64_t s2c_rx_data_bytes;

    // control packet
    uint64_t c2s_rx_ctrl_pkts;
    uint64_t s2c_rx_ctrl_pkts;

    uint64_t c2s_rx_ctrl_bytes;
    uint64_t s2c_rx_ctrl_bytes;

    // TCP segment
    uint64_t c2s_rx_tcp_seg;
    uint64_t s2c_rx_tcp_seg;

    uint64_t c2s_rx_tcp_bytes;
    uint64_t s2c_rx_tcp_bytes;

    // UDP payload
    uint64_t c2s_rx_udp_payload;
    uint64_t s2c_rx_udp_payload;

    uint64_t c2s_rx_udp_bytes;
    uint64_t s2c_rx_udp_bytes;

    // hexdump TCP segment
    int c2s_tcp_seg_hexdump_fd;
    int s2c_tcp_seg_hexdump_fd;

    // hexdump UDP payload
    int c2s_udp_payload_hexdump_fd;
    int s2c_udp_payload_hexdump_fd;
};

static void session_debugger_log(int fd, const char *fmt, ...)
{
    static const char *module = "session debugger";
    static unsigned char weekday_str[7][4] = {"Sun", "Mon", "Tue", "Wed", "Thu", "Fri", "Sat"};
    static unsigned char month_str[12][4] = {"Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"};

    int nwrite;
    char buf[PATH_MAX * 2] = {0};
    char *p = buf;
    char *end = buf + sizeof(buf);
    va_list arg;
    struct tm local;

    time_t t;
    time(&t);
    localtime_r(&t, &local);

    p += snprintf(p, end - p, "%s %s %d %02d:%02d:%02d %d ",
                  weekday_str[local.tm_wday], month_str[local.tm_mon], local.tm_mday, local.tm_hour, local.tm_min, local.tm_sec, local.tm_year + 1900);
    p += snprintf(p, end - p, "%lu ", pthread_self());
    p += snprintf(p, end - p, "(%s), ", module);
    va_start(arg, fmt);
    p += vsnprintf(p, end - p, fmt, arg);
    va_end(arg);
    p += snprintf(p, end - p, "\n");

    do
    {
        nwrite = write(fd, buf, p - buf);
    } while (nwrite == -1 && errno == EINTR);
}

static struct session_debugger_exdata *session_debugger_exdata_new(struct session_debugger *dbg, struct session *sess)
{
    char buff[PATH_MAX];
    struct session_debugger_exdata *exdata = (struct session_debugger_exdata *)calloc(1, sizeof(struct session_debugger_exdata));

    exdata->dbg = dbg;
    exdata->sess = sess;
    if (session_get_type(sess) == SESSION_TYPE_TCP)
    {
        memset(buff, 0, sizeof(buff));
        sprintf(buff, "./log/session_debugger.TCP_%s_C2S.hexdump", session_get_readable_addr(sess));
        exdata->c2s_tcp_seg_hexdump_fd = open(buff, O_WRONLY | O_APPEND | O_CREAT, 0644);

        memset(buff, 0, sizeof(buff));
        sprintf(buff, "./log/session_debugger.TCP_%s_S2C.hexdump", session_get_readable_addr(sess));
        exdata->s2c_tcp_seg_hexdump_fd = open(buff, O_WRONLY | O_APPEND | O_CREAT, 0644);
    }

    if (session_get_type(sess) == SESSION_TYPE_UDP)
    {
        memset(buff, 0, sizeof(buff));
        sprintf(buff, "./log/session_debugger.UDP_%s_C2S.hexdump", session_get_readable_addr(sess));
        exdata->c2s_udp_payload_hexdump_fd = open(buff, O_WRONLY | O_APPEND | O_CREAT, 0644);

        memset(buff, 0, sizeof(buff));
        sprintf(buff, "./log/session_debugger.UDP_%s_S2C.hexdump", session_get_readable_addr(sess));
        exdata->s2c_udp_payload_hexdump_fd = open(buff, O_WRONLY | O_APPEND | O_CREAT, 0644);
    }

    memset(buff, 0, sizeof(buff));
    session_to_str(sess, 1, buff, sizeof(buff) - 1);
    session_debugger_log(dbg->fd, "sess new: %s", buff);

    return exdata;
}

static void session_debugger_exdata_free(struct session_debugger_exdata *exdata)
{
    if (exdata)
    {
        if (exdata->c2s_tcp_seg_hexdump_fd > 0)
        {
            close(exdata->c2s_tcp_seg_hexdump_fd);
        }
        if (exdata->s2c_tcp_seg_hexdump_fd > 0)
        {
            close(exdata->s2c_tcp_seg_hexdump_fd);
        }
        if (exdata->c2s_udp_payload_hexdump_fd > 0)
        {
            close(exdata->c2s_udp_payload_hexdump_fd);
        }
        if (exdata->s2c_udp_payload_hexdump_fd > 0)
        {
            close(exdata->s2c_udp_payload_hexdump_fd);
        }

        free(exdata);
    }
}

static void session_debugger_exdata_free_cb(int idx, void *ex_ptr, void *arg)
{
    __attribute__((unused)) struct session_debugger *dbg = (struct session_debugger *)arg;
    assert(idx == dbg->sess_exdata_idx);

    session_debugger_exdata_free((struct session_debugger_exdata *)ex_ptr);
}

static void session_debugger_on_session_new(struct session_debugger *dbg, struct session *sess)
{
    struct session_debugger_exdata *exdata = session_debugger_exdata_new(dbg, sess);
    session_set_exdata(sess, dbg->sess_exdata_idx, exdata);
}

static void session_debugger_on_session_free(struct session_debugger *dbg, struct session *sess)
{
    struct session_debugger_exdata *exdata = (struct session_debugger_exdata *)session_get_exdata(sess, dbg->sess_exdata_idx);

    char buff[PATH_MAX] = {0};
    session_to_str(exdata->sess, 0, buff, sizeof(buff) - 1);
    session_debugger_log(exdata->dbg->fd, "session free: %s", buff);

    snprintf(buff, sizeof(buff),
             "==========================================================\n"
             "C2S Data Packets    : %6lu | C2S Data Bytes    : %6lu\n"
             "S2C Data Packets    : %6lu | S2C Data Bytes    : %6lu\n"
             "----------------------------------------------------------\n"
             "C2S Control Packets : %6lu | C2S Control Bytes : %6lu\n"
             "S2C Control Packets : %6lu | S2C Control Bytes : %6lu\n"
             "----------------------------------------------------------\n"
             "C2S TCP Segments    : %6lu | C2S TCP Bytes     : %6lu\n"
             "S2C TCP Segments    : %6lu | S2C TCP Bytes     : %6lu\n"
             "----------------------------------------------------------\n"
             "C2S UDP Payload     : %6lu | C2S UDP Bytes     : %6lu\n"
             "S2C UDP Payload     : %6lu | S2C UDP Bytes     : %6lu",
             exdata->c2s_rx_data_pkts, exdata->c2s_rx_data_bytes,
             exdata->s2c_rx_data_pkts, exdata->s2c_rx_data_bytes,
             exdata->c2s_rx_ctrl_pkts, exdata->c2s_rx_ctrl_bytes,
             exdata->s2c_rx_ctrl_pkts, exdata->s2c_rx_ctrl_bytes,
             exdata->c2s_rx_tcp_seg, exdata->c2s_rx_tcp_bytes,
             exdata->s2c_rx_tcp_seg, exdata->s2c_rx_tcp_bytes,
             exdata->c2s_rx_udp_payload, exdata->c2s_rx_udp_bytes,
             exdata->s2c_rx_udp_payload, exdata->s2c_rx_udp_bytes);
    session_debugger_log(exdata->dbg->fd, "session %lu %s statistics:\n%s", session_get_id(exdata->sess), session_get_readable_addr(exdata->sess), buff);
}

static void session_debugger_on_tcp_payload(struct session_debugger *dbg, struct session *sess, const char *tcp_payload, uint32_t tcp_payload_len)
{
    enum flow_type flow = session_get_flow_type(sess);
    struct session_debugger_exdata *exdata = (struct session_debugger_exdata *)session_get_exdata(sess, dbg->sess_exdata_idx);

    char buff[PATH_MAX];
    memset(buff, 0, sizeof(buff));
    session_to_str(sess, 1, buff, sizeof(buff) - 1);
    session_debugger_log(dbg->fd, "on TCP stream: %s", buff);

    pthread_spin_lock(&dbg->lock);
    if (flow == FLOW_TYPE_C2S)
    {
        session_debugger_log(dbg->fd, "rx C2S TCP segment: len: %d, data: %p", tcp_payload_len, tcp_payload);
        hexdump_to_fd(dbg->fd, exdata->c2s_rx_tcp_bytes, tcp_payload, tcp_payload_len);
        hexdump_to_fd(exdata->c2s_tcp_seg_hexdump_fd, exdata->c2s_rx_tcp_bytes, tcp_payload, tcp_payload_len);

        exdata->c2s_rx_tcp_seg++;
        exdata->c2s_rx_tcp_bytes += tcp_payload_len;
    }
    else
    {
        session_debugger_log(dbg->fd, "rx S2C TCP segment: len: %d, data: %p", tcp_payload_len, tcp_payload);
        hexdump_to_fd(dbg->fd, exdata->s2c_rx_tcp_bytes, tcp_payload, tcp_payload_len);
        hexdump_to_fd(exdata->s2c_tcp_seg_hexdump_fd, exdata->s2c_rx_tcp_bytes, tcp_payload, tcp_payload_len);

        exdata->s2c_rx_tcp_seg++;
        exdata->s2c_rx_tcp_bytes += tcp_payload_len;
    }
    pthread_spin_unlock(&dbg->lock);
}

static void session_debugger_on_udp_payload(struct session_debugger *dbg, struct session *sess, struct packet *pkt)
{
    enum flow_type flow = session_get_flow_type(sess);
    struct session_debugger_exdata *exdata = (struct session_debugger_exdata *)session_get_exdata(sess, dbg->sess_exdata_idx);

    char buff[PATH_MAX];
    const char *udp_payload = packet_get_payload_data(pkt);
    uint32_t udp_payload_len = packet_get_payload_len(pkt);
    if (udp_payload_len == 0)
    {
        return;
    }

    memset(buff, 0, sizeof(buff));
    session_to_str(sess, 1, buff, sizeof(buff) - 1);
    session_debugger_log(dbg->fd, "on UDP payload: %s", buff);

    pthread_spin_lock(&dbg->lock);
    if (flow == FLOW_TYPE_C2S)
    {
        session_debugger_log(dbg->fd, "rx C2S UDP payload: len: %d, data: %p", udp_payload_len, udp_payload);
        hexdump_to_fd(dbg->fd, exdata->c2s_rx_udp_bytes, udp_payload, udp_payload_len);
        hexdump_to_fd(exdata->c2s_udp_payload_hexdump_fd, exdata->c2s_rx_udp_bytes, udp_payload, udp_payload_len);
        exdata->c2s_rx_udp_payload++;
        exdata->c2s_rx_udp_bytes += udp_payload_len;
    }
    else
    {
        session_debugger_log(dbg->fd, "rx S2C UDP payload: len: %d, data: %p", udp_payload_len, udp_payload);
        hexdump_to_fd(dbg->fd, exdata->s2c_rx_udp_bytes, udp_payload, udp_payload_len);
        hexdump_to_fd(exdata->s2c_udp_payload_hexdump_fd, exdata->s2c_rx_udp_bytes, udp_payload, udp_payload_len);
        exdata->s2c_rx_udp_payload++;
        exdata->s2c_rx_udp_bytes += udp_payload_len;
    }
    pthread_spin_unlock(&dbg->lock);
}

void session_debugger_on_packet_forward(struct packet *pkt, struct module *mod)
{
    struct session_debugger *dbg = module_get_ctx(mod);
    assert(dbg);
    struct session_manager *sess_mgr = dbg->sess_mgr;

    struct session *sess = packet_exdata_to_session(sess_mgr, pkt);
    if (sess == NULL)
    {
        assert(0);
        return;
    }

    uint64_t pkt_tag_key = 0;
    uint64_t pkt_tag_val = 0;
    packet_tag_get(pkt, &pkt_tag_key, &pkt_tag_val);

    if (!(pkt_tag_key & PKT_TAG_KEY_SESS))
    {
        assert(0);
        return;
    }

    if (pkt_tag_val & PKT_TAG_VAL_SESS_FREE)
    {
        session_debugger_on_session_free(dbg, sess);
        return;
    }

    if (pkt_tag_val & PKT_TAG_VAL_SESS_NEW)
    {
        session_debugger_on_session_new(dbg, sess);
    }

    struct session_debugger_exdata *exdata = (struct session_debugger_exdata *)session_get_exdata(sess, dbg->sess_exdata_idx);
    enum flow_type flow = session_get_flow_type(sess);
    assert(flow == FLOW_TYPE_C2S || flow == FLOW_TYPE_S2C);

    int is_pseudo = (packet_get_type(pkt) == PACKET_TYPE_PSEUDO);
    if (flow == FLOW_TYPE_C2S)
    {
        if (is_pseudo)
        {
            exdata->c2s_rx_ctrl_pkts++;
            exdata->c2s_rx_ctrl_bytes += packet_get_raw_len(pkt);
        }
        else
        {
            exdata->c2s_rx_data_pkts++;
            exdata->c2s_rx_data_bytes += packet_get_raw_len(pkt);
        }
    }
    else
    {
        if (is_pseudo)
        {
            exdata->s2c_rx_ctrl_pkts++;
            exdata->s2c_rx_ctrl_bytes += packet_get_raw_len(pkt);
        }
        else
        {
            exdata->s2c_rx_data_pkts++;
            exdata->s2c_rx_data_bytes += packet_get_raw_len(pkt);
        }
    }

    char buff[PATH_MAX];
    memset(buff, 0, sizeof(buff));
    session_to_str(sess, 1, buff, sizeof(buff) - 1);
    session_debugger_log(dbg->fd, "on %s %s packet: %s", session_type_to_str(session_get_type(sess)), (is_pseudo ? "pseudo" : "data"), buff);

    memset(buff, 0, sizeof(buff));
    packet_dump_str(pkt, buff, sizeof(buff) - 1);
    session_debugger_log(dbg->fd, "rx %s %s packet\n%s", session_type_to_str(session_get_type(sess)), (is_pseudo ? "pseudo" : "data"), buff);

    pthread_spin_lock(&dbg->lock);
    packet_dump_hex(pkt, dbg->fd);
    pthread_spin_unlock(&dbg->lock);

    if (session_get_type(sess) == SESSION_TYPE_TCP)
    {
        struct tcp_segment *seg = packet_exdata_to_tcp_segment(sess_mgr, pkt);
        while (seg)
        {
            session_debugger_on_tcp_payload(dbg, sess, seg->data, seg->len);
            seg = seg->next;
        }
    }
    else
    {
        session_debugger_on_udp_payload(dbg, sess, pkt);
    }
}

static void session_debugger_free(struct session_debugger *dbg)
{
    if (dbg)
    {
        if (dbg->fd > 0)
        {
            close(dbg->fd);
        }
        pthread_spin_destroy(&dbg->lock);
        free(dbg);
    }
}

static struct session_debugger *session_debugger_new(struct session_manager *sess_mgr, struct logger *logger)
{
    struct session_debugger *dbg = (struct session_debugger *)calloc(1, sizeof(struct session_debugger));
    if (dbg == NULL)
    {
        session_debugger_log(STDERR_FILENO, "new failed\n");
        return NULL;
    }

    pthread_spin_init(&dbg->lock, PTHREAD_PROCESS_PRIVATE);

    dbg->logger = logger;
    dbg->sess_mgr = sess_mgr;
    dbg->fd = open("./log/session_debugger.log", O_WRONLY | O_APPEND | O_CREAT, 0644);
    if (dbg->fd == -1)
    {
        session_debugger_log(STDERR_FILENO, "open log file failed: %s\n", strerror(errno));
        goto error_out;
    }

    dbg->sess_exdata_idx = session_manager_new_session_exdata_index(dbg->sess_mgr, "session_debugger_exdata", session_debugger_exdata_free_cb, dbg);
    if (dbg->sess_exdata_idx == -1)
    {
        session_debugger_log(STDERR_FILENO, "new session exdata index failed\n");
        goto error_out;
    }

    return dbg;

error_out:
    session_debugger_free(dbg);
    return NULL;
}

/******************************************************************************
 * Plugin API
 ******************************************************************************/

struct module *session_debugger_on_init(struct module_manager *mod_mgr)
{
    struct session_manager *sess_mgr = module_to_session_manager(module_manager_get_module(mod_mgr, SESSION_MANAGER_MODULE_NAME));
    assert(sess_mgr);
    struct logger *logger = module_manager_get_logger(mod_mgr);
    assert(logger);
    const char *toml = module_manager_get_toml_path(mod_mgr);
    assert(toml);

    uint64_t enable = 0;
    if (load_toml_integer_config(toml, "session_debugger.enable", &enable, 0, 1) != 0)
    {
        return NULL;
    }
    if (enable == 0)
    {
        return NULL;
    }

    struct session_debugger *dbg = session_debugger_new(sess_mgr, logger);
    if (dbg == NULL)
    {
        return NULL;
    }

    struct module *dbg_mod = module_new(SESSION_DEBUGGER_MODULE_NAME, dbg);
    if (dbg_mod == NULL)
    {
        session_debugger_free(dbg);
        return NULL;
    }

    STELLAR_LOG_FATAL(dbg->logger, "session debugger", "session_debugger init")

    return dbg_mod;
}

void session_debugger_on_exit(struct module_manager *mod_mgr, struct module *mod)
{
    if (mod)
    {
        struct session_debugger *dbg = module_get_ctx(mod);
        if (dbg)
        {
            STELLAR_LOG_FATAL(dbg->logger, "session debugger", "session_debugger exit")
            session_debugger_free(dbg);
        }
        module_free(mod);
    }
}