#define _GNU_SOURCE

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

#include <cJSON.h>
#include <nostr_core/nostr_common.h>

#include "../src/auth_envelope.h"

#define MAX_MSG_SIZE 65536

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 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 send_framed(int fd, const char *payload) {
    uint32_t len = (uint32_t)strlen(payload);
    uint32_t be_len = htonl(len);

    if (write_full(fd, &be_len, sizeof(be_len)) != 0) {
        return -1;
    }
    if (write_full(fd, payload, len) != 0) {
        return -1;
    }
    return 0;
}

static int recv_framed(int fd, char **out_payload) {
    unsigned char *buf = NULL;
    size_t used = 0;
    size_t cap = 0;

    if (out_payload == NULL) {
        return -1;
    }
    *out_payload = NULL;

    for (;;) {
        unsigned char chunk[1024];
        ssize_t n = read(fd, chunk, sizeof(chunk));
        if (n == 0) {
            break;
        }
        if (n < 0) {
            if (errno == EINTR) {
                continue;
            }
            free(buf);
            return -1;
        }

        if (used + (size_t)n > cap) {
            size_t new_cap = (cap == 0) ? 2048 : cap * 2;
            while (new_cap < used + (size_t)n) {
                new_cap *= 2;
            }
            {
                unsigned char *tmp = (unsigned char *)realloc(buf, new_cap);
                if (tmp == NULL) {
                    free(buf);
                    return -1;
                }
                buf = tmp;
                cap = new_cap;
            }
        }

        memcpy(buf + used, chunk, (size_t)n);
        used += (size_t)n;
    }

    if (used < 4) {
        free(buf);
        return -1;
    }

    {
        size_t off;
        for (off = 0; off + 4 <= used; ++off) {
            uint32_t be_len;
            uint32_t len;

            memcpy(&be_len, buf + off, sizeof(be_len));
            len = ntohl(be_len);
            if (len == 0 || len > MAX_MSG_SIZE) {
                continue;
            }
            if (off + 4U + (size_t)len > used) {
                continue;
            }
            if (buf[off + 4] != '{') {
                continue;
            }

            {
                char *payload = (char *)malloc((size_t)len + 1U);
                if (payload == NULL) {
                    free(buf);
                    return -1;
                }
                memcpy(payload, buf + off + 4, (size_t)len);
                payload[len] = '\0';
                *out_payload = payload;
                free(buf);
                return 0;
            }
        }
    }

    free(buf);
    return -1;
}

static char *build_auth_request(const unsigned char privkey[32],
                                const char *req_id,
                                const char *method) {
    cJSON *root = NULL;
    cJSON *params = NULL;
    cJSON *auth = NULL;
    char *printed = NULL;

    params = cJSON_CreateArray();
    if (params == NULL) {
        return NULL;
    }
    cJSON_AddItemToArray(params, cJSON_CreateString(""));

    if (auth_envelope_build_for_request(req_id,
                                        method,
                                        params,
                                        privkey,
                                        "qrexec-test",
                                        time(NULL),
                                        &auth) != 0) {
        cJSON_Delete(params);
        return NULL;
    }

    root = cJSON_CreateObject();
    if (root == NULL) {
        cJSON_Delete(params);
        cJSON_Delete(auth);
        return NULL;
    }

    cJSON_AddStringToObject(root, "id", req_id);
    cJSON_AddStringToObject(root, "method", method);
    cJSON_AddItemToObject(root, "params", params);
    cJSON_AddItemToObject(root, "auth", auth);

    printed = cJSON_PrintUnformatted(root);
    cJSON_Delete(root);
    return printed;
}

