#define _GNU_SOURCE

#include <arpa/inet.h>
#include <errno.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#define VALID_MNEMONIC "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about"
#define LONG_INPUT_LEN 1024

static int g_failures = 0;

static void check_condition(const char *name, int condition) {
    if (condition) {
        printf("PASS: %s\n", name);
    } else {
        printf("FAIL: %s\n", name);
        g_failures++;
    }
}

static int sleep_ms(int ms) {
    struct timespec ts;
    ts.tv_sec = ms / 1000;
    ts.tv_nsec = (long)(ms % 1000) * 1000000L;
    return nanosleep(&ts, NULL);
}

static int write_full(int fd, const void *buf, size_t len) {
    const unsigned char *p = (const unsigned char *)buf;
    size_t off = 0;

    while (off < len) {
        ssize_t n = write(fd, p + off, len - off);
        if (n < 0) {
            if (errno == EINTR) {
                continue;
            }
            return -1;
        }
        off += (size_t)n;
    }

    return 0;
}

static int connect_socket_retry(const char *name, int timeout_ms) {
    int elapsed = 0;

    while (elapsed < timeout_ms) {
        int fd;
        struct sockaddr_un addr;
        socklen_t addr_len;

        fd = socket(AF_UNIX, SOCK_STREAM, 0);
        if (fd < 0) {
            return -1;
        }

        memset(&addr, 0, sizeof(addr));
        addr.sun_family = AF_UNIX;
        addr.sun_path[0] = '\0';
        strncpy(&addr.sun_path[1], name, sizeof(addr.sun_path) - 2);
        addr.sun_path[sizeof(addr.sun_path) - 1] = '\0';
        addr_len = (socklen_t)(sizeof(sa_family_t) + 1 + strlen(name));

        if (connect(fd, (struct sockaddr *)&addr, addr_len) == 0) {
            return fd;
        }

        close(fd);
        sleep_ms(100);
        elapsed += 100;
    }

    return -1;
}

static pid_t spawn_with_stdin_and_fd(const char *const argv[],
                                     const char *stdin_data,
                                     int close_stdin_after_write,
                                     const char *fd_data,
                                     int fd_num,
                                     int close_fd_after_write,
                                     int *stdin_write_end,
                                     int *fd_write_end,
                                     int *capture_read_end) {
    int stdin_pipe[2] = {-1, -1};
    int fd_pipe[2] = {-1, -1};
    int capture_pipe[2] = {-1, -1};
    pid_t child;

    if (pipe(stdin_pipe) != 0) {
        return -1;
    }
    if (pipe(capture_pipe) != 0) {
        close(stdin_pipe[0]);
        close(stdin_pipe[1]);
        return -1;
    }
    if (fd_data != NULL) {
        if (pipe(fd_pipe) != 0) {
            close(stdin_pipe[0]);
            close(stdin_pipe[1]);
            close(capture_pipe[0]);
            close(capture_pipe[1]);
            return -1;
        }
    }

    child = fork();
    if (child < 0) {
        close(stdin_pipe[0]);
        close(stdin_pipe[1]);
        close(capture_pipe[0]);
        close(capture_pipe[1]);
        if (fd_pipe[0] >= 0) close(fd_pipe[0]);
        if (fd_pipe[1] >= 0) close(fd_pipe[1]);
        return -1;
    }

    if (child == 0) {
        dup2(stdin_pipe[0], STDIN_FILENO);
        dup2(capture_pipe[1], STDOUT_FILENO);
        dup2(capture_pipe[1], STDERR_FILENO);

        close(stdin_pipe[0]);
        close(stdin_pipe[1]);
        close(capture_pipe[0]);
        close(capture_pipe[1]);

        if (fd_data != NULL) {
            dup2(fd_pipe[0], fd_num);
            close(fd_pipe[0]);
            close(fd_pipe[1]);
        }

        execv("./build/nsigner", (char *const *)argv);
        _exit(127);
    }

    close(stdin_pipe[0]);
    close(capture_pipe[1]);

    if (fd_data != NULL) {
        close(fd_pipe[0]);
    }

    if (stdin_data != NULL) {
        (void)write_full(stdin_pipe[1], stdin_data, strlen(stdin_data));
        if (close_stdin_after_write) {
            close(stdin_pipe[1]);
            stdin_pipe[1] = -1;
        }
    }

    if (fd_data != NULL) {
        (void)write_full(fd_pipe[1], fd_data, strlen(fd_data));
        if (close_fd_after_write) {
            close(fd_pipe[1]);
            fd_pipe[1] = -1;
        }
    }

    if (stdin_write_end != NULL) {
        *stdin_write_end = stdin_pipe[1];
    } else if (stdin_pipe[1] >= 0) {
        close(stdin_pipe[1]);
    }

    if (fd_write_end != NULL) {
        *fd_write_end = fd_pipe[1];
    } else if (fd_pipe[1] >= 0) {
        close(fd_pipe[1]);
    }

    if (capture_read_end != NULL) {
        *capture_read_end = capture_pipe[0];
    } else {
        close(capture_pipe[0]);
    }

    return child;
}

