#include <gtest/gtest.h>
#include <stdio.h>
#include <errno.h>
#include <unistd.h>
#include <string.h>
#include <signal.h>
#include <stdlib.h>
#include <netinet/tcp.h>

#include "logo.h"
#include "times.h"
#include "config.h"
#include "id_generator.h"
#include "stellar_priv.h"
#include "session_priv.h"
#include "inject_priv.h"
#include "stellar/tuple.h"
#include "stellar/session_mq.h"

/******************************************************************************
 * packet_injector
 ******************************************************************************/

enum condition
{
    AFTER_RECV_C2S_N_PACKET = 1,
    AFTER_RECV_S2C_N_PACKET = 2,
};

enum inject_type
{
    INJECT_TYPE_TCP_RST = 1,
    INJECT_TYPE_TCP_FIN = 2,
    INJECT_TYPE_TCP_PAYLOAD = 3,
    INJECT_TYPE_TCP_PAYLOAD_FIN_RST = 4,
    INJECT_TYPE_UDP_PAYLOAD = 5,
    INJECT_TYPE_CTRL_MSG = 6,
};

struct inject_rule
{
    int ip_type;
    struct in_addr v4;  /* network order */
    struct in6_addr v6; /* network order */
    int port;           /* network order */

    enum inject_type inject_type;
    enum condition count_dir;
    uint64_t count_num;
} rule;

#define INJECT_PACKET_PLUGIN_LOG_DEBUG(format, ...) LOG_DEBUG("inject packet plugin", format, ##__VA_ARGS__)

static void usage(char *cmd)
{
    printf("Usage: %s [options]\n\n", cmd);
    printf("Options:\n");
    printf("  -h <ip>       Host IP address\n");
    printf("  -p <port>     Port number\n");
    printf("  -t <type>     Type of manipulation\n");
    printf("                   Options: tcp-rst, tcp-fin, tcp-payload, tcp-payload-fin-rst, udp-payload, ctrl-msg\n");
    printf("  -c <condition> Condition for manipulation\n");
    printf("                   Options: c2s-packet, s2c-packet\n");
    printf("  -n <number>   Number of packets received before injecting action\n\n");
    printf("Example:\n");
    printf("  %s -h 192.168.1.100 -p 8080 -t tcp-payload -c c2s-packet -n 5\n", cmd);
    printf("  %s -h 2001:db8::1   -p 8080 -t tcp-rst     -c s2c-packet -n 10\n", cmd);
    printf("\n");
}

