#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

#include "log_private.h"
#include "stellar_stat.h"
#include "fieldstat/fieldstat_easy.h"
#include "fieldstat/fieldstat_exporter.h"

#define STAT_LOG_ERROR(format, ...) STELLAR_LOG_ERROR(__thread_local_logger, "stat", format, ##__VA_ARGS__)
#define STAT_LOG_INFO(format, ...) STELLAR_LOG_INFO(__thread_local_logger, "stat", format, ##__VA_ARGS__)

#define IS_FREE 0
#define IS_BUSY 0xf

enum stat_type
{
    // device packet
    STAT_TYPE_PKTS_RX,
    STAT_TYPE_BYTES_RX,
    STAT_TYPE_PKTS_TX,
    STAT_TYPE_BYTES_TX,

    // keep-alive packet
    STAT_TYPE_KEEP_ALIVE_PKTS,
    STAT_TYPE_KEEP_ALIVE_BYTES,

    // raw packet
    STAT_TYPE_RAW_PKTS_RX,
    STAT_TYPE_RAW_BYTES_RX,
    STAT_TYPE_RAW_PKTS_TX,
    STAT_TYPE_RAW_BYTES_TX,

    // drop packet
    STAT_TYPE_PKTS_DROPPED,
    STAT_TYPE_BYTES_DROPPED,

    // inject packet
    STAT_TYPE_PKTS_INJECTED,
    STAT_TYPE_BYTES_INJECTED,

    // ctrl packet
    STAT_TYPE_CTRL_PKTS_RX,
    STAT_TYPE_CTRL_BYTES_RX,
    STAT_TYPE_CTRL_PKTS_TX,
    STAT_TYPE_CTRL_BYTES_TX,

    // ipv4 reassembly
    STAT_TYPE_IP4_DEFRAGS_EXPECTED,
    STAT_TYPE_IP4_DEFRAGS_SUCCEED,
    STAT_TYPE_IP4_DEFRAGS_FAILED_TIMEOUT,
    STAT_TYPE_IP4_DEFRAGS_FAILED_INVALID_LENGTH,
    STAT_TYPE_IP4_DEFRAGS_FAILED_OVERLAP,
    STAT_TYPE_IP4_DEFRAGS_FAILED_TOO_MANY_FRAG,

    STAT_TYPE_IP4_FRAGS,
    STAT_TYPE_IP4_FRAGS_FREED,
    STAT_TYPE_IP4_FRAGS_BUFFERED,
    STAT_TYPE_IP4_FRAGS_BYPASS_NO_BUFFER,
    STAT_TYPE_IP4_FRAGS_BYPASS_DUP_FIST_FRAG,
    STAT_TYPE_IP4_FRAGS_BYPASS_DUP_LAST_FRAG,

    // ipv6 reassembly
    STAT_TYPE_IP6_DEFRAGS_EXPECTED,
    STAT_TYPE_IP6_DEFRAGS_SUCCEED,
    STAT_TYPE_IP6_DEFRAGS_FAILED_TIMEOUT,
    STAT_TYPE_IP6_DEFRAGS_FAILED_INVALID_LENGTH,
    STAT_TYPE_IP6_DEFRAGS_FAILED_OVERLAP,
    STAT_TYPE_IP6_DEFRAGS_FAILED_TOO_MANY_FRAG,

    STAT_TYPE_IP6_FRAGS,
    STAT_TYPE_IP6_FRAGS_FREED,
    STAT_TYPE_IP6_FRAGS_BUFFERED,
    STAT_TYPE_IP6_FRAGS_BYPASS_NO_BUFFER,
    STAT_TYPE_IP6_FRAGS_BYPASS_DUP_FIST_FRAG,
    STAT_TYPE_IP6_FRAGS_BYPASS_DUP_LAST_FRAG,

    // TCP session
    STAT_TYPE_HISTORY_TCP_SESSIONS,
    STAT_TYPE_TCP_SESS_USED,
    STAT_TYPE_TCP_SESS_OPENING,
    STAT_TYPE_TCP_SESS_ACTIVE,
    STAT_TYPE_TCP_SESS_CLOSING,
    STAT_TYPE_TCP_SESS_DISCARD,
    STAT_TYPE_TCP_SESS_CLOSED,

