#include <Library/UefiLib.h>
#include <Library/BaseLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/BaseMemoryLib.h>
#include <Uefi.h>

#include "api.h"
#include "params.h"
#include "wots.h"
#include "fors.h"
#include "hash.h"
#include "thash.h"
#include "address.h"
#include "utils.h"
#include "merkle.h"

/*
 * Returns the length of a secret key, in bytes
 */
UINT64 crypto_sign_secretkeybytes(VOID)
{
    return CRYPTO_SECRETKEYBYTES;
}

/*
 * Returns the length of a public key, in bytes
 */
UINT64 crypto_sign_publickeybytes(VOID)
{
    return CRYPTO_PUBLICKEYBYTES;
}

/*
 * Returns the length of a signature, in bytes
 */
UINT64 crypto_sign_bytes(VOID)
{
    return CRYPTO_BYTES;
}

/*
 * Returns the length of the seed required to generate a key pair, in bytes
 */
UINT64 crypto_sign_seedbytes(VOID)
{
    return CRYPTO_SEEDBYTES;
}

/*
 * Generates an SPX key pair given a seed of length
 * Format sk: [SK_SEED || SK_PRF || PUB_SEED || root]
 * Format pk: [PUB_SEED || root]
 */
INTN crypto_sign_seed_keypair(UINT8 *pk, UINT8 *sk,
                             const UINT8 *seed)
{
    spx_ctx ctx;

    /* Initialize SK_SEED, SK_PRF and PUB_SEED from seed. */
    CopyMem(sk, seed, CRYPTO_SEEDBYTES);

    CopyMem(pk, sk + 2*SPX_N, SPX_N);

    CopyMem(ctx.pub_seed, pk, SPX_N);
    CopyMem(ctx.sk_seed, sk, SPX_N);

    /* This hook allows the hash function instantiation to do whatever
       preparation or computation it needs, based on the public seed. */
    initialize_hash_function(&ctx);

    /* Compute root node of the top-most subtree. */
    merkle_gen_root(sk + 3*SPX_N, &ctx);

    CopyMem(pk + SPX_N, sk + 3*SPX_N, SPX_N);

    return 0;
}

/*
 * Generates an SPX key pair.
 * Format sk: [SK_SEED || SK_PRF || PUB_SEED || root]
 * Format pk: [PUB_SEED || root]
 */
#ifdef BP_CRYPTO_DEFINE_RANDOM_FLAG
INTN crypto_sign_keypair(UINT8 *pk, UINT8 *sk)
{
  UINT8 seed[CRYPTO_SEEDBYTES];
  randombytes(seed, CRYPTO_SEEDBYTES);
  crypto_sign_seed_keypair(pk, sk, seed);

  return 0;
}

/**
 * Returns an array containing a detached signature.
 */
INTN crypto_sign_signature(UINT8 *sig, UINTN *siglen,
                          const UINT8 *m, UINTN mlen, const UINT8 *sk)
{
    spx_ctx ctx;

    const UINT8 *sk_prf = sk + SPX_N;
    const UINT8 *pk = sk + 2*SPX_N;

    UINT8 optrand[SPX_N];
    UINT8 mhash[SPX_FORS_MSG_BYTES];
    UINT8 root[SPX_N];
    UINT32 i;
    UINT64 tree;
    UINT32 idx_leaf;
    UINT32 wots_addr[8] = {0};
    UINT32 tree_addr[8] = {0};

    CopyMem(ctx.sk_seed, sk, SPX_N);
    CopyMem(ctx.pub_seed, pk, SPX_N);

    /* This hook allows the hash function instantiation to do whatever
       preparation or computation it needs, based on the public seed. */
    initialize_hash_function(&ctx);

    set_type(wots_addr, SPX_ADDR_TYPE_WOTS);
    set_type(tree_addr, SPX_ADDR_TYPE_HASHTREE);

    /* Optionally, signing can be made non-deterministic using optrand.
       This can help counter side-channel attacks that would benefit from
       getting a large number of traces when the signer uses the same nodes. */
    randombytes(optrand, SPX_N);
    /* Compute the digest randomization value. */
    gen_message_random(sig, sk_prf, optrand, m, mlen, &ctx);

    /* Derive the message digest and leaf index from R, PK and M. */
    hash_message(mhash, &tree, &idx_leaf, sig, pk, m, mlen, &ctx);
    sig += SPX_N;

    set_tree_addr(wots_addr, tree);
    set_keypair_addr(wots_addr, idx_leaf);

    /* Sign the message hash using FORS. */
    fors_sign(sig, root, mhash, &ctx, wots_addr);
    sig += SPX_FORS_BYTES;

    for (i = 0; i < SPX_D; i++) {
        set_layer_addr(tree_addr, i);
        set_tree_addr(tree_addr, tree);

        copy_subtree_addr(wots_addr, tree_addr);
        set_keypair_addr(wots_addr, idx_leaf);

        merkle_sign(sig, root, &ctx, wots_addr, tree_addr, idx_leaf);
        sig += SPX_WOTS_BYTES + SPX_TREE_HEIGHT * SPX_N;

        /* Update the indices for the next layer. */
        idx_leaf = (tree & ((1 << SPX_TREE_HEIGHT)-1));
        tree = tree >> SPX_TREE_HEIGHT;
    }

    *siglen = SPX_BYTES;

    return 0;
}
#endif
/**
 * Verifies a detached signature and message under a given public key.
 */