static int read_capture_nonblocking(int fd, char *buf, size_t buf_sz) {
    ssize_t n;
    if (buf == NULL || buf_sz == 0) {
        return -1;
    }
    n = read(fd, buf, buf_sz - 1);
    if (n < 0) {
        if (errno == EAGAIN || errno == EWOULDBLOCK) {
            buf[0] = '\0';
            return 0;
        }
        return -1;
    }
    buf[n] = '\0';
    return 0;
}

static int wait_exit_timeout(pid_t pid, int timeout_ms, int *status) {
    int elapsed = 0;
    while (elapsed < timeout_ms) {
        pid_t rc = waitpid(pid, status, WNOHANG);
        if (rc == pid) {
            return 0;
        }
        if (rc < 0) {
            return -1;
        }
        sleep_ms(50);
        elapsed += 50;
    }
    return -1;
}

static void test_stdin_happy_path(void) {
    const char *const argv[] = {
        "./build/nsigner", "--socket-name", "nsigner_test_mi_stdin_ok", "--mnemonic-stdin", NULL
    };
    int cap_fd = -1;
    int status = 0;
    int fd;
    pid_t child = spawn_with_stdin_and_fd(argv,
                                          VALID_MNEMONIC "\n",
                                          1,
                                          NULL,
                                          -1,
                                          0,
                                          NULL,
                                          NULL,
                                          &cap_fd);

    check_condition("spawn --mnemonic-stdin happy-path child", child > 0);
    if (child <= 0) {
        return;
    }

    fd = connect_socket_retry("nsigner_test_mi_stdin_ok", 5000);
    check_condition("--mnemonic-stdin reaches running server", fd >= 0);
    if (fd >= 0) {
        close(fd);
    }

    (void)kill(child, SIGTERM);
    (void)waitpid(child, &status, 0);
    if (cap_fd >= 0) {
        close(cap_fd);
    }
}

static void test_stdin_eof_without_newline(void) {
    const char *const argv[] = {
        "./build/nsigner", "--socket-name", "nsigner_test_mi_stdin_eof", "--mnemonic-stdin", NULL
    };
    int cap_fd = -1;
    int status = 0;
    int fd;
    pid_t child = spawn_with_stdin_and_fd(argv,
                                          VALID_MNEMONIC,
                                          1,
                                          NULL,
                                          -1,
                                          0,
                                          NULL,
                                          NULL,
                                          &cap_fd);

    check_condition("spawn --mnemonic-stdin EOF-without-newline child", child > 0);
    if (child <= 0) {
        return;
    }

    fd = connect_socket_retry("nsigner_test_mi_stdin_eof", 5000);
    check_condition("--mnemonic-stdin EOF without newline accepted", fd >= 0);
    if (fd >= 0) {
        close(fd);
    }

    (void)kill(child, SIGTERM);
    (void)waitpid(child, &status, 0);
    if (cap_fd >= 0) {
        close(cap_fd);
    }
}

