// SPDX-FileCopyrightText: © 2022—2026 Vladimir Zorin <vladimir@deviant.guru>
// SPDX-License-Identifier: LicenseRef-OWL-1.0-or-later

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "litls.h"

#include <lauxlib.h>
#include <lua.h>

#define RETURN_CUSTOM_ERR(L, msg) \
    do {                          \
        lua_pushnil(L);           \
        lua_pushstring(L, msg);   \
        return 2;                 \
    } while (0)

int lua_base64_decode(lua_State *L) {
    size_t inLen;
    const char *in = lua_tolstring(L, 1, &inLen);
    size_t outLen  = (inLen * 3) / 4 + 4;
    uint8_t *out   = (uint8_t *)malloc(outLen);
    if (out == NULL) {
        RETURN_CUSTOM_ERR(L, "memory allocation failed");
    }
    if (litls_base64_decode(in, inLen, out, &outLen) != 0) {
        free(out);
        RETURN_CUSTOM_ERR(L, "failed to decode base64");
    }
    lua_pushlstring(L, (char *)out, outLen);
    free(out);
    return 1;
}

int lua_base64_encode(lua_State *L) {
    size_t inLen;
    const char *in = lua_tolstring(L, 1, &inLen);
    size_t outCap  = ((inLen + 2) / 3) * 4 + 1;
    char *out      = (char *)malloc(outCap);
    if (out == NULL) {
        RETURN_CUSTOM_ERR(L, "memory allocation failed");
    }
    size_t n = litls_base64_encode((const uint8_t *)in, inLen, out, outCap);
    if (n == 0 && inLen > 0) {
        free(out);
        RETURN_CUSTOM_ERR(L, "failed to encode base64");
    }
    lua_pushlstring(L, out, n);
    free(out);
    return 1;
}

int lua_sha1(lua_State *L) {
    size_t data_size;
    const char *data = lua_tolstring(L, 1, &data_size);
    uint8_t hash[LITLS_SHA1_DIGEST_SIZE];
    litls_sha1((const uint8_t *)data, data_size, hash);
    lua_pushlstring(L, (char *)hash, LITLS_SHA1_DIGEST_SIZE);
    return 1;
}

int lua_sha256(lua_State *L) {
    size_t data_size;
    const char *data = lua_tolstring(L, 1, &data_size);
    uint8_t hash[LITLS_SHA256_DIGEST_SIZE];
    litls_sha256((const uint8_t *)data, data_size, hash);
    lua_pushlstring(L, (char *)hash, LITLS_SHA256_DIGEST_SIZE);
    return 1;
}

int lua_hmac(lua_State *L) {
    size_t secret_size, msg_size;
    const char *secret = lua_tolstring(L, 1, &secret_size);
    const char *msg    = lua_tolstring(L, 2, &msg_size);
    uint8_t hash[LITLS_SHA256_DIGEST_SIZE];
    litls_hmac_sha256((const uint8_t *)secret, secret_size, (const uint8_t *)msg, msg_size, hash);
    lua_pushlstring(L, (char *)hash, LITLS_SHA256_DIGEST_SIZE);
    return 1;
}

int lua_ecc_generate_key(lua_State *L) {
    uint8_t private_key[32];
    uint8_t public_key[65];

    if (litls_p256_keygen(private_key, public_key) != 0) {
        RETURN_CUSTOM_ERR(L, "failed to generate ECC key");
    }

    /* Return: private (32B), public (65B), x (32B), y (32B) */
    lua_pushlstring(L, (char *)private_key, 32);
    lua_pushlstring(L, (char *)public_key, 65);
    lua_pushlstring(L, (char *)(public_key + 1), 32);  /* x = pub[1..32] */
    lua_pushlstring(L, (char *)(public_key + 33), 32); /* y = pub[33..64] */
    return 4;
}

