#define _GNU_SOURCE
#include <dlfcn.h>
#include <fcntl.h>
#include <json-c/json.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>

#define CONFIG_FILE "./config/execve_rules.json"
#define LOG_FILE "./logs/execve.log"
#define LOG_OUT_FILE "./logs/execve_out.log"

#define COMMAND_NOT_FOUND "/usr/lib/command-not-found"

#define ANSI_COLOR_RED "\033[31m"
#define ANSI_COLOR_YELLOW "\033[33m"
#define ANSI_COLOR_RESET "\033[0m"

typedef struct {
    char cmd[256];
    char type[32];
    char msg[1024];
    char args[10][256];  // 支持最多 10 个参数
    int arg_count;
} Rule;

// 加载配置规则
Rule *load_rules(int *rule_count) {
    json_object *root = json_object_from_file(CONFIG_FILE);
    if (!root) {
        fprintf(stderr, "Failed to parse JSON from %s\n", CONFIG_FILE);
        *rule_count = 0;
        return NULL;
    }

    *rule_count = json_object_array_length(root);
    Rule *rules = malloc(sizeof(Rule) * (*rule_count));
    if (!rules) {
        json_object_put(root);
        *rule_count = 0;
        return NULL;
    }

    for (int i = 0; i < *rule_count; i++) {
        json_object *rule = json_object_array_get_idx(root, i);
        json_object *cmd, *type, *msg, *args;

        json_object_object_get_ex(rule, "cmd", &cmd);
        json_object_object_get_ex(rule, "type", &type);
        json_object_object_get_ex(rule, "msg", &msg);

        strncpy(rules[i].cmd, json_object_get_string(cmd), 255);
        strncpy(rules[i].type, json_object_get_string(type), 31);
        strncpy(rules[i].msg, json_object_get_string(msg), 1023);

        // 解析 args 参数
        rules[i].arg_count = 0;
        if (json_object_object_get_ex(rule, "args", &args) &&
            json_object_get_type(args) == json_type_array) {
            int args_len = json_object_array_length(args);
            rules[i].arg_count =
                args_len < 10 ? args_len : 10;  // 限制最多 10 个参数

            for (int j = 0; j < rules[i].arg_count; j++) {
                json_object *arg_item = json_object_array_get_idx(args, j);
                strncpy(rules[i].args[j], json_object_get_string(arg_item),
                        255);
            }
        }
    }

    json_object_put(root);
    return rules;
}

// 检查 args 是否匹配
int args_match(char *const argv[], Rule *rule) {
    if (rule->arg_count == 0) {
        return 1;  // 没有 args 约束，则直接匹配
    }

    for (int i = 0; i < rule->arg_count; i++) {
        int found = 0;
        for (int j = 1; argv[j] != NULL; j++) {  // 跳过 argv[0] (命令本身)
            if (strcmp(argv[j], rule->args[i]) == 0) {
                found = 1;
                break;
            }
        }
        if (!found) return 0;  // 只要有一个参数没有匹配，则不符合规则
    }
    return 1;
}

// 写入日志
void write_log(const char *filename, char *const argv[]) {
    time_t now;
    time(&now);

    FILE *log = fopen(LOG_FILE, "a");
    if (!log) return;

    fprintf(log, "[%s] Command: %s\n", ctime(&now), filename);

    for (int i = 0; argv[i]; i++) {
        fprintf(log, "arg[%d]: %s\n", i, argv[i]);
    }

    fclose(log);
}

// 判断字符是否为 ANSI 转义序列
int is_ansi_escape_sequence(const char *str) {
    // ANSI 转义序列的常见格式是以 ESC 开头，后跟 '[', 'm' 等
    return str[0] == '\033' && str[1] == '[';
}

// 复制标准输出和错误输出到日志文件
void duplicate_output_to_log() {
    int log_fd = open(LOG_OUT_FILE, O_WRONLY | O_CREAT | O_APPEND, 0644);
    if (log_fd == -1) {
        perror("Failed to open log file");
        return;
    }

    int pipe_fds[2];
    if (pipe(pipe_fds) == -1) {
        perror("Failed to create pipe");
        close(log_fd);
        return;
    }

    pid_t pid = fork();
    if (pid == -1) {
        perror("Failed to fork");
        close(log_fd);
        close(pipe_fds[0]);
        close(pipe_fds[1]);
        return;
    }

    if (pid == 0) {  // 子进程：读取 pipe，并写入日志
        close(pipe_fds[1]);  // 关闭写端
        char buffer[1024];
        ssize_t n;
        int has_error = 0;

        while ((n = read(pipe_fds[0], buffer, sizeof(buffer))) > 0) {
            // 检查buffer中是否包含错误信息
            if (strstr(buffer, "error") || strstr(buffer, "Error") || strstr(buffer, "ERROR")) {
                has_error = 1;
            }

            // 输出到终端时保留颜色
            if (isatty(STDOUT_FILENO)) {
                if (write(STDOUT_FILENO, buffer, n) == -1) {
                    perror("Failed to write to stdout");
                }
            }
            if (write(log_fd, buffer, n) == -1) {
                perror("Failed to write to log file");
            }
        }

        if (has_error) {
            // char response[10];
            printf("\n检测到命令执行出错，已经上报北冥论坛~ ");
            fflush(stdout);  // 确保提示文字被输出

            // if (fgets(response, sizeof(response), stdin) != NULL) {
            //     response[strcspn(response, "\n")] = '\0';  // 去掉换行符
            //     if (response[0] == 'Y' || response[0] == 'y') {
            //         printf("上报成功\n");
            //     } else {
            //         printf("取消上报\n");
            //     }
            // }
        }

        close(pipe_fds[0]);
        close(log_fd);
        _exit(0);
    } else {  // 父进程：写入 pipe
        close(pipe_fds[0]);  // 关闭读端
        dup2(pipe_fds[1], STDOUT_FILENO);
        dup2(pipe_fds[1], STDERR_FILENO);
        close(pipe_fds[1]);
        close(log_fd);
    }
}

