#include "qjs_crypto.h"
#include "quickjs.h"
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "monocypher.h"

#include <stdint.h>
#include <stddef.h>

#if defined(_WIN32)
// ------- Windows: use BCryptGenRandom -------
#include <windows.h>
#include <bcrypt.h>

int randombytes(void *buf, size_t n) {
    NTSTATUS status = BCryptGenRandom(NULL, (PUCHAR)buf, (ULONG)n, BCRYPT_USE_SYSTEM_PREFERRED_RNG);
    return (status == 0) ? 0 : -1;
}

#elif defined(__linux__)
// ------- Linux: try getrandom, fall back to /dev/urandom -------
#include <unistd.h>
#include <sys/syscall.h>
#include <fcntl.h>
#include <errno.h>

// If we have a new enough libc and kernel, getrandom is available.
// Otherwise, we’ll do a /dev/urandom fallback.
#include <sys/stat.h>

static int randombytes_fallback(void *buf, size_t n) {
    int fd = open("/dev/urandom", O_RDONLY);
    if (fd < 0) return -1;
    ssize_t r = read(fd, buf, n);
    close(fd);
    return (r == (ssize_t)n) ? 0 : -1;
}

int randombytes(void *buf, size_t n) {
#ifdef SYS_getrandom
    // Try getrandom(2) if available
    ssize_t ret = syscall(SYS_getrandom, buf, n, 0);
    if (ret < 0) {
        // If getrandom is not supported or fails, fall back
        if (errno == ENOSYS) {
            return randombytes_fallback(buf, n);
        }
        return -1;
    }
    return (ret == (ssize_t)n) ? 0 : -1;
#else
    // getrandom not available, just fallback
    return randombytes_fallback(buf, n);
#endif
}

#else
// ------- Other Unix: read from /dev/urandom -------
#include <fcntl.h>
#include <unistd.h>

int randombytes(void *buf, size_t n) {
    int fd = open("/dev/urandom", O_RDONLY);
    if (fd < 0) return -1;
    ssize_t r = read(fd, buf, n);
    close(fd);
    return (r == (ssize_t)n) ? 0 : -1;
}
#endif

static inline void to_hex(const uint8_t *in, size_t in_len, char *out)
{
    static const char hexchars[] = "0123456789abcdef";
    for (size_t i = 0; i < in_len; i++) {
        out[2*i    ] = hexchars[(in[i] >> 4) & 0x0F];
        out[2*i + 1] = hexchars[ in[i]       & 0x0F];
    }
    out[2 * in_len] = '\0'; // null-terminate
}

static inline int nibble_from_char(char c, uint8_t *nibble)
{
    if (c >= '0' && c <= '9') { *nibble = (uint8_t)(c - '0');      return 0; }
    if (c >= 'a' && c <= 'f') { *nibble = (uint8_t)(c - 'a' + 10); return 0; }
    if (c >= 'A' && c <= 'F') { *nibble = (uint8_t)(c - 'A' + 10); return 0; }
    return -1; // invalid char
}

static inline int from_hex(const char *hex, uint8_t *out, size_t out_len)
{
    for (size_t i = 0; i < out_len; i++) {
        uint8_t hi, lo;
        if (nibble_from_char(hex[2*i],     &hi) < 0) return -1;
        if (nibble_from_char(hex[2*i + 1], &lo) < 0) return -1;
        out[i] = (uint8_t)((hi << 4) | lo);
    }
    return 0;
}

// Convert a JSValue containing a 64-character hex string into a 32-byte array.
static inline void js2crypto(JSContext *js, JSValue v, uint8_t *crypto)
{
    size_t hex_len;
    const char *hex_str = JS_ToCStringLen(js, &hex_len, v);
    if (!hex_str)
        return;

    if (hex_len != 64) {
        JS_FreeCString(js, hex_str);
        JS_ThrowTypeError(js, "js2crypto: expected 64-hex-char string");
        return;
    }

    if (from_hex(hex_str, crypto, 32) < 0) {
        JS_FreeCString(js, hex_str);
        JS_ThrowTypeError(js, "js2crypto: invalid hex encoding");
        return;
    }

    JS_FreeCString(js, hex_str);
}