INTN crypto_sign_verify(const UINT8 *sig, UINTN siglen,
                       const UINT8 *m, UINTN mlen, const UINT8 *pk)
{
    spx_ctx ctx;
    const UINT8 *pub_root = pk + SPX_N;
    UINT8 mhash[SPX_FORS_MSG_BYTES];
    UINT8 wots_pk[SPX_WOTS_BYTES];
    UINT8 root[SPX_N];
    UINT8 leaf[SPX_N];
    UINT32 i;
    UINT64 tree;
    UINT32 idx_leaf;
    UINT32 wots_addr[8] = {0};
    UINT32 tree_addr[8] = {0};
    UINT32 wots_pk_addr[8] = {0};

    if (siglen != SPX_BYTES) {
        return -1;
    }

    CopyMem(ctx.pub_seed, pk, SPX_N);

    /* This hook allows the hash function instantiation to do whatever
       preparation or computation it needs, based on the public seed. */
    initialize_hash_function(&ctx);

    set_type(wots_addr, SPX_ADDR_TYPE_WOTS);
    set_type(tree_addr, SPX_ADDR_TYPE_HASHTREE);
    set_type(wots_pk_addr, SPX_ADDR_TYPE_WOTSPK);

    /* Derive the message digest and leaf index from R || PK || M. */
    /* The additional SPX_N is a result of the hash domain separator. */
    hash_message(mhash, &tree, &idx_leaf, sig, pk, m, mlen, &ctx);
    sig += SPX_N;

    /* Layer correctly defaults to 0, so no need to set_layer_addr */
    set_tree_addr(wots_addr, tree);
    set_keypair_addr(wots_addr, idx_leaf);

    fors_pk_from_sig(root, sig, mhash, &ctx, wots_addr);
    sig += SPX_FORS_BYTES;

    /* For each subtree.. */
    for (i = 0; i < SPX_D; i++) {
        set_layer_addr(tree_addr, i);
        set_tree_addr(tree_addr, tree);

        copy_subtree_addr(wots_addr, tree_addr);
        set_keypair_addr(wots_addr, idx_leaf);

        copy_keypair_addr(wots_pk_addr, wots_addr);

        /* The WOTS public key is only correct if the signature was correct. */
        /* Initially, root is the FORS pk, but on subsequent iterations it is
           the root of the subtree below the currently processed subtree. */
        wots_pk_from_sig(wots_pk, sig, root, &ctx, wots_addr);
        sig += SPX_WOTS_BYTES;

        /* Compute the leaf node using the WOTS public key. */
        thash(leaf, wots_pk, SPX_WOTS_LEN, &ctx, wots_pk_addr);

        /* Compute the root node of this subtree. */
        compute_root(root, leaf, idx_leaf, 0, sig, SPX_TREE_HEIGHT,
                     &ctx, tree_addr);
        sig += SPX_TREE_HEIGHT * SPX_N;

        /* Update the indices for the next layer. */
        idx_leaf = (tree & ((1 << SPX_TREE_HEIGHT)-1));
        tree = tree >> SPX_TREE_HEIGHT;
    }

    /* Check if the root node equals the root node in the public key. */
    if (CompareMem(root, pub_root, SPX_N)) {
        return -1;
    }

    return 0;
}


/**
 * Returns an array containing the signature followed by the message.
 */
INTN crypto_sign(UINT8 *sm, UINT64 *smlen,
                const UINT8 *m, UINT64 mlen,
                const UINT8 *sk)
{
    UINTN siglen;

    crypto_sign_signature(sm, &siglen, m, (UINTN)mlen, sk);

    CopyMem(sm + SPX_BYTES, m, mlen);
    *smlen = siglen + mlen;

    return 0;
}

/**
 * Verifies a given signature-message pair under a given public key.
 */
INTN crypto_sign_open(UINT8 *m, UINT64 *mlen,
                     const UINT8 *sm, UINT64 smlen,
                     const UINT8 *pk)
{
    /* The API caller does not necessarily know what size a signature should be
       but SPHINCS+ signatures are always exactly SPX_BYTES. */
    if (smlen < SPX_BYTES) {
        SetMem(m, smlen, 0);
        *mlen = 0;
        return -1;
    }

    *mlen = smlen - SPX_BYTES;

    if (crypto_sign_verify(sm, SPX_BYTES, sm + SPX_BYTES, (UINTN)*mlen, pk)) {
        SetMem(m, smlen, 0);
        *mlen = 0;
        return -1;
    }

    /* If verification was successful, move the message to the right place. */
    CopyMem(m, sm + SPX_BYTES, *mlen);

    return 0;
}