static int packet_injector_on_init(int argc, char **argv, struct inject_rule *rule)
{
    memset(rule, 0, sizeof(struct inject_rule));

    int opt = 0;
    const char *host = NULL;
    const char *type = NULL;
    const char *condition = NULL;
    while ((opt = getopt(argc, argv, "h:p:t:c:n:")) != -1)
    {
        switch (opt)
        {
        case 'h':
            host = optarg;
            break;
        case 'p':
            rule->port = htons(atoi(optarg));
            break;
        case 't':
            type = optarg;
            break;
        case 'c':
            condition = optarg;
            break;
        case 'n':
            rule->count_num = atoi(optarg);
            break;
        default:
            usage(argv[0]);
            break;
        }
    }

    if (host)
    {
        if (inet_pton(AF_INET, host, &rule->v4) != 1)
        {
            if (inet_pton(AF_INET6, host, &rule->v6) != 1)
            {
                printf("unable to convert host %s to IPv4 / IPv6\n", host);
                return -1;
            }
            else
            {
                rule->ip_type = 6;
            }
        }
        else
        {
            rule->ip_type = 4;
        }
    }

    if (type == NULL)
    {
        usage(argv[0]);
        printf("invalid type\n");
        return -1;
    }
    else if (strcmp(type, "tcp-rst") == 0)
    {
        rule->inject_type = INJECT_TYPE_TCP_RST;
    }
    else if (strcmp(type, "tcp-fin") == 0)
    {
        rule->inject_type = INJECT_TYPE_TCP_FIN;
    }
    else if (strcmp(type, "tcp-payload") == 0)
    {
        rule->inject_type = INJECT_TYPE_TCP_PAYLOAD;
    }
    else if (strcmp(type, "tcp-payload-fin-rst") == 0)
    {
        rule->inject_type = INJECT_TYPE_TCP_PAYLOAD_FIN_RST;
    }
    else if (strcmp(type, "udp-payload") == 0)
    {
        rule->inject_type = INJECT_TYPE_UDP_PAYLOAD;
    }
    else if (strcmp(type, "ctrl-msg") == 0)
    {
        rule->inject_type = INJECT_TYPE_CTRL_MSG;
    }
    else
    {
        usage(argv[0]);
        printf("invalid type\n");
        return -1;
    }

    if (condition == NULL)
    {
        usage(argv[0]);
        printf("invalid condition\n");
        return -1;
    }
    else if (strcmp(condition, "c2s-packet") == 0)
    {
        rule->count_dir = AFTER_RECV_C2S_N_PACKET;
    }
    else if (strcmp(condition, "s2c-packet") == 0)
    {
        rule->count_dir = AFTER_RECV_S2C_N_PACKET;
    }
    else
    {
        usage(argv[0]);
        printf("invalid condition\n");
        return -1;
    }

    if (rule->count_num <= 0)
    {
        usage(argv[0]);
        printf("invalid count\n");
        return -1;
    }
    printf("%s load inject rule:\n", argv[0]);
    printf("      host      : %s\n", host);
    printf("      port      : %d\n", ntohs(rule->port));
    printf("      type      : %s\n", type);
    printf("      condition : %s\n", condition);
    printf("      count     : %lu\n\n", rule->count_num);

    return 0;
}

static void packet_injector_on_msg(struct session *sess, int topic_id, const void *msg, void *per_session_ctx, void *plugin_env)
{
    char buffer[1024] = {0};
    struct inject_rule *p_rule = &rule;
    // struct packet *pkt = (struct packet *)msg;
    const struct tuple6 *tuple = session_get_tuple6(sess);
    if (p_rule->ip_type == 4 &&
        memcmp(&tuple->src_addr.v4, &p_rule->v4, sizeof(struct in_addr)) &&
        memcmp(&tuple->dst_addr.v4, &p_rule->v4, sizeof(struct in_addr)))
    {
        return;
    }
    if (p_rule->ip_type == 6 &&
        memcmp(&tuple->src_addr.v6, &p_rule->v6, sizeof(struct in6_addr)) &&
        memcmp(&tuple->dst_addr.v6, &p_rule->v6, sizeof(struct in6_addr)))
    {
        return;
    }
    if (p_rule->port != 0 && tuple->src_port != p_rule->port && tuple->dst_port != p_rule->port)
    {
        return;
    }
    if (session_get_stat(sess, FLOW_DIRECTION_C2S, STAT_INJECTED_PACKETS_SUCCESS) > 0 ||
        session_get_stat(sess, FLOW_DIRECTION_S2C, STAT_INJECTED_PACKETS_SUCCESS) > 0)
    {
        return;
    }
    if (p_rule->count_dir == AFTER_RECV_C2S_N_PACKET && session_get_stat(sess, FLOW_DIRECTION_C2S, STAT_RAW_PACKETS_RECEIVED) != p_rule->count_num)
    {
        return;
    }
    if (p_rule->count_dir == AFTER_RECV_S2C_N_PACKET && session_get_stat(sess, FLOW_DIRECTION_S2C, STAT_RAW_PACKETS_RECEIVED) != p_rule->count_num)
    {
        return;
    }
    switch (p_rule->inject_type)
    {
    case INJECT_TYPE_TCP_RST:
        EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_C2S) > 0);
        EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_S2C) > 0);
        session_set_discard(sess);
        break;
    case INJECT_TYPE_TCP_FIN:
        EXPECT_TRUE(stellar_inject_tcp_fin(sess, FLOW_DIRECTION_C2S) > 0);
        EXPECT_TRUE(stellar_inject_tcp_fin(sess, FLOW_DIRECTION_S2C) > 0);
        session_set_discard(sess);
        break;
    case INJECT_TYPE_TCP_PAYLOAD:
        snprintf(buffer, sizeof(buffer), "HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n%s", 5 + 5 + 2, "Hello");
        EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_S2C, buffer, strlen(buffer)) > 0); // inject payload to client
        EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_S2C, "World\r\n", 7) > 0);         // inject payload to client
        EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_S2C) > 0);                         // inject RST to client
        EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_C2S) > 0);                         // inject RST to server
        session_set_discard(sess);
        break;
    case INJECT_TYPE_TCP_PAYLOAD_FIN_RST:
        snprintf(buffer, sizeof(buffer), "HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n%s", 5 + 5 + 2, "Hello");
        EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_S2C, buffer, strlen(buffer)) > 0); // inject payload to client
        EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_S2C, "World\r\n", 7) > 0);         // inject payload to client
        EXPECT_TRUE(stellar_inject_tcp_fin(sess, FLOW_DIRECTION_S2C) > 0);                         // inject FIN to client
        EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_S2C) > 0);                         // inject RST to client
        EXPECT_TRUE(stellar_inject_tcp_fin(sess, FLOW_DIRECTION_C2S) > 0);                         // inject FIN to server
        EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_C2S) > 0);                         // inject RST to server
        session_set_discard(sess);
        break;
    case INJECT_TYPE_UDP_PAYLOAD:
        EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_C2S, "Hello Server", 12) > 0);
        EXPECT_TRUE(stellar_inject_payload(sess, FLOW_DIRECTION_S2C, "Hello Client", 12) > 0);
        session_set_discard(sess);
        break;
    case INJECT_TYPE_CTRL_MSG:
        // TOOD
        break;
    default:
        break;
    }
}