static int run_qrexec_once(const char *request_json, char **out_response) {
    int in_pipe[2];
    int out_pipe[2];
    pid_t child;
    int status;
    const char *mnemonic = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about\n";

    if (out_response == NULL || request_json == NULL) {
        return -1;
    }
    *out_response = NULL;

    if (pipe(in_pipe) != 0) {
        return -1;
    }
    if (pipe(out_pipe) != 0) {
        close(in_pipe[0]);
        close(in_pipe[1]);
        return -1;
    }

    child = fork();
    if (child < 0) {
        close(in_pipe[0]);
        close(in_pipe[1]);
        close(out_pipe[0]);
        close(out_pipe[1]);
        return -1;
    }

    if (child == 0) {
        int null_fd = open("/dev/null", O_WRONLY);
        dup2(in_pipe[0], STDIN_FILENO);
        dup2(out_pipe[1], STDOUT_FILENO);
        if (null_fd >= 0) {
            dup2(null_fd, STDERR_FILENO);
            close(null_fd);
        }

        close(in_pipe[0]);
        close(in_pipe[1]);
        close(out_pipe[0]);
        close(out_pipe[1]);

        (void)setenv("QREXEC_REMOTE_DOMAIN", "personal", 1);

        execl("./build/nsigner",
              "./build/nsigner",
              "--listen", "qrexec",
              "--auth", "optional",
              "--preapprove", "caller=qubes:personal,role=main",
              (char *)NULL);
        _exit(127);
    }

    close(in_pipe[0]);
    close(out_pipe[1]);

    if (write_full(in_pipe[1], mnemonic, strlen(mnemonic)) != 0) {
        close(in_pipe[1]);
        close(out_pipe[0]);
        (void)kill(child, SIGKILL);
        (void)waitpid(child, &status, 0);
        return -1;
    }

    /* Give prompt_load_mnemonic()/fgets() time to consume only the mnemonic line
     * before we write framed request bytes to the same stdin stream. */
    {
        struct timespec ts;
        ts.tv_sec = 0;
        ts.tv_nsec = 250 * 1000 * 1000;
        (void)nanosleep(&ts, NULL);
    }

    if (send_framed(in_pipe[1], request_json) != 0) {
        close(in_pipe[1]);
        close(out_pipe[0]);
        (void)kill(child, SIGKILL);
        (void)waitpid(child, &status, 0);
        return -1;
    }
    close(in_pipe[1]);

    if (recv_framed(out_pipe[0], out_response) != 0) {
        close(out_pipe[0]);
        (void)kill(child, SIGKILL);
        (void)waitpid(child, &status, 0);
        return -1;
    }

    close(out_pipe[0]);
    if (waitpid(child, &status, 0) <= 0) {
        free(*out_response);
        *out_response = NULL;
        return -1;
    }

    return 0;
}

int main(void) {
    char *resp = NULL;
    char *auth_req = NULL;
    unsigned char privkey[32] = {
        0x01,0x00,0x00,0x00,0x00,0x00,0x00,0x00,
        0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,
        0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,
        0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00
    };

    check_condition("nostr_init for auth signing", nostr_init() == 0);

    /* no auth => legacy qubes:personal preapprove should match */
    if (run_qrexec_once("{\"id\":\"1\",\"method\":\"nostr_get_public_key\",\"params\":[\"\"]}", &resp) == 0) {
        check_condition("qrexec optional/no-auth still allows legacy qubes caller preapprove",
                        strstr(resp, "\"result\":") != NULL);
        free(resp);
        resp = NULL;
    } else {
        check_condition("qrexec optional/no-auth request roundtrip", 0);
    }

    auth_req = build_auth_request(privkey, "2", "nostr_get_public_key");
    check_condition("build qrexec auth request", auth_req != NULL);
    if (auth_req != NULL) {
        if (run_qrexec_once(auth_req, &resp) == 0) {
            check_condition("qrexec optional/auth request uses composite identity and misses legacy preapprove",
                            strstr(resp, "\"code\":2001") != NULL && strstr(resp, "policy_denied") != NULL);
            free(resp);
            resp = NULL;
        } else {
            check_condition("qrexec optional/auth request roundtrip", 0);
        }
        free(auth_req);
        auth_req = NULL;
    }

    auth_req = build_auth_request(privkey, "3", "nostr_get_public_key");
    check_condition("build second qrexec auth request", auth_req != NULL);
    if (auth_req != NULL) {
        cJSON *root = cJSON_Parse(auth_req);
        char *tampered = NULL;
        cJSON *method_item;

        if (root != NULL) {
            method_item = cJSON_GetObjectItemCaseSensitive(root, "method");
            if (cJSON_IsString(method_item)) {
                cJSON_SetValuestring(method_item, "nostr_sign_event");
            }
            tampered = cJSON_PrintUnformatted(root);
            cJSON_Delete(root);
        }

        if (tampered != NULL && run_qrexec_once(tampered, &resp) == 0) {
            check_condition("qrexec optional/malformed-auth rejects with auth-layer code",
                            strstr(resp, "\"code\":2015") != NULL && strstr(resp, "auth_envelope_mismatch") != NULL);
            free(resp);
            resp = NULL;
        } else {
            check_condition("qrexec optional/malformed-auth request roundtrip", 0);
        }

        free(tampered);
        free(auth_req);
    }

    nostr_cleanup();

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

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