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

#include "utils.h"
#include "packet_dump.h"
#include "session_internal.h"

#include "stellar/log.h"
#include "stellar/module_manager.h"
#include "stellar/session_manager.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;

    uint64_t c2s_rx_pkts;
    uint64_t s2c_rx_pkts;

    uint64_t c2s_rx_bytes;
    uint64_t s2c_rx_bytes;

    uint64_t c2s_rx_tcp_seg;
    uint64_t s2c_rx_tcp_seg;

    uint64_t c2s_rx_tcp_bytes;
    uint64_t s2c_rx_tcp_bytes;

    int c2s_tcp_seg_hexdump_fd;
    int s2c_tcp_seg_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 args;
    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(args, fmt);
    p += vsnprintf(p, end - p, fmt, args);
    va_end(args);
    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.%s_c2s.hexdump", session_get0_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.%s_s2c.hexdump", session_get0_readable_addr(sess));
        exdata->s2c_tcp_seg_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)
    {
        char buff[PATH_MAX] = {0};
        session_to_str(exdata->sess, 0, buff, sizeof(buff) - 1);
        session_debugger_log(exdata->dbg->fd, "sess free: %s", buff);
        session_debugger_log(exdata->dbg->fd, "session %lu %s stat:\n"
                                              "C2S rx      packets: %6lu, C2S rx     bytes: %6lu\n"
                                              "S2C rx      packets: %6lu, S2C rx     bytes: %6lu\n"
                                              "C2S rx TCP segments: %6lu, C2S rx TCP bytes: %6lu\n"
                                              "S2C rx TCP segments: %6lu, S2C rx TCP bytes: %6lu\n",
                             session_get_id(exdata->sess), session_get0_readable_addr(exdata->sess),
                             exdata->c2s_rx_pkts, exdata->c2s_rx_bytes,
                             exdata->s2c_rx_pkts, exdata->s2c_rx_bytes,
                             exdata->c2s_rx_tcp_seg, exdata->c2s_rx_tcp_bytes,
                             exdata->s2c_rx_tcp_seg, exdata->s2c_rx_tcp_bytes);

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

        free(exdata);
    }
}

static void session_debugger_exdata_free_callback(int idx, void *ex_ptr, void *arg)
{
    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 on_sess_msg(struct session *sess, struct packet *pkt, void *args)
{
    struct session_debugger *dbg = (struct session_debugger *)args;

    char buff[PATH_MAX];
    struct session_debugger_exdata *exdata = (struct session_debugger_exdata *)session_get_exdata(sess, dbg->sess_exdata_idx);
    if (exdata == NULL)
    {
        exdata = session_debugger_exdata_new(dbg, sess);
        session_set_exdata(sess, dbg->sess_exdata_idx, exdata);
    }

    if (session_get_flow_type(sess) == FLOW_TYPE_C2S)
    {
        exdata->c2s_rx_pkts++;
        exdata->c2s_rx_bytes += packet_get_raw_len(pkt);
    }
    else
    {
        exdata->s2c_rx_pkts++;
        exdata->s2c_rx_bytes += packet_get_raw_len(pkt);
    }

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

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

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

static void on_tcp_stream_msg(struct session *sess, const char *tcp_payload, uint32_t tcp_payload_len, void *args)
{
    struct session_debugger *dbg = (struct session_debugger *)args;

    char buff[PATH_MAX];
    struct session_debugger_exdata *exdata = (struct session_debugger_exdata *)session_get_exdata(sess, dbg->sess_exdata_idx);

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

    pthread_spin_lock(&dbg->lock);
    if (session_get_flow_type(sess) == 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_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, "DEBUG_MODULE_SESS_EXDATA", session_debugger_exdata_free_callback, dbg);
    if (dbg->sess_exdata_idx == -1)
    {
        session_debugger_log(STDERR_FILENO, "new session exdata index failed\n");
        goto error_out;
    }

    if (session_manager_subscribe_tcp(sess_mgr, on_sess_msg, dbg) == -1)
    {
        session_debugger_log(STDERR_FILENO, "subscribe tcp failed\n");
        goto error_out;
    }
    if (session_manager_subscribe_udp(sess_mgr, on_sess_msg, dbg) == -1)
    {
        session_debugger_log(STDERR_FILENO, "subscribe udp failed\n");
        goto error_out;
    }
    if (session_manager_subscribe_control_packet(sess_mgr, on_sess_msg, dbg) == -1)
    {
        session_debugger_log(STDERR_FILENO, "subscribe control packet failed\n");
        goto error_out;
    }
    if (session_manager_subscribe_tcp_stream(sess_mgr, on_tcp_stream_msg, dbg) == -1)
    {
        session_debugger_log(STDERR_FILENO, "subscribe tcp stream failed\n");
        goto error_out;
    }

    return dbg;

error_out:
    session_debugger_free(dbg);
    return NULL;
}

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

struct stellar_module *session_debugger_on_init(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);
    assert(sess_mgr_mod);
    struct session_manager *sess_mgr = stellar_module_get_ctx(sess_mgr_mod);
    assert(sess_mgr);
    struct logger *logger = stellar_module_manager_get_logger(mod_mgr);
    assert(logger);

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

    struct stellar_module *dbg_mod = stellar_module_new("session_debugger", NULL);
    if (dbg_mod == NULL)
    {
        session_debugger_free(dbg);
        return NULL;
    }
    stellar_module_set_ctx(dbg_mod, dbg);

    STELLAR_LOG_FATAL(dbg->logger, "session_debugger", "session_debugger initialized")

    return dbg_mod;
}

void session_debugger_on_exit(struct stellar_module_manager *mod_mgr, struct stellar_module *mod)
{
    if (mod)
    {
        struct session_debugger *dbg = stellar_module_get_ctx(mod);
        if (dbg)
        {
            STELLAR_LOG_FATAL(dbg->logger, "session_debugger", "session_debugger exited")
            session_debugger_free(dbg);
        }
        stellar_module_free(mod);
    }
}