    // UDP session
    STAT_TYPE_HISTORY_UDP_SESSIONS,
    STAT_TYPE_UDP_SESS_USED,
    STAT_TYPE_UDP_SESS_OPENING,
    STAT_TYPE_UDP_SESS_ACTIVE,
    STAT_TYPE_UDP_SESS_CLOSING,
    STAT_TYPE_UDP_SESS_DISCARD,
    STAT_TYPE_UDP_SESS_CLOSED,

    // Evicted session
    STAT_TYPE_TCP_SESS_EVICTED,
    STAT_TYPE_UDP_SESS_EVICTED,

    // Packet
    STAT_TYPE_UDP_PKTS_BYPASS_TABLE_FULL,
    STAT_TYPE_TCP_PKTS_BYPASS_TABLE_FULL,
    STAT_TYPE_TCP_PKTS_BYPASS_SESSION_NOT_FOUND,
    STAT_TYPE_TCP_PKTS_BYPASS_DUPLICATED,
    STAT_TYPE_UDP_PKTS_BYPASS_DUPLICATED,
    STAT_TYPE_UDP_PKTS_BYPASS_SESSION_EVICTED,

    // TCP segments
    STAT_TYPE_TCP_SEGS_INPUT,
    STAT_TYPE_TCP_SEGS_CONSUMED,
    STAT_TYPE_TCP_SEGS_TIMEOUT,
    STAT_TYPE_TCP_SEGS_RETRANSMITED,
    STAT_TYPE_TCP_SEGS_OVERLAPPED,
    STAT_TYPE_TCP_SEGS_OMITTED_TOO_MANY,
    STAT_TYPE_TCP_SEGS_INORDER,
    STAT_TYPE_TCP_SEGS_REORDERED,
    STAT_TYPE_TCP_SEGS_BUFFERED,
    STAT_TYPE_TCP_SEGS_FREED,

    // end
    STAT_TYPE_MAX,
};

const char *name[] = {
    // device packet
    "pkts_rx",
    "bytes_rx",
    "pkts_tx",
    "bytes_tx",

    // keep-alive packet
    "keep_alive_pkts",
    "keep_alive_bytes",

    // raw packet
    "raw_pkts_rx",
    "raw_bytes_rx",
    "raw_pkts_tx",
    "raw_bytes_tx",

    // drop packet
    "pkts_dropped",
    "bytes_dropped",

    // inject packet
    "pkts_injected",
    "bytes_injected",

    // ctrl packet
    "ctrl_pkts_rx",
    "ctrl_bytes_rx",
    "ctrl_pkts_tx",
    "ctrl_bytes_tx",

    // ipv4 reassembly
    "ip4_defrags_expected",
    "ip4_defrags_succeed",
    "ip4_defrags_failed_timeout",
    "ip4_defrags_failed_invalid_length",
    "ip4_defrags_failed_overlap",
    "ip4_defrags_failed_too_many_frag",

    "ip4_frags",
    "ip4_frags_freed",
    "ip4_frags_buffered",
    "ip4_frags_bypass_no_buffer",
    "ip4_frags_bypass_dup_fist_frag",
    "ip4_frags_bypass_dup_last_frag",

    // ipv6 reassembly
    "ip6_defrags_expected",
    "ip6_defrags_succeed",
    "ip6_defrags_failed_timeout",
    "ip6_defrags_failed_invalid_length",
    "ip6_defrags_failed_overlap",
    "ip6_defrags_failed_too_many_frag",

    "ip6_frags",
    "ip6_frags_freed",
    "ip6_frags_buffered",
    "ip6_frags_bypass_no_buffer",
    "ip6_frags_bypass_dup_fist_frag",
    "ip6_frags_bypass_dup_last_frag",

    // TCP session
    "history_tcp_sessions",
    "tcp_sess_used",
    "tcp_sess_opening",
    "tcp_sess_active",
    "tcp_sess_closing",
    "tcp_sess_discard",
    "tcp_sess_closed",

    // UDP session
    "history_udp_sessions",
    "udp_sess_used",
    "udp_sess_opening",
    "udp_sess_active",
    "udp_sess_closing",
    "udp_sess_discard",
    "udp_sess_closed",

    // Evicted session
    "tcp_sess_evicted",
    "udp_sess_evicted",

    // Packet
    "udp_pkts_bypass_table_full",
    "tcp_pkts_bypass_table_full",
    "tcp_pkts_bypass_session_not_found",
    "tcp_pkts_bypass_duplicated",
    "udp_pkts_bypass_duplicated",
    "udp_pkts_bypass_session_evicted",

    // TCP segments
    "tcp_segs_input",
    "tcp_segs_consumed",
    "tcp_segs_timeout",
    "tcp_segs_retransmited",
    "tcp_segs_overlapped",
    "tcp_segs_omitted_too_many",
    "tcp_segs_inorder",
    "tcp_segs_reordered",
    "tcp_segs_buffered",
    "tcp_segs_freed",
};

