#include "cell.h"
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "monocypher.h"

/*
  Crypto Module Documentation
  
  This module provides cryptographic functions using the Monocypher library.
  All inputs and outputs are Blobs.
  
  Functions:
  
  - keypair() -> { public: Blob(256 bits), private: Blob(256 bits) }
    Generates a new random X25519 keypair.
    
  - shared(public_key, private_key) -> Blob(256 bits)
    Computes a shared secret from your private key and another's public key (X25519).
    Input keys must be 256 bits (32 bytes).
    
  - blake2(data, [hash_size_bytes=32]) -> Blob
    Computes the BLAKE2b hash of the data.
    Default hash size is 32 bytes (256 bits). Supports 1-64 bytes.
    
  - sign(secret_key, message) -> Blob(512 bits)
    Signs a message using EdDSA. 
    secret_key must be 512 bits (64 bytes). 
    (Note: If you have a 32-byte seed, extend it first or use appropriate key generation).
    Returns a 64-byte signature.
    
  - verify(signature, public_key, message) -> bool
    Verifies an EdDSA signature.
    signature: 512 bits (64 bytes).
    public_key: 256 bits (32 bytes).
    Returns true if valid, false otherwise.
    
  - lock(key, nonce, message, [ad]) -> Blob
    Encrypts and authenticates a message using XChaCha20-Poly1305.
    key: 256 bits (32 bytes).
    nonce: 192 bits (24 bytes).
    ad: Optional associated data (Blob).
    Returns a blob containing the ciphertext followed by the 16-byte MAC.
    
  - unlock(key, nonce, ciphertext_with_mac, [ad]) -> Blob or null
    Decrypts and verifies a message.
    key: 256 bits (32 bytes).
    nonce: 192 bits (24 bytes).
    ciphertext_with_mac: Must include the 16-byte MAC at the end.
    ad: Optional associated data (Blob).
    Returns the plaintext Blob if successful, or null if verification fails.
*/

// Helper to get blob data and check exact bit length
static void *get_blob_check_bits(JSContext *js, JSValue val, size_t expected_bits, const char *name) {
    size_t bits;
    void* result = js_get_blob_data_bits(js, &bits, val);
    if (result == -1) {
        return NULL; // Exception already thrown by js_get_blob_data_bits
    }

    if (bits != expected_bits) {
        JS_ThrowTypeError(js, "%s: expected %zu bits, got %zu", name, expected_bits, bits);
        return NULL;
    }
    return result;
}

// Helper to get any blob data (checking it is a stoned blob)
static void *get_blob_any(JSContext *js, JSValue val, size_t *out_bits, const char *name) {
    void *result = js_get_blob_data_bits(js, out_bits, val);
    if (result == -1)
        return NULL;
    return result;
}


JSValue js_crypto_shared(JSContext *js, JSValue self, int argc, JSValue *argv)
{
    if (argc < 2) {
        return JS_ThrowTypeError(js, "crypto.shared: expected public_key, private_key");
    }

    uint8_t *pub = get_blob_check_bits(js, argv[0], 256, "crypto.shared public_key");
    if (!pub) return JS_EXCEPTION;

    uint8_t *priv = get_blob_check_bits(js, argv[1], 256, "crypto.shared private_key");
    if (!priv) return JS_EXCEPTION;

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

    return js_new_blob_stoned_copy(js, shared, 32);
}

JSValue js_crypto_blake2(JSContext *js, JSValue self, int argc, JSValue *argv)
{
    if (argc < 1)
        return JS_ThrowTypeError(js, "crypto.blake2: expected data blob");
    
    size_t data_bits;
    uint8_t *data = get_blob_any(js, argv[0], &data_bits, "crypto.blake2 data");
    if (!data) return JS_EXCEPTION;
    
    int32_t hash_len = 32;
    if (argc > 1) {
        if (JS_ToInt32(js, &hash_len, argv[1]))
            return JS_EXCEPTION;
        if (hash_len < 1 || hash_len > 64)
            return JS_ThrowRangeError(js, "crypto.blake2: hash length must be between 1 and 64 bytes");
    }
    
    uint8_t hash[64];
    // Use (bits + 7) / 8 to get byte length covering all bits
    crypto_blake2b(hash, hash_len, data, (data_bits + 7) / 8);
    
    return js_new_blob_stoned_copy(js, hash, hash_len);
}

JSValue js_crypto_sign(JSContext *js, JSValue self, int argc, JSValue *argv) {
    if (argc < 2) return JS_ThrowTypeError(js, "crypto.sign: expected secret_key, message");
    
    uint8_t *sk = get_blob_check_bits(js, argv[0], 512, "crypto.sign secret_key");
    if (!sk) return JS_EXCEPTION;
    
    size_t msg_bits;
    uint8_t *msg = get_blob_any(js, argv[1], &msg_bits, "crypto.sign message");
    if (!msg) return JS_EXCEPTION;
    
    uint8_t sig[64];
    crypto_eddsa_sign(sig, sk, msg, (msg_bits + 7) / 8);
    
    return js_new_blob_stoned_copy(js, sig, 64);
}