static void *packet_injector_on_sess_new(struct session *sess, void *plugin_env)
{
    return NULL;
}

static void packet_injector_on_sess_free(struct session *sess, void *ctx, void *plugin_env)
{
    char buff[4096] = {0};
    session_to_json(sess, buff, sizeof(buff));
    INJECT_PACKET_PLUGIN_LOG_DEBUG("=> session: %s", buff);
}

/******************************************************************************
 * main
 ******************************************************************************/

static const char *log_config_file = "./conf/log.toml";
static const char *stellar_config_file = "./conf/stellar.toml";

static void signal_handler(int signo)
{
    if (signo == SIGINT)
    {
        STELLAR_LOG_STATE("SIGINT received, notify threads to exit !!!");
        ATOMIC_SET(&runtime->need_exit, 1);
    }

    if (signo == SIGQUIT)
    {
        STELLAR_LOG_STATE("SIGQUIT received, notify threads to exit !!!");
        ATOMIC_SET(&runtime->need_exit, 1);
    }

    if (signo == SIGTERM)
    {
        STELLAR_LOG_STATE("SIGTERM received, notify threads to exit !!!");
        ATOMIC_SET(&runtime->need_exit, 1);
    }

    if (signo == SIGHUP)
    {
        STELLAR_LOG_STATE("SIGHUP received, reload log level !!!");
        log_reload_level(log_config_file);
    }
}

static int all_session_have_freed(void)
{
    for (int i = 0; i < config->io_opts.nr_threads; i++)
    {
        struct session_manager *sess_mgr = runtime->threads[i].sess_mgr;
        struct session_manager_stat *sess_stat = session_manager_stat(sess_mgr);

        if (ATOMIC_READ(&sess_stat->curr_nr_tcp_sess_used) != 0)
        {
            return 0;
        }

        if (ATOMIC_READ(&sess_stat->curr_nr_udp_sess_used) != 0)
        {
            return 0;
        }
    }

    return 1;
}

static int all_stat_have_output(void)
{
    static int count = 0;
    if (runtime->stat_last_output_ts == stellar_get_monotonic_time_msec())
    {
        count++;
    }

    return count == 2;
}