/*
 * This naming convention maintains consistency by using a clear, structured pattern:
 *      •	tcp_sessions_ or ipv6_frags_ as the prefix to indicate the type of data.
 *      •	Action or status (e.g., bypassed, active, dropped) as the middle part.
 *      •	Cause or condition (e.g., full_table, buffer_limit) as the suffix for additional clarity.
 */

struct stellar_stat
{
    uint16_t nr_thread;
    char output_file[2048];
    struct fieldstat_easy *fs;

    int flag[MAX_THREAD_NUM]; // IS_FREE or IS_BUSY
    struct thread_stat thr_stat[MAX_THREAD_NUM];

    uint64_t stat_idx[STAT_TYPE_MAX];
    uint64_t stat_val[STAT_TYPE_MAX];
};

uint64_t get_stat_value_by_idx(const struct thread_stat *thr_stat, size_t idx)
{
    switch (idx)
    {
        // device packet
    case STAT_TYPE_PKTS_RX:
        return thr_stat->packet_io->pkts_rx;
    case STAT_TYPE_BYTES_RX:
        return thr_stat->packet_io->bytes_rx;
    case STAT_TYPE_PKTS_TX:
        return thr_stat->packet_io->pkts_tx;
    case STAT_TYPE_BYTES_TX:
        return thr_stat->packet_io->bytes_tx;

        // keep-alive packet
    case STAT_TYPE_KEEP_ALIVE_PKTS:
        return thr_stat->packet_io->keep_alive_pkts;
    case STAT_TYPE_KEEP_ALIVE_BYTES:
        return thr_stat->packet_io->keep_alive_bytes;

        // raw packet
    case STAT_TYPE_RAW_PKTS_RX:
        return thr_stat->packet_io->raw_pkts_rx;
    case STAT_TYPE_RAW_BYTES_RX:
        return thr_stat->packet_io->raw_bytes_rx;
    case STAT_TYPE_RAW_PKTS_TX:
        return thr_stat->packet_io->raw_pkts_tx;
    case STAT_TYPE_RAW_BYTES_TX:
        return thr_stat->packet_io->raw_bytes_tx;

        // drop packet
    case STAT_TYPE_PKTS_DROPPED:
        return thr_stat->packet_io->pkts_dropped;
    case STAT_TYPE_BYTES_DROPPED:
        return thr_stat->packet_io->bytes_dropped;

        // inject packet
    case STAT_TYPE_PKTS_INJECTED:
        return thr_stat->packet_io->pkts_injected;
    case STAT_TYPE_BYTES_INJECTED:
        return thr_stat->packet_io->bytes_injected;

        // ctrl packet
    case STAT_TYPE_CTRL_PKTS_RX:
        return thr_stat->packet_io->ctrl_pkts_rx;
    case STAT_TYPE_CTRL_BYTES_RX:
        return thr_stat->packet_io->ctrl_bytes_rx;
    case STAT_TYPE_CTRL_PKTS_TX:
        return thr_stat->packet_io->ctrl_pkts_tx;
    case STAT_TYPE_CTRL_BYTES_TX:
        return thr_stat->packet_io->ctrl_bytes_tx;

        // ipv4 reassembly
    case STAT_TYPE_IP4_DEFRAGS_EXPECTED:
        return thr_stat->ip_reassembly->ip4_defrags_expected;
    case STAT_TYPE_IP4_DEFRAGS_SUCCEED:
        return thr_stat->ip_reassembly->ip4_defrags_succeed;
    case STAT_TYPE_IP4_DEFRAGS_FAILED_TIMEOUT:
        return thr_stat->ip_reassembly->ip4_defrags_failed_timeout;
    case STAT_TYPE_IP4_DEFRAGS_FAILED_INVALID_LENGTH:
        return thr_stat->ip_reassembly->ip4_defrags_failed_invalid_length;
    case STAT_TYPE_IP4_DEFRAGS_FAILED_OVERLAP:
        return thr_stat->ip_reassembly->ip4_defrags_failed_overlap;
    case STAT_TYPE_IP4_DEFRAGS_FAILED_TOO_MANY_FRAG:
        return thr_stat->ip_reassembly->ip4_defrags_failed_too_many_frag;
    case STAT_TYPE_IP4_FRAGS:
        return thr_stat->ip_reassembly->ip4_frags;
    case STAT_TYPE_IP4_FRAGS_FREED:
        return thr_stat->ip_reassembly->ip4_frags_freed;
    case STAT_TYPE_IP4_FRAGS_BUFFERED:
        return thr_stat->ip_reassembly->ip4_frags_buffered;
    case STAT_TYPE_IP4_FRAGS_BYPASS_NO_BUFFER:
        return thr_stat->ip_reassembly->ip4_frags_bypass_no_buffer;
    case STAT_TYPE_IP4_FRAGS_BYPASS_DUP_FIST_FRAG:
        return thr_stat->ip_reassembly->ip4_frags_bypass_dup_fist_frag;
    case STAT_TYPE_IP4_FRAGS_BYPASS_DUP_LAST_FRAG:
        return thr_stat->ip_reassembly->ip4_frags_bypass_dup_last_frag;

        // ipv6 reassembly
    case STAT_TYPE_IP6_DEFRAGS_EXPECTED:
        return thr_stat->ip_reassembly->ip6_defrags_expected;
    case STAT_TYPE_IP6_DEFRAGS_SUCCEED:
        return thr_stat->ip_reassembly->ip6_defrags_succeed;
    case STAT_TYPE_IP6_DEFRAGS_FAILED_TIMEOUT:
        return thr_stat->ip_reassembly->ip6_defrags_failed_timeout;
    case STAT_TYPE_IP6_DEFRAGS_FAILED_INVALID_LENGTH:
        return thr_stat->ip_reassembly->ip6_defrags_failed_invalid_length;
    case STAT_TYPE_IP6_DEFRAGS_FAILED_OVERLAP:
        return thr_stat->ip_reassembly->ip6_defrags_failed_overlap;
    case STAT_TYPE_IP6_DEFRAGS_FAILED_TOO_MANY_FRAG:
        return thr_stat->ip_reassembly->ip6_defrags_failed_too_many_frag;
    case STAT_TYPE_IP6_FRAGS:
        return thr_stat->ip_reassembly->ip6_frags;
    case STAT_TYPE_IP6_FRAGS_FREED:
        return thr_stat->ip_reassembly->ip6_frags_freed;
    case STAT_TYPE_IP6_FRAGS_BUFFERED:
        return thr_stat->ip_reassembly->ip6_frags_buffered;
    case STAT_TYPE_IP6_FRAGS_BYPASS_NO_BUFFER:
        return thr_stat->ip_reassembly->ip6_frags_bypass_no_buffer;
    case STAT_TYPE_IP6_FRAGS_BYPASS_DUP_FIST_FRAG:
        return thr_stat->ip_reassembly->ip6_frags_bypass_dup_fist_frag;
    case STAT_TYPE_IP6_FRAGS_BYPASS_DUP_LAST_FRAG:
        return thr_stat->ip_reassembly->ip6_frags_bypass_dup_last_frag;

        // TCP session
    case STAT_TYPE_HISTORY_TCP_SESSIONS:
        return thr_stat->session_mgr->history_tcp_sessions;
    case STAT_TYPE_TCP_SESS_USED:
        return thr_stat->session_mgr->tcp_sess_used;
    case STAT_TYPE_TCP_SESS_OPENING:
        return thr_stat->session_mgr->tcp_sess_opening;
    case STAT_TYPE_TCP_SESS_ACTIVE:
        return thr_stat->session_mgr->tcp_sess_active;
    case STAT_TYPE_TCP_SESS_CLOSING:
        return thr_stat->session_mgr->tcp_sess_closing;
    case STAT_TYPE_TCP_SESS_DISCARD:
        return thr_stat->session_mgr->tcp_sess_discard;
    case STAT_TYPE_TCP_SESS_CLOSED:
        return thr_stat->session_mgr->tcp_sess_closed;

        // UDP session
    case STAT_TYPE_HISTORY_UDP_SESSIONS:
        return thr_stat->session_mgr->history_udp_sessions;
    case STAT_TYPE_UDP_SESS_USED:
        return thr_stat->session_mgr->udp_sess_used;
    case STAT_TYPE_UDP_SESS_OPENING:
        return thr_stat->session_mgr->udp_sess_opening;
    case STAT_TYPE_UDP_SESS_ACTIVE:
        return thr_stat->session_mgr->udp_sess_active;
    case STAT_TYPE_UDP_SESS_CLOSING:
        return thr_stat->session_mgr->udp_sess_closing;
    case STAT_TYPE_UDP_SESS_DISCARD:
        return thr_stat->session_mgr->udp_sess_discard;
    case STAT_TYPE_UDP_SESS_CLOSED:
        return thr_stat->session_mgr->udp_sess_closed;

        // Evicted session
    case STAT_TYPE_TCP_SESS_EVICTED:
        return thr_stat->session_mgr->tcp_sess_evicted;
    case STAT_TYPE_UDP_SESS_EVICTED:
        return thr_stat->session_mgr->udp_sess_evicted;

        // Packet
    case STAT_TYPE_UDP_PKTS_BYPASS_TABLE_FULL:
        return thr_stat->session_mgr->udp_pkts_bypass_table_full;
    case STAT_TYPE_TCP_PKTS_BYPASS_TABLE_FULL:
        return thr_stat->session_mgr->tcp_pkts_bypass_table_full;
    case STAT_TYPE_TCP_PKTS_BYPASS_SESSION_NOT_FOUND:
        return thr_stat->session_mgr->tcp_pkts_bypass_session_not_found;
    case STAT_TYPE_TCP_PKTS_BYPASS_DUPLICATED:
        return thr_stat->session_mgr->tcp_pkts_bypass_duplicated;
    case STAT_TYPE_UDP_PKTS_BYPASS_DUPLICATED:
        return thr_stat->session_mgr->udp_pkts_bypass_duplicated;
    case STAT_TYPE_UDP_PKTS_BYPASS_SESSION_EVICTED:
        return thr_stat->session_mgr->udp_pkts_bypass_session_evicted;

        // TCP segments
    case STAT_TYPE_TCP_SEGS_INPUT:
        return thr_stat->session_mgr->tcp_segs_input;
    case STAT_TYPE_TCP_SEGS_CONSUMED:
        return thr_stat->session_mgr->tcp_segs_consumed;
    case STAT_TYPE_TCP_SEGS_TIMEOUT:
        return thr_stat->session_mgr->tcp_segs_timeout;
    case STAT_TYPE_TCP_SEGS_RETRANSMITED:
        return thr_stat->session_mgr->tcp_segs_retransmited;
    case STAT_TYPE_TCP_SEGS_OVERLAPPED:
        return thr_stat->session_mgr->tcp_segs_overlapped;
    case STAT_TYPE_TCP_SEGS_OMITTED_TOO_MANY:
        return thr_stat->session_mgr->tcp_segs_omitted_too_many;
    case STAT_TYPE_TCP_SEGS_INORDER:
        return thr_stat->session_mgr->tcp_segs_inorder;
    case STAT_TYPE_TCP_SEGS_REORDERED:
        return thr_stat->session_mgr->tcp_segs_reordered;
    case STAT_TYPE_TCP_SEGS_BUFFERED:
        return thr_stat->session_mgr->tcp_segs_buffered;
    case STAT_TYPE_TCP_SEGS_FREED:
        return thr_stat->session_mgr->tcp_segs_freed;

    default:
        assert(0);
        return 0;
    }
}

