#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 "config.h"
#include "timestamp.h"
#include "id_generator.h"
#include "stellar_priv.h"
#include "session_priv.h"
#include "inject_priv.h"
#include "stellar/tuple.h"

#define MOCK_PLUGIN_LOG_DEBUG(format, ...) LOG_DEBUG("mock plugin", format, ##__VA_ARGS__)

/******************************************************************************
 * gmock plugin manager
 ******************************************************************************/

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_UDP_PAYLOAD = 4,
    INJECT_TYPE_CTRL_MSG = 5,
};

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;

static void inject_packet_plugin(struct session *sess, struct packet *pkt, struct inject_rule *rule)
{
    const struct tuple6 *tuple = session_get_tuple(sess);
    if (rule->ip_type == 4 && memcmp(&tuple->src_addr.v4, &rule->v4, sizeof(struct in_addr)) && memcmp(&tuple->dst_addr.v4, &rule->v4, sizeof(struct in_addr)))
    {
        return;
    }
    if (rule->ip_type == 6 && memcmp(&tuple->src_addr.v6, &rule->v6, sizeof(struct in6_addr)) && memcmp(&tuple->dst_addr.v6, &rule->v6, sizeof(struct in6_addr)))
    {
        return;
    }
    if (rule->port != 0 && tuple->src_port != rule->port && tuple->dst_port != rule->port)
    {
        return;
    }
    if (rule->count_dir == AFTER_RECV_C2S_N_PACKET && session_get_stat(sess, FLOW_DIRECTION_C2S, STAT_RAW_PACKETS_RECEIVED) != rule->count_num)
    {
        return;
    }
    if (rule->count_dir == AFTER_RECV_S2C_N_PACKET && session_get_stat(sess, FLOW_DIRECTION_S2C, STAT_RAW_PACKETS_RECEIVED) != rule->count_num)
    {
        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;
    }
    switch (rule->inject_type)
    {
    case INJECT_TYPE_TCP_RST:
        packet_set_action(pkt, PACKET_ACTION_DROP);
        EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_C2S) > 0);
        EXPECT_TRUE(stellar_inject_tcp_rst(sess, FLOW_DIRECTION_S2C) > 0);
        break;
    case INJECT_TYPE_TCP_FIN:
        packet_set_action(pkt, PACKET_ACTION_DROP);
        EXPECT_TRUE(stellar_inject_tcp_fin(sess, FLOW_DIRECTION_C2S) > 0);
        EXPECT_TRUE(stellar_inject_tcp_fin(sess, FLOW_DIRECTION_S2C) > 0);
        break;
    case INJECT_TYPE_TCP_PAYLOAD:
    case INJECT_TYPE_UDP_PAYLOAD:
        packet_set_action(pkt, PACKET_ACTION_DROP);
        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);
        break;
    case INJECT_TYPE_CTRL_MSG:
        // TOOD
        break;
    default:
        break;
    }
}

/******************************************************************************
 * mock plugin manager
 ******************************************************************************/

struct plugin_manager
{
};

void *plugin_manager_new_ctx(struct session *sess)
{
    return sess;
}

void plugin_manager_free_ctx(void *ctx)
{
    struct session *sess = (struct session *)ctx;

    char buff[4096] = {0};
    session_to_json(sess, buff, sizeof(buff));
    MOCK_PLUGIN_LOG_DEBUG("=> session: %s", buff);
}

struct plugin_manager *plugin_manager_new(void)
{
    static struct plugin_manager mgr;
    return &mgr;
}

void plugin_manager_free(struct plugin_manager *mgr)
{
}

void plugin_manager_dispatch_session(struct plugin_manager *mgr, struct session *sess, struct packet *pkt)
{
    struct tcp_segment *seg;
    enum session_type type = session_get_type(sess);
    uint16_t thr_idx = stellar_get_current_thread_index();
    MOCK_PLUGIN_LOG_DEBUG("=> thread: %d, session: %lu %s, type: %s, state: %s, c2s packet received: %lu, s2c packet received: %lu", thr_idx,
                          session_get_id(sess), session_get_tuple_str(sess),
                          session_type_to_str(type), session_state_to_str(session_get_state(sess)),
                          session_get_stat(sess, FLOW_DIRECTION_C2S, STAT_RAW_PACKETS_RECEIVED),
                          session_get_stat(sess, FLOW_DIRECTION_S2C, STAT_RAW_PACKETS_RECEIVED));

    if (packet_is_ctrl(pkt))
    {
    }
    else
    {
        switch (type)
        {
        case SESSION_TYPE_TCP:
            while ((seg = session_get_tcp_segment(sess)) != NULL)
            {
                session_free_tcp_segment(sess, seg);
            }
            break;
        case SESSION_TYPE_UDP:
            break;
        default:
            assert(0);
            break;
        }

        inject_packet_plugin(sess, pkt, &rule);
    }
}

void plugin_manager_dispatch_packet(struct plugin_manager *mgr, struct packet *pkt)
{
}

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

// curl -v http://http.badssl.selftest.gdnt-cloud.website --resolve "http.badssl.selftest.gdnt-cloud.website:80:192.0.2.110"

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 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, 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 parse_cmdline(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, "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("load inject rule:\n");
    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;
}

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

    timestamp_update();

    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;
    }
    STELLAR_LOG_STATE("start stellar (version: %s)\n %s", __stellar_version, logo_str);

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

    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 = timestamp_get_msec();
    while (!ATOMIC_READ(&runtime->need_exit))
    {
        timestamp_update();
        if (timestamp_get_msec() - runtime->stat_last_output_ts > 2000)
        {
            runtime->stat_last_output_ts = timestamp_get_msec();
            stellar_stat_output(runtime->stat);
        }
        usleep(1000); // 1ms
    }

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

    return 0;
}