JSValue js_crypto_verify(JSContext *js, JSValue self, int argc, JSValue *argv) {
    if (argc < 3) return JS_ThrowTypeError(js, "crypto.verify: expected signature, public_key, message");
    
    uint8_t *sig = get_blob_check_bits(js, argv[0], 512, "crypto.verify signature");
    if (!sig) return JS_EXCEPTION;
    
    uint8_t *pk = get_blob_check_bits(js, argv[1], 256, "crypto.verify public_key");
    if (!pk) return JS_EXCEPTION;
    
    size_t msg_bits;
    uint8_t *msg = get_blob_any(js, argv[2], &msg_bits, "crypto.verify message");
    if (!msg) return JS_EXCEPTION;
    
    int ret = crypto_eddsa_check(sig, pk, msg, (msg_bits + 7) / 8);
    return JS_NewBool(js, ret == 0);
}

JSValue js_crypto_lock(JSContext *js, JSValue self, int argc, JSValue *argv) {
    if (argc < 3) return JS_ThrowTypeError(js, "crypto.lock: expected key, nonce, message, [ad]");
    
    uint8_t *key = get_blob_check_bits(js, argv[0], 256, "crypto.lock key");
    if (!key) return JS_EXCEPTION;
    
    uint8_t *nonce = get_blob_check_bits(js, argv[1], 192, "crypto.lock nonce");
    if (!nonce) return JS_EXCEPTION;
    
    size_t msg_bits;
    uint8_t *msg = get_blob_any(js, argv[2], &msg_bits, "crypto.lock message");
    if (!msg) return JS_EXCEPTION;
    size_t msg_len = (msg_bits + 7) / 8;
    
    size_t ad_len = 0;
    uint8_t *ad = NULL;
    if (argc > 3 && !JS_IsNull(argv[3])) {
        size_t ad_bits;
        ad = get_blob_any(js, argv[3], &ad_bits, "crypto.lock ad");
        if (!ad) return JS_EXCEPTION;
        ad_len = (ad_bits + 7) / 8;
    }
    
    size_t out_len = msg_len + 16;
    uint8_t *out = malloc(out_len);
    if (!out) return JS_ThrowOutOfMemory(js);
    
    // Output: [Ciphertext (msg_len)] [MAC (16)]
    crypto_aead_lock(out, out + msg_len, key, nonce, ad, ad_len, msg, msg_len);
    
    JSValue ret = js_new_blob_stoned_copy(js, out, out_len);
    free(out);
    return ret;
}

JSValue js_crypto_unlock(JSContext *js, JSValue self, int argc, JSValue *argv) {
    if (argc < 3) return JS_ThrowTypeError(js, "crypto.unlock: expected key, nonce, ciphertext, [ad]");
    
    uint8_t *key = get_blob_check_bits(js, argv[0], 256, "crypto.unlock key");
    if (!key) return JS_EXCEPTION;
    
    uint8_t *nonce = get_blob_check_bits(js, argv[1], 192, "crypto.unlock nonce");
    if (!nonce) return JS_EXCEPTION;
    
    size_t cipher_bits;
    uint8_t *cipher = get_blob_any(js, argv[2], &cipher_bits, "crypto.unlock ciphertext");
    if (!cipher) return JS_EXCEPTION;
    
    size_t cipher_len = (cipher_bits + 7) / 8;
    if (cipher_len < 16) return JS_ThrowTypeError(js, "crypto.unlock: ciphertext too short (min 16 bytes)");
    
    size_t msg_len = cipher_len - 16;
    
    size_t ad_len = 0;
    uint8_t *ad = NULL;
    if (argc > 3 && !JS_IsNull(argv[3])) {
        size_t ad_bits;
        ad = get_blob_any(js, argv[3], &ad_bits, "crypto.unlock ad");
        if (!ad) return JS_EXCEPTION;
        ad_len = (ad_bits + 7) / 8;
    }
    
    uint8_t *out = malloc(msg_len > 0 ? msg_len : 1);
    if (!out) return JS_ThrowOutOfMemory(js);
    
    // MAC is at cipher + msg_len
    const uint8_t *mac = cipher + msg_len;
    
    if (crypto_aead_unlock(out, mac, key, nonce, ad, ad_len, cipher, msg_len) != 0) {
        free(out);
        return JS_NULL;
    }
    
    JSValue ret = js_new_blob_stoned_copy(js, out, msg_len);
    free(out);
    return ret;
}

static const JSCFunctionListEntry js_crypto_funcs[] = {
  JS_CFUNC_DEF("shared", 2, js_crypto_shared),
  JS_CFUNC_DEF("blake2", 2, js_crypto_blake2),
  JS_CFUNC_DEF("sign", 2, js_crypto_sign),
  JS_CFUNC_DEF("verify", 3, js_crypto_verify),
  JS_CFUNC_DEF("lock", 3, js_crypto_lock),
  JS_CFUNC_DEF("unlock", 3, js_crypto_unlock),
};

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;
}