int lua_ecc_sign(lua_State *L) {
    size_t key_size, pub_key_size, msg_size;
    const char *private_key = lua_tolstring(L, 1, &key_size);
    const char *public_key  = lua_tolstring(L, 2, &pub_key_size);
    const char *msg         = lua_tolstring(L, 3, &msg_size);

    /* Hash the message first */
    uint8_t hash[LITLS_SHA256_DIGEST_SIZE];
    litls_sha256((const uint8_t *)msg, msg_size, hash);

    /* Sign the hash */
    uint8_t sig[64];
    size_t sig_len = 0;
    if (litls_p256_ecdsa_sign(hash, sizeof(hash), (const uint8_t *)private_key, sig, &sig_len) != 0) {
        RETURN_CUSTOM_ERR(L, "failed to sign hash");
    }

    lua_pushlstring(L, (char *)sig, 64);
    return 1;
}

int lua_ecc_verify(lua_State *L) {
    size_t public_size, msg_size, sig_size;
    const char *pub_key = lua_tolstring(L, 1, &public_size);
    const char *msg     = lua_tolstring(L, 2, &msg_size);
    const char *sig     = lua_tolstring(L, 3, &sig_size);

    /* Hash the message first */
    uint8_t hash[LITLS_SHA256_DIGEST_SIZE];
    litls_sha256((const uint8_t *)msg, msg_size, hash);

    int ret = litls_p256_ecdsa_verify(hash, sizeof(hash), (const uint8_t *)pub_key, (const uint8_t *)sig, sig_size);
    lua_pushboolean(L, ret == 0);
    return 1;
}

int lua_ed25519_generate_key(lua_State *L) {
    uint8_t seed[32];
    uint8_t public_key[32];
    uint8_t secret_key[64];

    if (litls_random_bytes(seed, 32) != 0) {
        RETURN_CUSTOM_ERR(L, "failed to generate random bytes");
    }

    litls_ed25519_keypair(public_key, secret_key, seed);

    /* Return (seed_32B, pub_32B) */
    lua_pushlstring(L, (char *)seed, 32);
    lua_pushlstring(L, (char *)public_key, 32);
    return 2;
}

int lua_ed25519_sign(lua_State *L) {
    size_t key_size, msg_size;
    const char *seed = lua_tolstring(L, 1, &key_size);
    const char *msg  = lua_tolstring(L, 2, &msg_size);

    /* Expand seed to full secret key + public key */
    uint8_t public_key[32];
    uint8_t secret_key[64];
    litls_ed25519_keypair(public_key, secret_key, (const uint8_t *)seed);

    uint8_t sig[64];
    if (litls_ed25519_sign(sig, (const uint8_t *)msg, msg_size, public_key, secret_key) != 0) {
        RETURN_CUSTOM_ERR(L, "failed to sign message");
    }

    lua_pushlstring(L, (char *)sig, 64);
    return 1;
}

int lua_ed25519_verify(lua_State *L) {
    size_t pub_key_size, msg_size, sig_size;
    const char *public_key = lua_tolstring(L, 1, &pub_key_size);
    const char *msg        = lua_tolstring(L, 2, &msg_size);
    const char *sig        = lua_tolstring(L, 3, &sig_size);

    int ret = litls_ed25519_verify((const uint8_t *)sig, (const uint8_t *)msg, msg_size, (const uint8_t *)public_key);
    lua_pushboolean(L, ret == 0);
    return 1;
}

