#include <gtest/gtest.h>

#include <stdarg.h>
#include <unistd.h>
#include <sys/stat.h>
#include "packet_injector_test_frame.h"

static void system_cmd(const char *cmd, ...)
{
    char buf[1024] = {0};
    va_list args;

    va_start(args, cmd);
    vsnprintf(buf, sizeof(buf), cmd, args);
    va_end(args);

    system(buf);
}

static int replace_file_string(const char *file, const char *old_str, const char *new_str)
{
#define BUFFER_SIZE 1024

    FILE *in_fp = fopen(file, "r");
    if (in_fp == NULL)
    {
        printf("Open file %s failed, %s\n", file, strerror(errno));
        return -1;
    }

    FILE *tmp_fp = tmpfile();
    if (tmp_fp == NULL)
    {
        printf("Create temporary file failed, %s\n", strerror(errno));
        fclose(in_fp);
        return -1;
    }

    size_t old_len = strlen(old_str);
    size_t new_len = strlen(new_str);
    char buff[BUFFER_SIZE];

    while (fgets(buff, BUFFER_SIZE, in_fp))
    {
        char *pos = buff;
        if ((pos = strstr(pos, old_str)))
        {
            fwrite(buff, 1, pos - buff, tmp_fp); // Write characters before the old_str
            fwrite(new_str, 1, new_len, tmp_fp); // Write the new_str
            pos += old_len;                      // Move past the old_str
            fwrite(pos, 1, strlen(pos), tmp_fp); // Write characters after the old_str
        }
        else
        {
            fputs(buff, tmp_fp); // Write the remaining part of the line
        }
    }

    fclose(in_fp);
    fseek(tmp_fp, 0, SEEK_SET);

    FILE *out_fp = fopen(file, "w");
    if (out_fp == NULL)
    {
        printf("Open file %s for writing failed, %s\n", file, strerror(errno));
        fclose(tmp_fp);
        return -1;
    }

    while (fgets(buff, BUFFER_SIZE, tmp_fp))
    {
        fputs(buff, out_fp); // Write the contents of the temporary file to the original file
    }

    fclose(tmp_fp);
    fclose(out_fp);

    return 0;
}

static void expect_cmp_inject(const char *work_dir,
                              const char *expect_dir_abs_path, const char *expect_pcap_file,
                              const char *output_dir_abs_path, const char *output_pcap_file,
                              const char *diff_skip_pattern, const char *flow_dir)
{
    printf("\033[32m ============================================= \033[0m\n");
    printf("\033[32m Compare [%s] expect pcap and output pcap \033[0m\n", flow_dir);
    printf("\033[32m ============================================= \033[0m\n");

    struct stat s;
    char expect_pcap_file_abs_path[1024] = {0};
    char output_pcap_file_abs_path[1024] = {0};
    char expect_json_file_abs_path[1024] = {0};
    char output_json_file_abs_path[1024] = {0};
    char diff_txt_file_abs_path[1024] = {0};

    // absulute path
    snprintf(expect_pcap_file_abs_path, sizeof(expect_pcap_file_abs_path), "%s/%s", expect_dir_abs_path, expect_pcap_file);
    snprintf(output_pcap_file_abs_path, sizeof(output_pcap_file_abs_path), "%s/%s", output_dir_abs_path, output_pcap_file);
    snprintf(expect_json_file_abs_path, sizeof(expect_json_file_abs_path), "%s/expect_%s.json", work_dir, flow_dir);
    snprintf(output_json_file_abs_path, sizeof(output_json_file_abs_path), "%s/output_%s.json", work_dir, flow_dir);
    snprintf(diff_txt_file_abs_path, sizeof(diff_txt_file_abs_path), "%s/diff_%s.txt", work_dir, flow_dir);

    // check pcap file size
    stat(expect_pcap_file_abs_path, &s);
    EXPECT_TRUE(s.st_size > 0);

    stat(output_pcap_file_abs_path, &s);
    EXPECT_TRUE(s.st_size > 0);

    // tcpdump
    printf("\033[32m tcpdump read [%s] expect pcap (%s) \033[0m\n", flow_dir, expect_pcap_file);
    system_cmd("tcpdump -r %s", expect_pcap_file_abs_path);

    printf("\033[32m tcpdump read [%s] output pcap (%s) \033[0m\n", flow_dir, output_pcap_file);
    system_cmd("tcpdump -r %s", output_pcap_file_abs_path);

    // tshark
    system_cmd("tshark -r %s -T json | jq >> %s", expect_pcap_file_abs_path, expect_json_file_abs_path);
    system_cmd("tshark -r %s -T json | jq >> %s", output_pcap_file_abs_path, output_json_file_abs_path);

    // check json file size
    stat(expect_json_file_abs_path, &s);
    EXPECT_TRUE(s.st_size > 0);

    stat(output_json_file_abs_path, &s);
    EXPECT_TRUE(s.st_size > 0);

    // diff
    system_cmd("diff %s %s %s >> %s", diff_skip_pattern, expect_json_file_abs_path, output_json_file_abs_path, diff_txt_file_abs_path);

    // check diff file size
    stat(diff_txt_file_abs_path, &s);
    EXPECT_TRUE(s.st_size == 0);
}