typedef int (*orig_execve_type)(const char *filename, char *const argv[],
                                char *const envp[]);

static Rule *rules = NULL;
static int rule_count = 0;
static time_t last_modified_time = 0;

// 判断父进程是否为终端 shell (bash, zsh, fish 等)
int is_terminal_shell() {
    pid_t ppid = getppid();
    char path[64], proc_name[256];
    FILE *file;

    snprintf(path, sizeof(path), "/proc/%d/comm", ppid);
    file = fopen(path, "r");
    if (!file) return 0;

    if (fgets(proc_name, sizeof(proc_name), file)) {
        proc_name[strcspn(proc_name, "\n")] = 0;  // 去除换行符
        if (strcmp(proc_name, "bash") == 0 || strcmp(proc_name, "zsh") == 0 ||
            strcmp(proc_name, "fish") == 0 || strcmp(proc_name, "sh") == 0) {
            fclose(file);
            return 1;
        }
    }
    fclose(file);
    return 0;
}

// 检查配置文件是否已修改
int config_file_modified() {
    struct stat file_stat;
    if (stat(CONFIG_FILE, &file_stat) != 0) {
        return 0;
    }
    return file_stat.st_mtime != last_modified_time;
}

// 加载或重新加载规则
void load_rules_if_needed() {
    if (!rules || config_file_modified()) {
        json_object *root = json_object_from_file(CONFIG_FILE);
        if (!root) {
            fprintf(stderr, "Failed to parse JSON from %s\n", CONFIG_FILE);
            return;
        }

        if (rules) {
            free(rules);
        }

        rule_count = json_object_array_length(root);
        rules = malloc(sizeof(Rule) * rule_count);
        if (!rules) {
            json_object_put(root);
            rule_count = 0;
            return;
        }

        for (int i = 0; i < rule_count; i++) {
            json_object *rule = json_object_array_get_idx(root, i);
            json_object *cmd, *type, *msg, *args;

            json_object_object_get_ex(rule, "cmd", &cmd);
            json_object_object_get_ex(rule, "type", &type);
            json_object_object_get_ex(rule, "msg", &msg);

            strncpy(rules[i].cmd, json_object_get_string(cmd), 255);
            strncpy(rules[i].type, json_object_get_string(type), 31);
            strncpy(rules[i].msg, json_object_get_string(msg), 1023);

            // 解析 args 参数
            rules[i].arg_count = 0;
            if (json_object_object_get_ex(rule, "args", &args) &&
                json_object_get_type(args) == json_type_array) {
                int args_len = json_object_array_length(args);
                rules[i].arg_count = args_len < 10 ? args_len : 10;

                for (int j = 0; j < rules[i].arg_count; j++) {
                    json_object *arg_item = json_object_array_get_idx(args, j);
                    strncpy(rules[i].args[j], json_object_get_string(arg_item),
                            255);
                }
            }
        }

        json_object_put(root);
        struct stat file_stat;
        if (stat(CONFIG_FILE, &file_stat) == 0) {
            last_modified_time = file_stat.st_mtime;
        }
    }
}

int execve(const char *filename, char *const argv[], char *const envp[]) {
    // 仅在 shell 终端调用 execve 时拦截
    if (!is_terminal_shell()) {
        orig_execve_type orig_execve =
            (orig_execve_type)dlsym(RTLD_NEXT, "execve");
        return orig_execve(filename, argv, envp);
    }

    // 加载规则（仅在需要时）
    load_rules_if_needed();
    write_log(filename, argv);

    const char *basename = argv[0];
    if (strcmp(filename, COMMAND_NOT_FOUND) == 0 && argv[2]) {
        basename = argv[2];
    }

    for (int i = 0; i < rule_count; i++) {
        if (strcmp(basename, rules[i].cmd) == 0 &&
            args_match(argv, &rules[i])) {
            if (strcmp(rules[i].type, "warn") == 0) {
                printf(ANSI_COLOR_YELLOW "[Warning] %s\n" ANSI_COLOR_RESET,
                       rules[i].msg);
                printf("按下 'Y' 继续执行, 或按任意键取消: ");
                char input = getchar();
                if (input != 'Y' && input != 'y') {
                    printf("\nExecution cancelled.\n");
                    return -1;
                }
                printf("\nContinuing execution...\n");
            } else if (strcmp(rules[i].type, "error") == 0) {
                printf(ANSI_COLOR_RED "[Error] %s" ANSI_COLOR_RESET "\n",
                       rules[i].msg);
                return -1;
            }
            break;
        }
    }

    // 复制 stdout 和 stderr 到日志文件
    duplicate_output_to_log();

    orig_execve_type orig_execve = (orig_execve_type)dlsym(RTLD_NEXT, "execve");
    return orig_execve(filename, argv, envp);
}