/*
 * LITLS -- Lua binding for cryptographic primitives
 * Registers as "litls.core" via luaopen_litls_core
 */

#include <alloca.h>
#include <lauxlib.h>
#include <lua.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include "litls.h"
#include "tls13/tls13.h"

/* Max size for stack-allocated VLA buffers in Lua bindings.
 * Inputs larger than this use heap allocation. */
#define LITLS_LUA_STACK_MAX 65536

/* ── Base64 ────────────────────────────────────────────── */

static int lua_litls_base64_encode(lua_State *L) {
    size_t in_len;
    const char *in = luaL_checklstring(L, 1, &in_len);
    size_t out_cap = ((in_len + 2) / 3) * 4 + 1;
    char *buf      = (out_cap <= LITLS_LUA_STACK_MAX) ? (char *)alloca(out_cap) : (char *)malloc(out_cap);
    if (!buf)
        return luaL_error(L, "out of memory");
    size_t n = litls_base64_encode((const uint8_t *)in, in_len, buf, out_cap);
    lua_pushlstring(L, buf, n);
    if (out_cap > LITLS_LUA_STACK_MAX)
        free(buf);
    return 1;
}

static int lua_litls_base64_decode(lua_State *L) {
    size_t in_len;
    const char *in = luaL_checklstring(L, 1, &in_len);
    size_t out_len = (in_len * 3) / 4 + 4;
    uint8_t *buf   = (out_len <= LITLS_LUA_STACK_MAX) ? (uint8_t *)alloca(out_len) : (uint8_t *)malloc(out_len);
    if (!buf)
        return luaL_error(L, "out of memory");
    if (litls_base64_decode(in, in_len, buf, &out_len) != 0) {
        if (out_len > LITLS_LUA_STACK_MAX)
            free(buf);
        lua_pushnil(L);
        lua_pushstring(L, "invalid base64 input");
        return 2;
    }
    lua_pushlstring(L, (const char *)buf, out_len);
    if ((in_len * 3) / 4 + 4 > LITLS_LUA_STACK_MAX)
        free(buf);
    return 1;
}

static int lua_litls_base64url_encode(lua_State *L) {
    size_t in_len;
    const char *in = luaL_checklstring(L, 1, &in_len);
    size_t out_cap = ((in_len + 2) / 3) * 4 + 1;
    char *buf      = (out_cap <= LITLS_LUA_STACK_MAX) ? (char *)alloca(out_cap) : (char *)malloc(out_cap);
    if (!buf)
        return luaL_error(L, "out of memory");
    size_t n = litls_base64url_encode((const uint8_t *)in, in_len, buf, out_cap);
    lua_pushlstring(L, buf, n);
    if (out_cap > LITLS_LUA_STACK_MAX)
        free(buf);
    return 1;
}

static int lua_litls_base64url_decode(lua_State *L) {
    size_t in_len;
    const char *in = luaL_checklstring(L, 1, &in_len);
    size_t out_cap = (in_len * 3) / 4 + 4;
    size_t out_len = out_cap;
    uint8_t *buf   = (out_cap <= LITLS_LUA_STACK_MAX) ? (uint8_t *)alloca(out_cap) : (uint8_t *)malloc(out_cap);
    if (!buf)
        return luaL_error(L, "out of memory");
    if (litls_base64url_decode(in, in_len, buf, &out_len) != 0) {
        if (out_cap > LITLS_LUA_STACK_MAX)
            free(buf);
        lua_pushnil(L);
        lua_pushstring(L, "invalid base64url input");
        return 2;
    }
    lua_pushlstring(L, (const char *)buf, out_len);
    if (out_cap > LITLS_LUA_STACK_MAX)
        free(buf);
    return 1;
}

/* ── RNG ───────────────────────────────────────────────── */