void packet_injector_test_frame_run(struct packet_injector_case *test)
{
    printf("\033[32mTest: %s\033[0m\n", test->descriptor);

    char config_file_abs_path[1024] = {0};
    char input_dir_abs_path[1024] = {0};
    char output_dir_abs_path[1024] = {0};
    char expect_dir_abs_path[1024] = {0};

    // absulute path
    snprintf(config_file_abs_path, sizeof(config_file_abs_path), "%s/conf/stellar.toml", test->work_dir);
    snprintf(input_dir_abs_path, sizeof(input_dir_abs_path), "%s/input/", test->work_dir);
    snprintf(output_dir_abs_path, sizeof(output_dir_abs_path), "%s/output/", test->work_dir);
    snprintf(expect_dir_abs_path, sizeof(expect_dir_abs_path), "%s/expect/", test->work_dir);

    // create directory
    system_cmd("rm -rf %s", test->work_dir);
    system_cmd("mkdir -p %s", input_dir_abs_path);
    system_cmd("mkdir -p %s", output_dir_abs_path);
    system_cmd("mkdir -p %s", expect_dir_abs_path);

    // copy file to work directory
    if (test->c2s_expect_pcap)
    {
        system_cmd("cp %s/%s %s", test->input_prefix, test->c2s_expect_pcap, expect_dir_abs_path);
    }
    if (test->s2c_expect_pcap)
    {
        system_cmd("cp %s/%s %s", test->input_prefix, test->s2c_expect_pcap, expect_dir_abs_path);
    }
    system_cmd("cp %s/%s %s", test->input_prefix, test->input_pcap, input_dir_abs_path);
    system_cmd("cp -r conf %s/", test->work_dir);
    system_cmd("cp packet_injector %s/", test->work_dir);

    // replace config file
    char temp[2048] = {0};
    snprintf(temp, sizeof(temp), "dumpfile_dir = \"%s\"", input_dir_abs_path);
    EXPECT_TRUE(replace_file_string(config_file_abs_path, "mode = marsio", "mode = dumpfile") == 0);
    EXPECT_TRUE(replace_file_string(config_file_abs_path, "dumpfile_dir = \"/tmp/dumpfile/\"", temp) == 0);

    // run packet injector
    chdir(test->work_dir);
    system_cmd(test->packet_injector_cmd);
    chdir("../");

    // compare pcap
    if (test->c2s_output_pcap && test->c2s_expect_pcap)
    {
        system_cmd("mv %s/%s %s", test->output_prefix, test->c2s_output_pcap, output_dir_abs_path);
        expect_cmp_inject(test->work_dir,
                          expect_dir_abs_path, test->c2s_expect_pcap,
                          output_dir_abs_path, test->c2s_output_pcap,
                          test->diff_skip_pattern, "C2S");
    }
    if (test->s2c_output_pcap && test->s2c_expect_pcap)
    {
        system_cmd("mv %s/%s %s", test->output_prefix, test->s2c_output_pcap, output_dir_abs_path);
        expect_cmp_inject(test->work_dir,
                          expect_dir_abs_path, test->s2c_expect_pcap,
                          output_dir_abs_path, test->s2c_output_pcap,
                          test->diff_skip_pattern, "S2C");
    }

    // clean work directory
    if (test->finish_clean_work_dir)
    {
        system_cmd("rm -rf %s", test->work_dir);
    }
}