int packet_injector_main(int argc, char **argv)
{
    if (packet_injector_on_init(argc, argv, &rule) != 0)
    {
        return -1;
    }

    stellar_update_time_cache();
    struct stellar st = {runtime};
    int sess_plug_id;
    int tcp_topic_id;
    int udp_topic_id;

    signal(SIGINT, signal_handler);
    signal(SIGQUIT, signal_handler);
    signal(SIGTERM, signal_handler);
    signal(SIGHUP, signal_handler);

    if (log_init(log_config_file) != 0)
    {
        STELLAR_LOG_ERROR("unable to init log");
        goto error_out;
    }

    if (stellar_load_config(stellar_config_file, config) != 0)
    {
        STELLAR_LOG_ERROR("unable to load config file");
        goto error_out;
    }
    stellar_print_config(config);
    STELLAR_LOG_DEBUG("sizeof(struct session) = %lu bytes", sizeof(struct session));

    if (id_generator_init(config->dev_opts.base, config->dev_opts.offset) != 0)
    {
        STELLAR_LOG_ERROR("unable to init id generator");
        goto error_out;
    }

    runtime->stat = stellar_stat_new(config->io_opts.nr_threads);
    if (runtime->stat == NULL)
    {
        STELLAR_LOG_ERROR("unable to create stellar stat");
        goto error_out;
    }
    runtime->plug_mgr = plugin_manager_init(&st, "./stellar_plugin/spec.toml");
    if (runtime->plug_mgr == NULL)
    {
        STELLAR_LOG_ERROR("unable to create plugin manager");
        goto error_out;
    }

    runtime->packet_io = packet_io_new(&config->io_opts);
    if (runtime->packet_io == NULL)
    {
        STELLAR_LOG_ERROR("unable to create packet io");
        goto error_out;
    }

    sess_plug_id = stellar_session_plugin_register(&st, packet_injector_on_sess_new, packet_injector_on_sess_free, NULL);
    tcp_topic_id = stellar_session_mq_get_topic_id(&st, TOPIC_TCP);
    udp_topic_id = stellar_session_mq_get_topic_id(&st, TOPIC_UDP);

    stellar_session_mq_subscribe(&st, tcp_topic_id, packet_injector_on_msg, sess_plug_id);
    stellar_session_mq_subscribe(&st, udp_topic_id, packet_injector_on_msg, sess_plug_id);

    if (stellar_thread_init(runtime, config) != 0)
    {
        STELLAR_LOG_ERROR("unable to init thread context");
        goto error_out;
    }

    if (stellar_thread_run(runtime, config) != 0)
    {
        STELLAR_LOG_ERROR("unable to create worker thread");
        goto error_out;
    }

    runtime->stat_last_output_ts = stellar_get_monotonic_time_msec();
    while (!ATOMIC_READ(&runtime->need_exit))
    {
        stellar_update_time_cache();
        if (stellar_get_monotonic_time_msec() - runtime->stat_last_output_ts > 2000)
        {
            runtime->stat_last_output_ts = stellar_get_monotonic_time_msec();
            stellar_stat_output(runtime->stat);
        }
        usleep(1000); // 1ms

        // Only available in dump file mode, automatically exits when all sessions have been released
        if (packet_io_wait_exit(runtime->packet_io) && all_session_have_freed() && all_stat_have_output())
        {
            STELLAR_LOG_STATE("all sessions have been released and all stat have been output, notify threads to exit !!!");
            ATOMIC_SET(&runtime->need_exit, 1);
        }
    }

error_out:
    stellar_thread_join(runtime, config);
    stellar_thread_clean(runtime, config);
    packet_io_free(runtime->packet_io);
    plugin_manager_exit(runtime->plug_mgr);
    stellar_stat_free(runtime->stat);
    STELLAR_LOG_STATE("stellar exit !!!\n");
    log_free();

    return 0;
}

int __attribute__((weak)) main(int argc, char **argv)
{
    return packet_injector_main(argc, argv);
}