static int lua_litls_random_bytes(lua_State *L) {
    int n = (int)luaL_checkinteger(L, 1);
    if (n < 0 || n > 65536) {
        lua_pushnil(L);
        lua_pushstring(L, "size must be 0-65536");
        return 2;
    }
    uint8_t buf[n > 0 ? n : 1];
    if (litls_random_bytes(buf, (size_t)n) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "getrandom failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)buf, (size_t)n);
    return 1;
}

/* ── Constant-time compare ─────────────────────────────── */

static int lua_litls_secure_memcmp(lua_State *L) {
    size_t a_len, b_len;
    const char *a = luaL_checklstring(L, 1, &a_len);
    const char *b = luaL_checklstring(L, 2, &b_len);
    if (a_len != b_len) {
        lua_pushboolean(L, 0);
        return 1;
    }
    lua_pushboolean(L, litls_secure_memcmp((const uint8_t *)a, (const uint8_t *)b, a_len) == 0);
    return 1;
}

/* ── SHA-1 ─────────────────────────────────────────────── */

static int lua_litls_sha1(lua_State *L) {
    size_t len;
    const char *data = luaL_checklstring(L, 1, &len);
    uint8_t out[20];
    litls_sha1((const uint8_t *)data, len, out);
    lua_pushlstring(L, (const char *)out, 20);
    return 1;
}

/* ── SHA-2 family ──────────────────────────────────────── */

static int lua_litls_sha256(lua_State *L) {
    size_t len;
    const char *data = luaL_checklstring(L, 1, &len);
    uint8_t out[32];
    litls_sha256((const uint8_t *)data, len, out);
    lua_pushlstring(L, (const char *)out, 32);
    return 1;
}

static int lua_litls_sha384(lua_State *L) {
    size_t len;
    const char *data = luaL_checklstring(L, 1, &len);
    uint8_t out[48];
    litls_sha384((const uint8_t *)data, len, out);
    lua_pushlstring(L, (const char *)out, 48);
    return 1;
}

static int lua_litls_sha512(lua_State *L) {
    size_t len;
    const char *data = luaL_checklstring(L, 1, &len);
    uint8_t out[64];
    litls_sha512((const uint8_t *)data, len, out);
    lua_pushlstring(L, (const char *)out, 64);
    return 1;
}

/* ── HMAC ──────────────────────────────────────────────── */

static int lua_litls_hmac_sha256(lua_State *L) {
    size_t key_len, msg_len;
    const char *key = luaL_checklstring(L, 1, &key_len);
    const char *msg = luaL_checklstring(L, 2, &msg_len);
    uint8_t out[32];
    litls_hmac_sha256((const uint8_t *)key, key_len, (const uint8_t *)msg, msg_len, out);
    lua_pushlstring(L, (const char *)out, 32);
    return 1;
}

static int lua_litls_hmac_sha384(lua_State *L) {
    size_t key_len, msg_len;
    const char *key = luaL_checklstring(L, 1, &key_len);
    const char *msg = luaL_checklstring(L, 2, &msg_len);
    uint8_t out[48];
    litls_hmac_sha384((const uint8_t *)key, key_len, (const uint8_t *)msg, msg_len, out);
    lua_pushlstring(L, (const char *)out, 48);
    return 1;
}

/* ── HKDF ──────────────────────────────────────────────── */

static int lua_litls_hkdf_extract_sha256(lua_State *L) {
    size_t salt_len, ikm_len;
    const char *salt = luaL_checklstring(L, 1, &salt_len);
    const char *ikm  = luaL_checklstring(L, 2, &ikm_len);
    uint8_t prk[32];
    litls_hkdf_extract_sha256((const uint8_t *)salt, salt_len, (const uint8_t *)ikm, ikm_len, prk);
    lua_pushlstring(L, (const char *)prk, 32);
    return 1;
}

static int lua_litls_hkdf_expand_sha256(lua_State *L) {
    size_t prk_len, info_len;
    const char *prk  = luaL_checklstring(L, 1, &prk_len);
    const char *info = luaL_checklstring(L, 2, &info_len);
    int okm_len      = (int)luaL_checkinteger(L, 3);
    if (prk_len < 32 || okm_len < 0 || okm_len > 8160) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid prk or output length");
        return 2;
    }
    uint8_t okm[okm_len > 0 ? okm_len : 1];
    if (litls_hkdf_expand_sha256((const uint8_t *)prk, (const uint8_t *)info, info_len, okm, (size_t)okm_len) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "HKDF expand failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)okm, (size_t)okm_len);
    return 1;
}

static int lua_litls_hkdf_extract_sha384(lua_State *L) {
    size_t salt_len, ikm_len;
    const char *salt = luaL_checklstring(L, 1, &salt_len);
    const char *ikm  = luaL_checklstring(L, 2, &ikm_len);
    uint8_t prk[48];
    litls_hkdf_extract_sha384((const uint8_t *)salt, salt_len, (const uint8_t *)ikm, ikm_len, prk);
    lua_pushlstring(L, (const char *)prk, 48);
    return 1;
}

static int lua_litls_hkdf_expand_sha384(lua_State *L) {
    size_t prk_len, info_len;
    const char *prk  = luaL_checklstring(L, 1, &prk_len);
    const char *info = luaL_checklstring(L, 2, &info_len);
    int okm_len      = (int)luaL_checkinteger(L, 3);
    if (prk_len < 48 || okm_len < 0 || okm_len > 12240) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid prk or output length");
        return 2;
    }
    uint8_t okm[okm_len > 0 ? okm_len : 1];
    if (litls_hkdf_expand_sha384((const uint8_t *)prk, (const uint8_t *)info, info_len, okm, (size_t)okm_len) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "HKDF expand failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)okm, (size_t)okm_len);
    return 1;
}

/* ── ChaCha20 ──────────────────────────────────────────── */

static int lua_litls_chacha20_block(lua_State *L) {
    size_t key_len, nonce_len;
    const char *key   = luaL_checklstring(L, 1, &key_len);
    const char *nonce = luaL_checklstring(L, 2, &nonce_len);
    int counter       = (int)luaL_checkinteger(L, 3);
    if (key_len != 32 || nonce_len != 12) {
        lua_pushnil(L);
        lua_pushstring(L, "key must be 32 bytes, nonce 12 bytes");
        return 2;
    }
    uint8_t out[64];
    litls_chacha20_block((const uint8_t *)key, (const uint8_t *)nonce, (uint32_t)counter, out);
    lua_pushlstring(L, (const char *)out, 64);
    return 1;
}

static int lua_litls_chacha20_encrypt(lua_State *L) {
    size_t key_len, nonce_len, in_len;
    const char *key   = luaL_checklstring(L, 1, &key_len);
    const char *nonce = luaL_checklstring(L, 2, &nonce_len);
    int counter       = (int)luaL_checkinteger(L, 3);
    const char *in    = luaL_checklstring(L, 4, &in_len);
    if (key_len != 32 || nonce_len != 12) {
        lua_pushnil(L);
        lua_pushstring(L, "key must be 32 bytes, nonce 12 bytes");
        return 2;
    }
    int heap     = in_len > LITLS_LUA_STACK_MAX;
    uint8_t *out = heap ? (uint8_t *)malloc(in_len > 0 ? in_len : 1) : (uint8_t *)alloca(in_len > 0 ? in_len : 1);
    if (!out)
        return luaL_error(L, "out of memory");
    litls_chacha20_encrypt((const uint8_t *)key, (const uint8_t *)nonce, (uint32_t)counter, (const uint8_t *)in, in_len,
                           out);
    lua_pushlstring(L, (const char *)out, in_len);
    if (heap)
        free(out);
    return 1;
}

/* ── Poly1305 ──────────────────────────────────────────── */

static int lua_litls_poly1305_auth(lua_State *L) {
    size_t key_len, msg_len;
    const char *key = luaL_checklstring(L, 1, &key_len);
    const char *msg = luaL_checklstring(L, 2, &msg_len);
    if (key_len != 32) {
        lua_pushnil(L);
        lua_pushstring(L, "key must be 32 bytes");
        return 2;
    }
    uint8_t tag[16];
    litls_poly1305_auth((const uint8_t *)key, (const uint8_t *)msg, msg_len, tag);
    lua_pushlstring(L, (const char *)tag, 16);
    return 1;
}

/* ── ChaCha20-Poly1305 AEAD ────────────────────────────── */

static int lua_litls_chacha20_poly1305_encrypt(lua_State *L) {
    size_t key_len, nonce_len, aad_len, pt_len;
    const char *key   = luaL_checklstring(L, 1, &key_len);
    const char *nonce = luaL_checklstring(L, 2, &nonce_len);
    const char *aad   = luaL_checklstring(L, 3, &aad_len);
    const char *pt    = luaL_checklstring(L, 4, &pt_len);
    if (key_len != 32 || nonce_len != 12) {
        lua_pushnil(L);
        lua_pushstring(L, "key must be 32 bytes, nonce 12 bytes");
        return 2;
    }
    int heap    = pt_len > LITLS_LUA_STACK_MAX;
    uint8_t *ct = heap ? (uint8_t *)malloc(pt_len > 0 ? pt_len : 1) : (uint8_t *)alloca(pt_len > 0 ? pt_len : 1);
    if (!ct)
        return luaL_error(L, "out of memory");
    uint8_t tag[16];
    litls_chacha20_poly1305_encrypt((const uint8_t *)key, (const uint8_t *)nonce, (const uint8_t *)aad, aad_len,
                                    (const uint8_t *)pt, pt_len, ct, tag);
    lua_pushlstring(L, (const char *)ct, pt_len);
    lua_pushlstring(L, (const char *)tag, 16);
    if (heap)
        free(ct);
    return 2;
}

static int lua_litls_chacha20_poly1305_decrypt(lua_State *L) {
    size_t key_len, nonce_len, aad_len, ct_len, tag_len;
    const char *key   = luaL_checklstring(L, 1, &key_len);
    const char *nonce = luaL_checklstring(L, 2, &nonce_len);
    const char *aad   = luaL_checklstring(L, 3, &aad_len);
    const char *ct    = luaL_checklstring(L, 4, &ct_len);
    const char *tag   = luaL_checklstring(L, 5, &tag_len);
    if (key_len != 32 || nonce_len != 12 || tag_len != 16) {
        lua_pushnil(L);
        lua_pushstring(L, "key 32B, nonce 12B, tag 16B required");
        return 2;
    }
    int heap    = ct_len > LITLS_LUA_STACK_MAX;
    uint8_t *pt = heap ? (uint8_t *)malloc(ct_len > 0 ? ct_len : 1) : (uint8_t *)alloca(ct_len > 0 ? ct_len : 1);
    if (!pt)
        return luaL_error(L, "out of memory");
    if (litls_chacha20_poly1305_decrypt((const uint8_t *)key, (const uint8_t *)nonce, (const uint8_t *)aad, aad_len,
                                        (const uint8_t *)ct, ct_len, (const uint8_t *)tag, pt) != 0) {
        if (heap)
            free(pt);
        lua_pushnil(L);
        lua_pushstring(L, "authentication failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)pt, ct_len);
    if (heap)
        free(pt);
    return 1;
}

/* ── AES-GCM ──────────────────────────────────────────── */

static int lua_litls_aes_gcm_encrypt(lua_State *L, int key_size) {
    size_t key_len, nonce_len, aad_len, pt_len;
    const char *key   = luaL_checklstring(L, 1, &key_len);
    const char *nonce = luaL_checklstring(L, 2, &nonce_len);
    const char *aad   = luaL_checklstring(L, 3, &aad_len);
    const char *pt    = luaL_checklstring(L, 4, &pt_len);
    if ((int)key_len != key_size || nonce_len != 12) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid key or nonce length");
        return 2;
    }
    int heap    = pt_len > LITLS_LUA_STACK_MAX;
    uint8_t *ct = heap ? (uint8_t *)malloc(pt_len > 0 ? pt_len : 1) : (uint8_t *)alloca(pt_len > 0 ? pt_len : 1);
    if (!ct)
        return luaL_error(L, "out of memory");
    uint8_t tag[16];
    int rc;
    if (key_size == 16)
        rc = litls_aes128_gcm_encrypt((const uint8_t *)key, (const uint8_t *)nonce, (const uint8_t *)aad, aad_len,
                                      (const uint8_t *)pt, pt_len, ct, tag);
    else
        rc = litls_aes256_gcm_encrypt((const uint8_t *)key, (const uint8_t *)nonce, (const uint8_t *)aad, aad_len,
                                      (const uint8_t *)pt, pt_len, ct, tag);
    if (rc != 0) {
        if (heap)
            free(ct);
        lua_pushnil(L);
        lua_pushstring(L, "AES-GCM encrypt failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)ct, pt_len);
    lua_pushlstring(L, (const char *)tag, 16);
    if (heap)
        free(ct);
    return 2;
}

static int lua_litls_aes_gcm_decrypt(lua_State *L, int key_size) {
    size_t key_len, nonce_len, aad_len, ct_len, tag_len;
    const char *key   = luaL_checklstring(L, 1, &key_len);
    const char *nonce = luaL_checklstring(L, 2, &nonce_len);
    const char *aad   = luaL_checklstring(L, 3, &aad_len);
    const char *ct    = luaL_checklstring(L, 4, &ct_len);
    const char *tag   = luaL_checklstring(L, 5, &tag_len);
    if ((int)key_len != key_size || nonce_len != 12 || tag_len != 16) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid key, nonce, or tag length");
        return 2;
    }
    int heap        = ct_len > LITLS_LUA_STACK_MAX;
    uint8_t *pt_out = heap ? (uint8_t *)malloc(ct_len > 0 ? ct_len : 1) : (uint8_t *)alloca(ct_len > 0 ? ct_len : 1);
    if (!pt_out)
        return luaL_error(L, "out of memory");
    int rc;
    if (key_size == 16)
        rc = litls_aes128_gcm_decrypt((const uint8_t *)key, (const uint8_t *)nonce, (const uint8_t *)aad, aad_len,
                                      (const uint8_t *)ct, ct_len, (const uint8_t *)tag, pt_out);
    else
        rc = litls_aes256_gcm_decrypt((const uint8_t *)key, (const uint8_t *)nonce, (const uint8_t *)aad, aad_len,
                                      (const uint8_t *)ct, ct_len, (const uint8_t *)tag, pt_out);
    if (rc != 0) {
        if (heap)
            free(pt_out);
        lua_pushnil(L);
        lua_pushstring(L, "authentication failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)pt_out, ct_len);
    if (heap)
        free(pt_out);
    return 1;
}

static int lua_litls_aes128_gcm_encrypt(lua_State *L) {
    return lua_litls_aes_gcm_encrypt(L, 16);
}
static int lua_litls_aes128_gcm_decrypt(lua_State *L) {
    return lua_litls_aes_gcm_decrypt(L, 16);
}
static int lua_litls_aes256_gcm_encrypt(lua_State *L) {
    return lua_litls_aes_gcm_encrypt(L, 32);
}
static int lua_litls_aes256_gcm_decrypt(lua_State *L) {
    return lua_litls_aes_gcm_decrypt(L, 32);
}

/* ── X25519 ────────────────────────────────────────────── */

static int lua_litls_x25519(lua_State *L) {
    size_t priv_len, pub_len;
    const char *priv = luaL_checklstring(L, 1, &priv_len);
    const char *pub  = luaL_checklstring(L, 2, &pub_len);
    if (priv_len != 32 || pub_len != 32) {
        lua_pushnil(L);
        lua_pushstring(L, "keys must be 32 bytes");
        return 2;
    }
    uint8_t shared[32];
    if (litls_x25519(shared, (const uint8_t *)priv, (const uint8_t *)pub) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "X25519 failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)shared, 32);
    return 1;
}

static int lua_litls_x25519_base(lua_State *L) {
    size_t priv_len;
    const char *priv = luaL_checklstring(L, 1, &priv_len);
    if (priv_len != 32) {
        lua_pushnil(L);
        lua_pushstring(L, "private key must be 32 bytes");
        return 2;
    }
    uint8_t pub[32];
    litls_x25519_base(pub, (const uint8_t *)priv);
    lua_pushlstring(L, (const char *)pub, 32);
    return 1;
}

/* ── Ed25519 ───────────────────────────────────────────── */

static int lua_litls_ed25519_keypair(lua_State *L) {
    size_t seed_len;
    const char *seed = luaL_checklstring(L, 1, &seed_len);
    if (seed_len != 32) {
        lua_pushnil(L);
        lua_pushstring(L, "seed must be 32 bytes");
        return 2;
    }
    uint8_t pk[32], sk[64];
    litls_ed25519_keypair(pk, sk, (const uint8_t *)seed);
    lua_pushlstring(L, (const char *)pk, 32);
    lua_pushlstring(L, (const char *)sk, 64);
    return 2;
}

static int lua_litls_ed25519_sign(lua_State *L) {
    size_t msg_len, pk_len, sk_len;
    const char *msg = luaL_checklstring(L, 1, &msg_len);
    const char *pk  = luaL_checklstring(L, 2, &pk_len);
    const char *sk  = luaL_checklstring(L, 3, &sk_len);
    if (pk_len != 32 || sk_len != 64) {
        lua_pushnil(L);
        lua_pushstring(L, "pk 32B, sk 64B required");
        return 2;
    }
    uint8_t sig[64];
    litls_ed25519_sign(sig, (const uint8_t *)msg, msg_len, (const uint8_t *)pk, (const uint8_t *)sk);
    lua_pushlstring(L, (const char *)sig, 64);
    return 1;
}

static int lua_litls_ed25519_verify(lua_State *L) {
    size_t sig_len, msg_len, pk_len;
    const char *sig = luaL_checklstring(L, 1, &sig_len);
    const char *msg = luaL_checklstring(L, 2, &msg_len);
    const char *pk  = luaL_checklstring(L, 3, &pk_len);
    if (sig_len != 64 || pk_len != 32) {
        lua_pushboolean(L, 0);
        return 1;
    }
    int ok = litls_ed25519_verify((const uint8_t *)sig, (const uint8_t *)msg, msg_len, (const uint8_t *)pk);
    lua_pushboolean(L, ok == 0);
    return 1;
}

static int lua_litls_openssh_parse_ed25519(lua_State *L) {
    size_t pem_len;
    const char *pem = luaL_checklstring(L, 1, &pem_len);

    /* Strip PEM armor: find base64 content between header/footer lines */
    const char *begin     = "-----BEGIN OPENSSH PRIVATE KEY-----";
    const char *end       = "-----END OPENSSH PRIVATE KEY-----";
    const char *b64_start = strstr(pem, begin);
    if (!b64_start) {
        lua_pushnil(L);
        lua_pushstring(L, "not an OpenSSH private key");
        return 2;
    }
    b64_start += strlen(begin);
    while (b64_start < pem + pem_len && (*b64_start == '\n' || *b64_start == '\r'))
        b64_start++;
    const char *b64_end = strstr(b64_start, end);
    if (!b64_end) {
        lua_pushnil(L);
        lua_pushstring(L, "missing PEM footer");
        return 2;
    }
    size_t b64_len = b64_end - b64_start;

    /* Strip whitespace from base64 content (PEM wraps lines) */
    char *b64_clean      = (char *)lua_newuserdata(L, b64_len);
    size_t b64_clean_len = 0;
    for (size_t i = 0; i < b64_len; i++) {
        char c = b64_start[i];
        if (c != '\n' && c != '\r' && c != ' ' && c != '\t')
            b64_clean[b64_clean_len++] = c;
    }

    /* Base64 decode */
    size_t raw_cap = b64_clean_len; /* decoded is always smaller */
    uint8_t *raw   = (uint8_t *)lua_newuserdata(L, raw_cap);
    size_t raw_len = raw_cap;

    uint8_t seed[32], pk[32];
    const char *err = NULL;
    int success     = 0;

    if (litls_base64_decode(b64_clean, b64_clean_len, raw, &raw_len) != 0) {
        err = "base64 decode failed";
        goto cleanup;
    }

    int rc = litls_openssh_parse_ed25519(raw, raw_len, seed, pk);
    if (rc == -2) {
        err = "encrypted SSH keys not supported";
        goto cleanup;
    } else if (rc == -4) {
        err = "not an ed25519 key";
        goto cleanup;
    } else if (rc != 0) {
        err = "failed to parse OpenSSH key";
        goto cleanup;
    }
    success = 1;

cleanup:
    /* Wipe both intermediate buffers: b64_clean holds the base64 of the
     * private-key blob, raw holds the decoded blob with the seed embedded. */
    litls_secure_zero(b64_clean, b64_len);
    litls_secure_zero(raw, raw_cap);

    if (!success) {
        litls_secure_zero(seed, 32);
        lua_pushnil(L);
        lua_pushstring(L, err);
        return 2;
    }

    lua_pushlstring(L, (const char *)seed, 32);
    lua_pushlstring(L, (const char *)pk, 32);
    litls_secure_zero(seed, 32);
    return 2;
}

static int lua_litls_ed25519_to_x25519_pk(lua_State *L) {
    size_t pk_len;
    const char *pk = luaL_checklstring(L, 1, &pk_len);
    if (pk_len != 32) {
        lua_pushnil(L);
        lua_pushstring(L, "ed25519 public key must be 32 bytes");
        return 2;
    }
    uint8_t x_pk[32];
    litls_ed25519_to_x25519_pk(x_pk, (const uint8_t *)pk);
    lua_pushlstring(L, (const char *)x_pk, 32);
    return 1;
}

static int lua_litls_ed25519_to_x25519_sk(lua_State *L) {
    size_t seed_len;
    const char *seed = luaL_checklstring(L, 1, &seed_len);
    if (seed_len != 32) {
        lua_pushnil(L);
        lua_pushstring(L, "ed25519 seed must be 32 bytes");
        return 2;
    }
    uint8_t x_sk[32];
    litls_ed25519_to_x25519_sk(x_sk, (const uint8_t *)seed);
    lua_pushlstring(L, (const char *)x_sk, 32);
    return 1;
}

/* ── ECDSA P-256 ───────────────────────────────────────── */

static int lua_litls_p256_keygen(lua_State *L) {
    uint8_t priv[32], pub[65];
    if (litls_p256_keygen(priv, pub) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "P-256 keygen failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)priv, 32);
    lua_pushlstring(L, (const char *)pub, 65);
    return 2;
}

static int lua_litls_p256_ecdsa_sign(lua_State *L) {
    size_t hash_len, priv_len;
    const char *hash = luaL_checklstring(L, 1, &hash_len);
    const char *priv = luaL_checklstring(L, 2, &priv_len);
    if (priv_len != 32) {
        lua_pushnil(L);
        lua_pushstring(L, "private key must be 32 bytes");
        return 2;
    }
    if (hash_len != 32) {
        lua_pushnil(L);
        lua_pushstring(L, "hash must be 32 bytes");
        return 2;
    }
    uint8_t sig[72];
    size_t sig_len = sizeof(sig);
    if (litls_p256_ecdsa_sign((const uint8_t *)hash, hash_len, (const uint8_t *)priv, sig, &sig_len) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "ECDSA sign failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)sig, sig_len);
    return 1;
}

static int lua_litls_p256_ecdsa_verify(lua_State *L) {
    size_t hash_len, pub_len, sig_len;
    const char *hash = luaL_checklstring(L, 1, &hash_len);
    const char *pub  = luaL_checklstring(L, 2, &pub_len);
    const char *sig  = luaL_checklstring(L, 3, &sig_len);
    if (pub_len != 65) {
        lua_pushboolean(L, 0);
        return 1;
    }
    int ok =
        litls_p256_ecdsa_verify((const uint8_t *)hash, hash_len, (const uint8_t *)pub, (const uint8_t *)sig, sig_len);
    lua_pushboolean(L, ok == 0);
    return 1;
}

static int lua_litls_p384_ecdsa_verify(lua_State *L) {
    size_t hash_len, pub_len, sig_len;
    const char *hash = luaL_checklstring(L, 1, &hash_len);
    const char *pub  = luaL_checklstring(L, 2, &pub_len);
    const char *sig  = luaL_checklstring(L, 3, &sig_len);
    if (pub_len != 97) {
        lua_pushboolean(L, 0);
        return 1;
    }
    int ok =
        litls_p384_ecdsa_verify((const uint8_t *)hash, hash_len, (const uint8_t *)pub, (const uint8_t *)sig, sig_len);
    lua_pushboolean(L, ok == 0);
    return 1;
}

/* ── RSA verify ────────────────────────────────────────── */

static int lua_litls_rsa_pkcs1_verify(lua_State *L) {
    size_t hash_len, sig_len, n_len, e_len;
    const char *hash = luaL_checklstring(L, 1, &hash_len);
    int hash_id      = (int)luaL_checkinteger(L, 2);
    const char *sig  = luaL_checklstring(L, 3, &sig_len);
    const char *n    = luaL_checklstring(L, 4, &n_len);
    const char *e    = luaL_checklstring(L, 5, &e_len);
    int ok           = litls_rsa_pkcs1_verify((const uint8_t *)hash, hash_len, hash_id, (const uint8_t *)sig, sig_len,
                                              (const uint8_t *)n, n_len, (const uint8_t *)e, e_len);
    lua_pushboolean(L, ok == 0);
    return 1;
}

/* ── Hex encode (test helper) ──────────────────────────── */

static int lua_litls_hex_encode(lua_State *L) {
    size_t len;
    const char *in = luaL_checklstring(L, 1, &len);
    size_t out_len = len * 2 + 1;
    int heap       = out_len > LITLS_LUA_STACK_MAX;
    char *out      = heap ? (char *)malloc(out_len) : (char *)alloca(out_len);
    if (!out)
        return luaL_error(L, "out of memory");
    for (size_t i = 0; i < len; i++) {
        static const char hex[] = "0123456789abcdef";
        out[i * 2]              = hex[((uint8_t)in[i]) >> 4];
        out[i * 2 + 1]          = hex[((uint8_t)in[i]) & 0x0f];
    }
    lua_pushlstring(L, out, len * 2);
    if (heap)
        free(out);
    return 1;
}

static int lua_litls_hex_decode(lua_State *L) {
    size_t len;
    const char *in = luaL_checklstring(L, 1, &len);
    if (len % 2 != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "hex string must have even length");
        return 2;
    }
    size_t out_len = len / 2;
    int heap       = out_len > LITLS_LUA_STACK_MAX;
    uint8_t *out   = heap ? (uint8_t *)malloc(out_len > 0 ? out_len : 1) : (uint8_t *)alloca(out_len > 0 ? out_len : 1);
    if (!out)
        return luaL_error(L, "out of memory");
    for (size_t i = 0; i < out_len; i++) {
        int hi, lo;
        char c = in[i * 2];
        if (c >= '0' && c <= '9')
            hi = c - '0';
        else if (c >= 'a' && c <= 'f')
            hi = c - 'a' + 10;
        else if (c >= 'A' && c <= 'F')
            hi = c - 'A' + 10;
        else {
            if (heap)
                free(out);
            lua_pushnil(L);
            lua_pushstring(L, "invalid hex");
            return 2;
        }
        c = in[i * 2 + 1];
        if (c >= '0' && c <= '9')
            lo = c - '0';
        else if (c >= 'a' && c <= 'f')
            lo = c - 'a' + 10;
        else if (c >= 'A' && c <= 'F')
            lo = c - 'A' + 10;
        else {
            if (heap)
                free(out);
            lua_pushnil(L);
            lua_pushstring(L, "invalid hex");
            return 2;
        }
        out[i] = (uint8_t)((hi << 4) | lo);
    }
    lua_pushlstring(L, (const char *)out, out_len);
    if (heap)
        free(out);
    return 1;
}

/* ── X.509 / PEM / CSR ─────────────────────────────────── */

static int lua_litls_pem_decode(lua_State *L) {
    size_t pem_len;
    const char *pem = luaL_checklstring(L, 1, &pem_len);
    uint8_t der[8192];
    size_t der_len = sizeof(der);
    if (litls_pem_decode((const uint8_t *)pem, pem_len, der, &der_len) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "PEM decode failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)der, der_len);
    return 1;
}

static int lua_litls_x509_parse(lua_State *L) {
    size_t der_len;
    const char *der = luaL_checklstring(L, 1, &der_len);
    litls_x509_cert_info info;
    if (litls_x509_parse((const uint8_t *)der, der_len, &info) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "failed to parse certificate");
        return 2;
    }
    lua_createtable(L, 0, 5);

    if (info.common_name_len > 0) {
        lua_pushlstring(L, info.common_name, info.common_name_len);
        lua_setfield(L, -2, "common_name");
    }

    /* Return timestamps as numeric strings for date_to_ts() compatibility */
    char ts_buf[24];
    snprintf(ts_buf, sizeof(ts_buf), "%lld", (long long)info.not_before);
    lua_pushstring(L, ts_buf);
    lua_setfield(L, -2, "not_before");
    snprintf(ts_buf, sizeof(ts_buf), "%lld", (long long)info.not_after);
    lua_pushstring(L, ts_buf);
    lua_setfield(L, -2, "not_after");

    if (info.san_count > 0) {
        lua_createtable(L, info.san_count, 0);
        for (int i = 0; i < info.san_count; i++) {
            lua_pushlstring(L, info.sans[i].name, info.sans[i].len);
            lua_rawseti(L, -2, i + 1);
        }
        lua_setfield(L, -2, "sans");
    }

    /* X.509 v3 extension summary (present only if the extension was seen in
     * the certificate; otherwise the field is absent from the returned table). */
    if (info.basic_constraints_present) {
        lua_pushboolean(L, info.is_ca);
        lua_setfield(L, -2, "is_ca");
        lua_pushinteger(L, info.path_len_constraint);
        lua_setfield(L, -2, "path_len_constraint");
    }
    if (info.key_usage_present) {
        lua_pushinteger(L, info.key_usage);
        lua_setfield(L, -2, "key_usage");
    }
    if (info.eku_present) {
        lua_pushboolean(L, info.eku_server_auth);
        lua_setfield(L, -2, "eku_server_auth");
        lua_pushboolean(L, info.eku_any);
        lua_setfield(L, -2, "eku_any");
    }
    return 1;
}

/* x509_verify_hostname(der, hostname) -> true on match, (nil, err) otherwise. */
static int lua_litls_x509_verify_hostname(lua_State *L) {
    size_t der_len;
    const char *der      = luaL_checklstring(L, 1, &der_len);
    const char *hostname = luaL_checkstring(L, 2);
    litls_x509_cert_info info;
    if (litls_x509_parse((const uint8_t *)der, der_len, &info) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "failed to parse certificate");
        return 2;
    }
    if (litls_x509_verify_hostname(&info, hostname) == LITLS_X509_OK) {
        lua_pushboolean(L, 1);
        return 1;
    }
    lua_pushnil(L);
    lua_pushstring(L, "hostname mismatch");
    return 2;
}

static int lua_litls_csr_generate(lua_State *L) {
    size_t key_len, pub_len, domain_len;
    const char *key    = luaL_checklstring(L, 1, &key_len);
    const char *pub    = luaL_checklstring(L, 2, &pub_len);
    const char *domain = luaL_checklstring(L, 3, &domain_len);
    if (key_len != 32 || pub_len != 65) {
        lua_pushnil(L);
        lua_pushstring(L, "private key must be 32B, public key 65B");
        return 2;
    }

    /* Collect SANs from optional table argument */
    const char *san_ptrs[LITLS_X509_MAX_SANS];
    int san_count = 0;
    if (lua_gettop(L) >= 4 && lua_istable(L, 4)) {
        lua_pushnil(L);
        while (lua_next(L, 4) != 0 && san_count < LITLS_X509_MAX_SANS) {
            if (lua_type(L, -1) != LUA_TSTRING) {
                lua_pushnil(L);
                lua_pushstring(L, "SAN entries must be strings");
                return 2;
            }
            /* Value stays referenced in the table (arg 4) on the stack, so the
             * string pointer remains valid through litls_csr_generate(). */
            san_ptrs[san_count++] = lua_tostring(L, -1);
            lua_pop(L, 1);
        }
    }

    litls_csr_params params;
    params.domain      = domain;
    params.sans        = san_count > 0 ? san_ptrs : NULL;
    params.san_count   = san_count;
    params.private_key = (const uint8_t *)key;
    params.public_key  = (const uint8_t *)pub;

    uint8_t der[LITLS_CSR_MAX_SIZE];
    int ret = litls_csr_generate(&params, der, sizeof(der));
    if (ret < 0) {
        lua_pushnil(L);
        lua_pushstring(L, "CSR generation failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)der, (size_t)ret);
    return 1;
}

static int lua_litls_x509_verify_chain(lua_State *L) {
    /* Arg 1: table of DER cert strings (leaf first) */
    luaL_checktype(L, 1, LUA_TTABLE);
    /* Arg 2: CA PEM bundle string */
    size_t ca_pem_len;
    const char *ca_pem = luaL_checklstring(L, 2, &ca_pem_len);

    /* Collect chain DERs */
    int chain_count = (int)lua_objlen(L, 1);
    if (chain_count < 1 || chain_count > LITLS_X509_MAX_CHAIN) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid chain length");
        return 2;
    }
    const uint8_t *chain_ders[LITLS_X509_MAX_CHAIN];
    size_t chain_lens[LITLS_X509_MAX_CHAIN];
    for (int i = 0; i < chain_count; i++) {
        lua_rawgeti(L, 1, i + 1);
        if (lua_type(L, -1) != LUA_TSTRING) {
            lua_pop(L, 1);
            lua_pushnil(L);
            lua_pushstring(L, "chain entries must be DER strings");
            return 2;
        }
        chain_ders[i] = (const uint8_t *)lua_tolstring(L, -1, &chain_lens[i]);
        lua_pop(L, 1);
    }

    /* Load trust anchors */
    litls_x509_trust_anchor anchors[LITLS_X509_MAX_ANCHORS];
    int anchor_count = litls_x509_load_anchors((const uint8_t *)ca_pem, ca_pem_len, anchors, LITLS_X509_MAX_ANCHORS);
    if (anchor_count <= 0) {
        lua_pushnil(L);
        lua_pushstring(L, "no trust anchors loaded");
        return 2;
    }

    int64_t now = (int64_t)time(NULL);
    int ret     = litls_x509_verify_chain(chain_ders, chain_lens, chain_count, anchors, anchor_count, now);
    if (ret == LITLS_X509_OK) {
        lua_pushboolean(L, 1);
        return 1;
    }

    lua_pushnil(L);
    switch (ret) {
    case LITLS_X509_ERR_PARSE:
        lua_pushstring(L, "certificate parse error");
        break;
    case LITLS_X509_ERR_SIGNATURE:
        lua_pushstring(L, "signature verification failed");
        break;
    case LITLS_X509_ERR_EXPIRED:
        lua_pushstring(L, "certificate expired");
        break;
    case LITLS_X509_ERR_NO_ANCHOR:
        lua_pushstring(L, "no matching trust anchor");
        break;
    case LITLS_X509_ERR_CHAIN:
        lua_pushstring(L, "invalid chain");
        break;
    case LITLS_X509_ERR_CONSTRAINT:
        lua_pushstring(L, "certificate constraint violation");
        break;
    default:
        lua_pushstring(L, "unknown error");
        break;
    }
    return 2;
}

/* ── RSA-PSS verify ────────────────────────────────────── */

static int lua_litls_rsa_pss_sha256_verify(lua_State *L) {
    size_t hash_len, sig_len, n_len, e_len;
    const char *hash = luaL_checklstring(L, 1, &hash_len);
    const char *sig  = luaL_checklstring(L, 2, &sig_len);
    const char *n    = luaL_checklstring(L, 3, &n_len);
    const char *e    = luaL_checklstring(L, 4, &e_len);
    if (hash_len != 32) {
        lua_pushboolean(L, 0);
        return 1;
    }
    int ok = litls_rsa_pss_sha256_verify((const uint8_t *)hash, (const uint8_t *)sig, sig_len, (const uint8_t *)n,
                                         n_len, (const uint8_t *)e, e_len);
    lua_pushboolean(L, ok == 0);
    return 1;
}

/* ── TLS 1.3 key schedule (test-facing) ────────────────── */

static int lua_litls_tls13_expand_label_256(lua_State *L) {
    size_t secret_len, ctx_len;
    const char *secret  = luaL_checklstring(L, 1, &secret_len);
    const char *label   = luaL_checkstring(L, 2);
    const char *context = luaL_checklstring(L, 3, &ctx_len);
    int out_len         = (int)luaL_checkinteger(L, 4);
    if (secret_len < 32 || out_len < 0 || out_len > 255) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid parameters");
        return 2;
    }
    uint8_t out[255];
    if (litls_tls13_expand_label_256((const uint8_t *)secret, label, (const uint8_t *)context, ctx_len, out,
                                     (size_t)out_len) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "expand label failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)out, (size_t)out_len);
    return 1;
}

static int lua_litls_tls13_derive_secret_256(lua_State *L) {
    size_t secret_len, hash_len;
    const char *secret = luaL_checklstring(L, 1, &secret_len);
    const char *label  = luaL_checkstring(L, 2);
    const char *hash   = luaL_checklstring(L, 3, &hash_len);
    if (secret_len < 32 || hash_len != 32) {
        lua_pushnil(L);
        lua_pushstring(L, "invalid parameters");
        return 2;
    }
    uint8_t out[32];
    if (litls_tls13_derive_secret_256((const uint8_t *)secret, label, (const uint8_t *)hash, out) != 0) {
        lua_pushnil(L);
        lua_pushstring(L, "derive secret failed");
        return 2;
    }
    lua_pushlstring(L, (const char *)out, 32);
    return 1;
}

/* ── Module registration ───────────────────────────────── */

static luaL_Reg litls_funcs[] = {
    /* base64 */
    {"base64_encode",             lua_litls_base64_encode            },
    {"base64_decode",             lua_litls_base64_decode            },
    {"base64url_encode",          lua_litls_base64url_encode         },
    {"base64url_decode",          lua_litls_base64url_decode         },
    /* rng */
    {"random_bytes",              lua_litls_random_bytes             },
    /* constant-time compare */
    {"secure_memcmp",             lua_litls_secure_memcmp            },
    /* sha */
    {"sha1",                      lua_litls_sha1                     },
    {"sha256",                    lua_litls_sha256                   },
    {"sha384",                    lua_litls_sha384                   },
    {"sha512",                    lua_litls_sha512                   },
    /* hmac */
    {"hmac_sha256",               lua_litls_hmac_sha256              },
    {"hmac_sha384",               lua_litls_hmac_sha384              },
    /* hkdf */
    {"hkdf_extract_sha256",       lua_litls_hkdf_extract_sha256      },
    {"hkdf_expand_sha256",        lua_litls_hkdf_expand_sha256       },
    {"hkdf_extract_sha384",       lua_litls_hkdf_extract_sha384      },
    {"hkdf_expand_sha384",        lua_litls_hkdf_expand_sha384       },
    /* chacha20 */
    {"chacha20_block",            lua_litls_chacha20_block           },
    {"chacha20_encrypt",          lua_litls_chacha20_encrypt         },
    /* poly1305 */
    {"poly1305_auth",             lua_litls_poly1305_auth            },
    /* chacha20-poly1305 */
    {"chacha20_poly1305_encrypt", lua_litls_chacha20_poly1305_encrypt},
    {"chacha20_poly1305_decrypt", lua_litls_chacha20_poly1305_decrypt},
    /* aes-gcm */
    {"aes128_gcm_encrypt",        lua_litls_aes128_gcm_encrypt       },
    {"aes128_gcm_decrypt",        lua_litls_aes128_gcm_decrypt       },
    {"aes256_gcm_encrypt",        lua_litls_aes256_gcm_encrypt       },
    {"aes256_gcm_decrypt",        lua_litls_aes256_gcm_decrypt       },
    /* x25519 */
    {"x25519",                    lua_litls_x25519                   },
    {"x25519_base",               lua_litls_x25519_base              },
    /* ed25519 */
    {"ed25519_keypair",           lua_litls_ed25519_keypair          },
    {"ed25519_sign",              lua_litls_ed25519_sign             },
    {"ed25519_verify",            lua_litls_ed25519_verify           },
    {"ed25519_to_x25519_pk",      lua_litls_ed25519_to_x25519_pk     },
    {"ed25519_to_x25519_sk",      lua_litls_ed25519_to_x25519_sk     },
    {"openssh_parse_ed25519",     lua_litls_openssh_parse_ed25519    },
    /* ecdsa p-256 */
    {"p256_keygen",               lua_litls_p256_keygen              },
    {"p256_ecdsa_sign",           lua_litls_p256_ecdsa_sign          },
    {"p256_ecdsa_verify",         lua_litls_p256_ecdsa_verify        },
    /* ecdsa p-384 */
    {"p384_ecdsa_verify",         lua_litls_p384_ecdsa_verify        },
    /* rsa */
    {"rsa_pkcs1_verify",          lua_litls_rsa_pkcs1_verify         },
    {"rsa_pss_sha256_verify",     lua_litls_rsa_pss_sha256_verify    },
    /* hex helpers */
    {"hex_encode",                lua_litls_hex_encode               },
    {"hex_decode",                lua_litls_hex_decode               },
    /* x509 / pem / csr */
    {"pem_decode",                lua_litls_pem_decode               },
    {"x509_parse",                lua_litls_x509_parse               },
    {"csr_generate",              lua_litls_csr_generate             },
    {"x509_verify_chain",         lua_litls_x509_verify_chain        },
    {"x509_verify_hostname",      lua_litls_x509_verify_hostname     },
    /* tls 1.3 key schedule (test-facing) */
    {"tls13_expand_label_256",    lua_litls_tls13_expand_label_256   },
    {"tls13_derive_secret_256",   lua_litls_tls13_derive_secret_256  },
    /* sentinel */
    {NULL,                        NULL                               }
};

int luaopen_litls_core(lua_State *L) {
    luaL_newlib(L, litls_funcs);
    lua_pushinteger(L, LITLS_HASH_SHA256);
    lua_setfield(L, -2, "HASH_SHA256");
    lua_pushinteger(L, LITLS_HASH_SHA384);
    lua_setfield(L, -2, "HASH_SHA384");
    lua_pushinteger(L, LITLS_HASH_SHA512);
    lua_setfield(L, -2, "HASH_SHA512");
    return 1;
}