static inline JSValue crypto2js(JSContext *js, const uint8_t *crypto)
{
    char hex[65];  // 32*2 + 1 for null terminator
    to_hex(crypto, 32, hex);
    return JS_NewString(js, hex);
}

JSValue js_crypto_keypair(JSContext *js, JSValue self, int argc, JSValue *argv) {
  JSValue ret = JS_NewObject(js);

  uint8_t public[32];
  uint8_t private[32];

  randombytes(private,32);

  private[0] &= 248;
  private[31] &= 127;
  private[31] |= 64;

  crypto_x25519_public_key(public,private);

  JS_SetPropertyStr(js, ret, "public", crypto2js(js, public));
  JS_SetPropertyStr(js, ret, "private", crypto2js(js,private));
  return ret;
}

JSValue js_crypto_shared(JSContext *js, JSValue self, int argc, JSValue *argv)
{
    if (argc < 1 || !JS_IsObject(argv[0])) {
        return JS_ThrowTypeError(js, "crypto.shared: expected an object argument");
    }

    JSValue obj = argv[0];

    JSValue val_pub = JS_GetPropertyStr(js, obj, "public");
    if (JS_IsException(val_pub)) {
        JS_FreeValue(js, val_pub);
        return JS_EXCEPTION;
    }

    JSValue val_priv = JS_GetPropertyStr(js, obj, "private");
    if (JS_IsException(val_priv)) {
        JS_FreeValue(js, val_pub);
        JS_FreeValue(js, val_priv);
        return JS_EXCEPTION;
    }

    uint8_t pub[32], priv[32];
    js2crypto(js, val_pub,  pub);
    js2crypto(js, val_priv, priv);

    JS_FreeValue(js, val_pub);
    JS_FreeValue(js, val_priv);

    uint8_t shared[32];
    crypto_x25519(shared, priv, pub);

    return crypto2js(js, shared);
}

JSValue js_crypto_random(JSContext *js, JSValue self, int argc, JSValue *argv)
{
    // 1) Pull 64 bits of cryptographically secure randomness
    uint64_t r;
    if (randombytes(&r, sizeof(r)) != 0) {
        // If something fails (extremely rare), throw an error
        return JS_ThrowInternalError(js, "crypto.random: unable to get random bytes");
    }

    // 2) Convert r to a double in the range [0,1). 
    //    We divide by (UINT64_MAX + 1.0) to ensure we never produce exactly 1.0.
    double val = (double)r / ((double)UINT64_MAX + 1.0);

    // 3) Return that as a JavaScript number
    return JS_NewFloat64(js, val);
}