// python3 -m pip install prettytable
// python3 -m pip install jinja2
// /opt/MESA/bin/fieldstat_exporter.py local -j log/stellar_fs4.json -e -l --clear-screen
struct stellar_stat *stellar_stat_new(uint16_t nr_thread)
{
    struct stellar_stat *stat = (struct stellar_stat *)calloc(1, sizeof(struct stellar_stat));
    if (stat == NULL)
    {
        return NULL;
    }

    snprintf(stat->output_file, sizeof(stat->output_file), "./log/stellar_fs4.json");

    stat->fs = fieldstat_easy_new(1, "stellar", NULL, 0);
    if (stat->fs == NULL)
    {
        STAT_LOG_ERROR("failed to create fieldstat_easy");
        goto error_out;
    }

    stat->nr_thread = nr_thread;
    for (int i = 0; i < MAX_THREAD_NUM; i++)
    {
        stat->flag[i] = IS_FREE;
    }

    for (size_t i = 0; i < STAT_TYPE_MAX; i++)
    {
        stat->stat_idx[i] = fieldstat_easy_register_counter(stat->fs, name[i]);
    }

    return stat;

error_out:
    stellar_stat_free(stat);
    return NULL;
}

void stellar_stat_free(struct stellar_stat *stat)
{
    if (stat)
    {
        if (stat->fs)
        {
            fieldstat_easy_free(stat->fs);
            stat->fs = NULL;
        }
        free(stat);
        stat = NULL;
    }
}