static void test_stdin_invalid_mnemonic(void) {
    const char *const argv[] = {
        "./build/nsigner", "--mnemonic-stdin", NULL
    };
    int cap_fd = -1;
    int status = 0;
    char cap[512];
    pid_t child = spawn_with_stdin_and_fd(argv,
                                          "not a valid mnemonic\n",
                                          1,
                                          NULL,
                                          -1,
                                          0,
                                          NULL,
                                          NULL,
                                          &cap_fd);

    check_condition("spawn --mnemonic-stdin invalid child", child > 0);
    if (child <= 0) {
        return;
    }

    check_condition("--mnemonic-stdin invalid exits", wait_exit_timeout(child, 2000, &status) == 0);
    check_condition("--mnemonic-stdin invalid exits non-zero", WIFEXITED(status) && WEXITSTATUS(status) != 0);

    if (cap_fd >= 0) {
        (void)read_capture_nonblocking(cap_fd, cap, sizeof(cap));
        check_condition("--mnemonic-stdin invalid emits validation failure", strstr(cap, "mnemonic validation failed") != NULL);
        close(cap_fd);
    }
}

static void test_stdin_overflow(void) {
    const char *const argv[] = {
        "./build/nsigner", "--mnemonic-stdin", NULL
    };
    char long_input[LONG_INPUT_LEN + 2];
    int i;
    int cap_fd = -1;
    int status = 0;
    char cap[512];
    pid_t child;

    for (i = 0; i < LONG_INPUT_LEN; ++i) {
        long_input[i] = 'a';
    }
    long_input[LONG_INPUT_LEN] = '\n';
    long_input[LONG_INPUT_LEN + 1] = '\0';

    child = spawn_with_stdin_and_fd(argv,
                                    long_input,
                                    1,
                                    NULL,
                                    -1,
                                    0,
                                    NULL,
                                    NULL,
                                    &cap_fd);

    check_condition("spawn --mnemonic-stdin overflow child", child > 0);
    if (child <= 0) {
        return;
    }

    check_condition("--mnemonic-stdin overflow exits", wait_exit_timeout(child, 2000, &status) == 0);
    check_condition("--mnemonic-stdin overflow exits non-zero", WIFEXITED(status) && WEXITSTATUS(status) != 0);

    if (cap_fd >= 0) {
        (void)read_capture_nonblocking(cap_fd, cap, sizeof(cap));
        check_condition("--mnemonic-stdin overflow emits length error", strstr(cap, "exceeded maximum length") != NULL);
        close(cap_fd);
    }
}

static void test_fd_happy_path(void) {
    const char *const argv[] = {
        "./build/nsigner", "--socket-name", "nsigner_test_mi_fd_ok", "--mnemonic-fd", "3", NULL
    };
    int cap_fd = -1;
    int status = 0;
    int fd;
    pid_t child = spawn_with_stdin_and_fd(argv,
                                          "",
                                          0,
                                          VALID_MNEMONIC "\n",
                                          3,
                                          1,
                                          NULL,
                                          NULL,
                                          &cap_fd);

    check_condition("spawn --mnemonic-fd happy-path child", child > 0);
    if (child <= 0) {
        return;
    }

    fd = connect_socket_retry("nsigner_test_mi_fd_ok", 5000);
    check_condition("--mnemonic-fd reaches running server", fd >= 0);
    if (fd >= 0) {
        close(fd);
    }

    (void)kill(child, SIGTERM);
    (void)waitpid(child, &status, 0);
    if (cap_fd >= 0) {
        close(cap_fd);
    }
}