JSValue js_crypto_encrypt_pk(JSContext *js, JSValue self, int argc, JSValue *argv)
{
    if (argc < 2) {
        return JS_ThrowTypeError(js, "crypto.encrypt_pk: expected 2 arguments (public_key, plaintext)");
    }

    uint8_t public_key[32];
    js2crypto(js, argv[0], public_key);

    size_t plaintext_len;
    uint8_t *plaintext;
    
    if (JS_IsString(argv[1])) {
        const char *str = JS_ToCStringLen(js, &plaintext_len, argv[1]);
        if (!str) return JS_EXCEPTION;
        plaintext = (uint8_t *)str;
    } else {
        plaintext = JS_GetArrayBuffer(js, &plaintext_len, argv[1]);
        if (!plaintext) {
            return JS_ThrowTypeError(js, "crypto.encrypt_pk: plaintext must be string or ArrayBuffer");
        }
    }

    // Generate ephemeral keypair
    uint8_t ephemeral_secret[32];
    uint8_t ephemeral_public[32];
    randombytes(ephemeral_secret, 32);
    ephemeral_secret[0] &= 248;
    ephemeral_secret[31] &= 127;
    ephemeral_secret[31] |= 64;
    crypto_x25519_public_key(ephemeral_public, ephemeral_secret);

    // Compute shared secret
    uint8_t shared[32];
    crypto_x25519(shared, ephemeral_secret, public_key);

    // Derive encryption key using BLAKE2b
    uint8_t key[32];
    crypto_blake2b(key, 32, shared, 32);

    // Generate random nonce
    uint8_t nonce[24];
    randombytes(nonce, 24);

    // Allocate output buffer: ephemeral_public(32) + nonce(24) + ciphertext + mac(16)
    size_t output_size = 32 + 24 + plaintext_len + 16;
    uint8_t *output = js_malloc(js, output_size);
    if (!output) {
        if (JS_IsString(argv[1])) JS_FreeCString(js, (const char *)plaintext);
        return JS_EXCEPTION;
    }

    // Copy ephemeral public key and nonce to output
    memcpy(output, ephemeral_public, 32);
    memcpy(output + 32, nonce, 24);

    // Encrypt
    crypto_aead_lock(output + 32 + 24, output + 32 + 24 + plaintext_len,
                     key, nonce, NULL, 0, plaintext, plaintext_len);

    if (JS_IsString(argv[1])) JS_FreeCString(js, (const char *)plaintext);

    // Wipe sensitive data
    crypto_wipe(ephemeral_secret, 32);
    crypto_wipe(shared, 32);
    crypto_wipe(key, 32);

    JSValue result = JS_NewArrayBufferCopy(js, output, output_size);
    js_free(js, output);
    return result;
}

JSValue js_crypto_decrypt_pk(JSContext *js, JSValue self, int argc, JSValue *argv)
{
    if (argc < 2) {
        return JS_ThrowTypeError(js, "crypto.decrypt_pk: expected 2 arguments (private_key, ciphertext)");
    }

    uint8_t private_key[32];
    js2crypto(js, argv[0], private_key);

    size_t ciphertext_len;
    uint8_t *ciphertext = JS_GetArrayBuffer(js, &ciphertext_len, argv[1]);
    if (!ciphertext) {
        return JS_ThrowTypeError(js, "crypto.decrypt_pk: ciphertext must be ArrayBuffer");
    }

    if (ciphertext_len < 32 + 24 + 16) {
        return JS_ThrowTypeError(js, "crypto.decrypt_pk: ciphertext too short");
    }

    // Extract ephemeral public key and nonce
    uint8_t ephemeral_public[32];
    uint8_t nonce[24];
    memcpy(ephemeral_public, ciphertext, 32);
    memcpy(nonce, ciphertext + 32, 24);

    // Compute shared secret
    uint8_t shared[32];
    crypto_x25519(shared, private_key, ephemeral_public);

    // Derive decryption key
    uint8_t key[32];
    crypto_blake2b(key, 32, shared, 32);

    // Decrypt
    size_t plaintext_len = ciphertext_len - 32 - 24 - 16;
    uint8_t *plaintext = js_malloc(js, plaintext_len);
    if (!plaintext) {
        crypto_wipe(shared, 32);
        crypto_wipe(key, 32);
        return JS_EXCEPTION;
    }

    int result = crypto_aead_unlock(plaintext, 
                                    ciphertext + ciphertext_len - 16,
                                    key, nonce, NULL, 0,
                                    ciphertext + 32 + 24, plaintext_len);

    crypto_wipe(shared, 32);
    crypto_wipe(key, 32);

    if (result != 0) {
        js_free(js, plaintext);
        return JS_ThrowTypeError(js, "crypto.decrypt_pk: decryption failed");
    }

    JSValue ret = JS_NewArrayBufferCopy(js, plaintext, plaintext_len);
    js_free(js, plaintext);
    return ret;
}