int lua_generate_csr(lua_State *L) {
    size_t key_size, pub_key_size, domain_size;
    const char *private_key = lua_tolstring(L, 1, &key_size);
    const char *public_key  = lua_tolstring(L, 2, &pub_key_size);
    const char *domain      = lua_tolstring(L, 3, &domain_size);

    litls_csr_params params;
    params.domain      = domain;
    params.private_key = (const uint8_t *)private_key;
    params.public_key  = (const uint8_t *)public_key;
    params.sans        = NULL;
    params.san_count   = 0;

    /* Collect SANs from optional table argument */
    const char *san_strs[64];
    if (lua_gettop(L) >= 4 && !lua_isnil(L, 4)) {
        if (!lua_istable(L, 4)) {
            RETURN_CUSTOM_ERR(L, "alternative names must be a table");
        }
        int n = (int)lua_objlen(L, 4);
        if (n > 64)
            n = 64;
        for (int i = 0; i < n; i++) {
            lua_rawgeti(L, 4, i + 1);
            san_strs[i] = lua_tostring(L, -1);
            lua_pop(L, 1);
        }
        params.sans      = san_strs;
        params.san_count = n;
    }

    uint8_t der[LITLS_CSR_MAX_SIZE];
    int ret = litls_csr_generate(&params, der, sizeof(der));
    if (ret < 0) {
        RETURN_CUSTOM_ERR(L, "failed to generate certificate request");
    }

    lua_pushlstring(L, (char *)der, (size_t)ret);
    return 1;
}

int lua_der_to_pem_ecc_key(lua_State *L) {
    size_t key_size, pub_key_size;
    const char *private_key = lua_tolstring(L, 1, &key_size);
    const char *public_key  = lua_tolstring(L, 2, &pub_key_size);

    char pem[4096];
    size_t pem_len = 0;
    if (litls_ec_privkey_to_pem((const uint8_t *)private_key, (const uint8_t *)public_key, pem, sizeof(pem),
                                &pem_len) != 0) {
        RETURN_CUSTOM_ERR(L, "failed to convert key to PEM");
    }

    lua_pushlstring(L, pem, pem_len);
    return 1;
}

int lua_parse_x509_cert(lua_State *L) {
    size_t cert_der_size;
    const char *cert_der = lua_tolstring(L, 1, &cert_der_size);

    if (!cert_der) {
        RETURN_CUSTOM_ERR(L, "no certificate provided");
    }

    litls_x509_cert_info info;
    memset(&info, 0, sizeof(info));
    int ret = litls_x509_parse((const uint8_t *)cert_der, cert_der_size, &info);
    if (ret != 0) {
        RETURN_CUSTOM_ERR(L, "failed to parse certificate");
    }

    /* Create result table */
    lua_createtable(L, 0, 5);

    /* Add common name */
    if (info.common_name_len > 0) {
        lua_pushstring(L, "common_name");
        lua_pushlstring(L, info.common_name, info.common_name_len);
        lua_settable(L, -3);
    }

    /* Add validity dates as Unix timestamp strings */
    char ts_buf[32];

    lua_pushstring(L, "not_before");
    snprintf(ts_buf, sizeof(ts_buf), "%lld", (long long)info.not_before);
    lua_pushstring(L, ts_buf);
    lua_settable(L, -3);

    lua_pushstring(L, "not_after");
    snprintf(ts_buf, sizeof(ts_buf), "%lld", (long long)info.not_after);
    lua_pushstring(L, ts_buf);
    lua_settable(L, -3);

    /* Add SANs as array */
    if (info.san_count > 0) {
        lua_pushstring(L, "sans");
        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_settable(L, -3);
    }

    return 1;
}

static luaL_Reg funcs[] = {
    {"sha1",                 lua_sha1                },
    {"sha256",               lua_sha256              },
    {"hmac",                 lua_hmac                },
    {"base64_decode",        lua_base64_decode       },
    {"base64_encode",        lua_base64_encode       },
    {"ecc_generate_key",     lua_ecc_generate_key    },
    {"ecc_sign",             lua_ecc_sign            },
    {"ecc_verify",           lua_ecc_verify          },
    {"ed25519_generate_key", lua_ed25519_generate_key},
    {"ed25519_sign",         lua_ed25519_sign        },
    {"ed25519_verify",       lua_ed25519_verify      },
    {"generate_csr",         lua_generate_csr        },
    {"der_to_pem_ecc_key",   lua_der_to_pem_ecc_key  },
    {"parse_x509_cert",      lua_parse_x509_cert     },
    {NULL,                   NULL                    }
};

int luaopen_crypto_core(lua_State *L) {
    /* Return the module */
    luaL_newlib(L, funcs);
    return 1;
}