static void test_mutual_exclusion_and_validation_errors(void) {
    struct {
        const char *name;
        const char *const *argv;
        const char *expect_substr;
    } cases[] = {
        {
            "mutual exclusion stdin+fd",
            (const char *const[]){"./build/nsigner", "--mnemonic-stdin", "--mnemonic-fd", "3", NULL},
            "mutually exclusive"
        },
        {
            "stdin with stdio listener",
            (const char *const[]){"./build/nsigner", "--listen", "stdio", "--mnemonic-stdin", NULL},
            "requires --listen unix or --listen tcp"
        },
        {
            "mnemonic-fd 1 rejected",
            (const char *const[]){"./build/nsigner", "--mnemonic-fd", "1", NULL},
            "--mnemonic-fd 1 is not allowed"
        },
        {
            "mnemonic-fd 0 rejected with stdio",
            (const char *const[]){"./build/nsigner", "--listen", "stdio", "--mnemonic-fd", "0", NULL},
            "cannot be used with --listen stdio"
        }
    };

    size_t i;
    for (i = 0; i < sizeof(cases) / sizeof(cases[0]); ++i) {
        int cap_fd = -1;
        int status = 0;
        char cap[512];
        pid_t child = spawn_with_stdin_and_fd(cases[i].argv,
                                              "",
                                              1,
                                              NULL,
                                              -1,
                                              0,
                                              NULL,
                                              NULL,
                                              &cap_fd);

        check_condition(cases[i].name, child > 0);
        if (child <= 0) {
            continue;
        }

        check_condition("expected parse-time failure exits", wait_exit_timeout(child, 2000, &status) == 0);
        check_condition("expected parse-time failure non-zero", WIFEXITED(status) && WEXITSTATUS(status) != 0);

        if (cap_fd >= 0) {
            (void)read_capture_nonblocking(cap_fd, cap, sizeof(cap));
            check_condition("expected parse-time error message", strstr(cap, cases[i].expect_substr) != NULL);
            close(cap_fd);
        }
    }
}

static void test_lock_reunlock_falls_back_to_tui_after_fd_startup(void) {
    const char *const argv[] = {
        "./build/nsigner", "--socket-name", "nsigner_test_mi_lock", "--mnemonic-fd", "3", NULL
    };
    int stdin_write = -1;
    int cap_fd = -1;
    int status = 0;
    int fd;
    pid_t child = spawn_with_stdin_and_fd(argv,
                                          "",
                                          0,
                                          VALID_MNEMONIC "\n",
                                          3,
                                          1,
                                          &stdin_write,
                                          NULL,
                                          &cap_fd);

    check_condition("spawn --mnemonic-fd lock test child", child > 0);
    if (child <= 0) {
        return;
    }

    fd = connect_socket_retry("nsigner_test_mi_lock", 5000);
    check_condition("--mnemonic-fd lock test initial running server", fd >= 0);
    if (fd >= 0) {
        close(fd);
    }

    (void)write_full(stdin_write, "l\n" VALID_MNEMONIC "\n", strlen("l\n" VALID_MNEMONIC "\n"));
    sleep_ms(800);

    fd = connect_socket_retry("nsigner_test_mi_lock", 2000);
    check_condition("lock+reunlock via TUI keeps server alive", fd >= 0);
    if (fd >= 0) {
        close(fd);
    }

    if (stdin_write >= 0) {
        close(stdin_write);
    }

    (void)kill(child, SIGTERM);
    (void)waitpid(child, &status, 0);
    if (cap_fd >= 0) {
        close(cap_fd);
    }
}

int main(void) {
    (void)signal(SIGPIPE, SIG_IGN);

    test_stdin_happy_path();
    test_stdin_eof_without_newline();
    test_stdin_invalid_mnemonic();
    test_stdin_overflow();
    test_fd_happy_path();
    test_mutual_exclusion_and_validation_errors();
    test_lock_reunlock_falls_back_to_tui_after_fd_startup();

    if (g_failures == 0) {
        printf("ALL TESTS PASSED\n");
        return 0;
    }

    printf("TESTS FAILED: %d\n", g_failures);
    return 1;
}