void stellar_stat_output(struct stellar_stat *stat)
{
    for (uint16_t i = 0; i < stat->nr_thread; i++)
    {
        if (ATOMIC_READ(&(stat->flag[i])) == IS_BUSY)
        {
            struct thread_stat *thr_stat = &stat->thr_stat[i];
            for (size_t j = 0; j < STAT_TYPE_MAX; j++)
            {
                stat->stat_val[j] += get_stat_value_by_idx(thr_stat, j);
            }
            memset(thr_stat, 0, sizeof(struct thread_stat));
            ATOMIC_SET(&(stat->flag[i]), IS_FREE);
        }
    }

    for (size_t j = 0; j < STAT_TYPE_MAX; j++)
    {
        fieldstat_easy_counter_set(stat->fs, 0, stat->stat_idx[j], NULL, 0, stat->stat_val[j]);
    }

    char *buff;
    size_t len;
    fieldstat_easy_output(stat->fs, &buff, &len);
    if (buff)
    {
        FILE *fp = fopen(stat->output_file, "w+");
        if (fp == NULL)
        {
            STAT_LOG_ERROR("failed to open file: %s, %s", stat->output_file, strerror(errno));
        }
        else
        {
            fwrite(buff, len, 1, fp);
            fflush(fp);
            fclose(fp);
        }
        free(buff);
    }

    for (size_t j = 0; j < STAT_TYPE_MAX; j++)
    {
        stat->stat_val[j] = 0;
    }
}

void stellar_stat_merge(struct stellar_stat *stat, const struct thread_stat *thr_stat, uint16_t thr_idx)
{
    if (ATOMIC_READ(&(stat->flag[thr_idx])) == IS_FREE)
    {
        memcpy(&stat->thr_stat[thr_idx], thr_stat, sizeof(struct thread_stat));
        ATOMIC_SET(&(stat->flag[thr_idx]), IS_BUSY);
    }
}

void stellar_stat_print(struct stellar_stat *stat __attribute__((unused)), const struct thread_stat *thr_stat, uint16_t thr_idx)
{
    for (size_t i = 0; i < STAT_TYPE_MAX; i++)
    {
        STAT_LOG_INFO("worker thread %lu => %-34s: %lu", thr_idx, name[i], get_stat_value_by_idx(thr_stat, i));
    }
}