JSValue js_crypto_encrypt(JSContext *js, JSValue self, int argc, JSValue *argv)
{
    if (argc < 2) {
        return JS_ThrowTypeError(js, "crypto.encrypt: expected 2 arguments (key, plaintext)");
    }

    uint8_t key[32];
    js2crypto(js, argv[0], key);

    size_t plaintext_len;
    uint8_t *plaintext;
    
    if (JS_IsString(argv[1])) {
        const char *str = JS_ToCStringLen(js, &plaintext_len, argv[1]);
        if (!str) return JS_EXCEPTION;
        plaintext = (uint8_t *)str;
    } else {
        plaintext = JS_GetArrayBuffer(js, &plaintext_len, argv[1]);
        if (!plaintext) {
            return JS_ThrowTypeError(js, "crypto.encrypt: plaintext must be string or ArrayBuffer");
        }
    }

    // Generate random nonce
    uint8_t nonce[24];
    randombytes(nonce, 24);

    // Allocate output buffer: nonce(24) + ciphertext + mac(16)
    size_t output_size = 24 + plaintext_len + 16;
    uint8_t *output = js_malloc(js, output_size);
    if (!output) {
        if (JS_IsString(argv[1])) JS_FreeCString(js, (const char *)plaintext);
        return JS_EXCEPTION;
    }

    // Copy nonce to output
    memcpy(output, nonce, 24);

    // Encrypt
    crypto_aead_lock(output + 24, output + 24 + plaintext_len,
                     key, nonce, NULL, 0, plaintext, plaintext_len);

    if (JS_IsString(argv[1])) JS_FreeCString(js, (const char *)plaintext);

    JSValue result = JS_NewArrayBufferCopy(js, output, output_size);
    js_free(js, output);
    return result;
}

JSValue js_crypto_decrypt(JSContext *js, JSValue self, int argc, JSValue *argv)
{
    if (argc < 2) {
        return JS_ThrowTypeError(js, "crypto.decrypt: expected 2 arguments (key, ciphertext)");
    }

    uint8_t key[32];
    js2crypto(js, argv[0], key);

    size_t ciphertext_len;
    uint8_t *ciphertext = JS_GetArrayBuffer(js, &ciphertext_len, argv[1]);
    if (!ciphertext) {
        return JS_ThrowTypeError(js, "crypto.decrypt: ciphertext must be ArrayBuffer");
    }

    if (ciphertext_len < 24 + 16) {
        return JS_ThrowTypeError(js, "crypto.decrypt: ciphertext too short");
    }

    // Extract nonce
    uint8_t nonce[24];
    memcpy(nonce, ciphertext, 24);

    // Decrypt
    size_t plaintext_len = ciphertext_len - 24 - 16;
    uint8_t *plaintext = js_malloc(js, plaintext_len);
    if (!plaintext) {
        return JS_EXCEPTION;
    }

    int result = crypto_aead_unlock(plaintext, 
                                    ciphertext + ciphertext_len - 16,
                                    key, nonce, NULL, 0,
                                    ciphertext + 24, plaintext_len);

    if (result != 0) {
        js_free(js, plaintext);
        return JS_ThrowTypeError(js, "crypto.decrypt: decryption failed");
    }

    JSValue ret = JS_NewArrayBufferCopy(js, plaintext, plaintext_len);
    js_free(js, plaintext);
    return ret;
}

static const JSCFunctionListEntry js_crypto_funcs[] = {
  JS_CFUNC_DEF("keypair", 0, js_crypto_keypair),
  JS_CFUNC_DEF("shared", 1, js_crypto_shared),
  JS_CFUNC_DEF("random", 0, js_crypto_random),
  JS_CFUNC_DEF("encrypt_pk", 2, js_crypto_encrypt_pk),
  JS_CFUNC_DEF("decrypt_pk", 2, js_crypto_decrypt_pk),
  JS_CFUNC_DEF("encrypt", 2, js_crypto_encrypt),
  JS_CFUNC_DEF("decrypt", 2, js_crypto_decrypt),
};

JSValue js_crypto_use(JSContext *js)
{
  JSValue obj = JS_NewObject(js);
  JS_SetPropertyFunctionList(js, obj, js_crypto_funcs, sizeof(js_crypto_funcs)/sizeof(js_crypto_funcs[0]));
  return obj;
}