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

#include "fors.h"
#include "utils.h"
#include "utilsx1.h"
#include "hash.h"
#include "thash.h"
#include "address.h"

static void fors_gen_sk(unsigned char *sk, const ascon_sign_ctx *ctx,
                        uint32_t fors_leaf_addr[8])
{
    prf_addr(sk, ctx, fors_leaf_addr);
}

static void fors_sk_to_leaf(unsigned char *leaf, const unsigned char *sk,
                            const ascon_sign_ctx *ctx,
                            uint32_t fors_leaf_addr[8])
{
    thash(leaf, sk, 1, ctx, fors_leaf_addr);
}

struct fors_gen_leaf_info {
    uint32_t leaf_addrx[8];
};

static void fors_gen_leafx1(unsigned char *leaf,
                            const ascon_sign_ctx *ctx,
                            uint32_t addr_idx, void *info)
{
    struct fors_gen_leaf_info *fors_info = info;
    uint32_t *fors_leaf_addr = fors_info->leaf_addrx;

    /* Only set the parts that the caller doesn't set */
    set_tree_index(fors_leaf_addr, addr_idx);
    set_type(fors_leaf_addr, ASGN_ADDR_TYPE_FORSPRF);
    fors_gen_sk(leaf, ctx, fors_leaf_addr);

    set_type(fors_leaf_addr, ASGN_ADDR_TYPE_FORSTREE);
    fors_sk_to_leaf(leaf, leaf,
                    ctx, fors_leaf_addr);
}

/**
 * Interprets m as ASGN_FORS_HEIGHT-bit unsigned integers.
 * Assumes m contains at least ASGN_FORS_HEIGHT * ASGN_FORS_TREES bits.
 * Assumes indices has space for ASGN_FORS_TREES integers.
 */
static void message_to_indices(uint32_t *indices, const unsigned char *m)
{
    unsigned int i, j;
    unsigned int offset = 0;

    for (i = 0; i < ASGN_FORS_TREES; i++) {
        indices[i] = 0;
        for (j = 0; j < ASGN_FORS_HEIGHT; j++) {
            indices[i] ^= ((m[offset >> 3] >> (offset & 0x7)) & 1u) << j;
            offset++;
        }
    }
}

/**
 * Signs a message m, deriving the secret key from sk_seed and the FTS address.
 * Assumes m contains at least ASGN_FORS_HEIGHT * ASGN_FORS_TREES bits.
 */
void fors_sign(unsigned char *sig, unsigned char *pk,
               const unsigned char *m,
               const ascon_sign_ctx *ctx,
               const uint32_t fors_addr[8])
{
    uint32_t indices[ASGN_FORS_TREES];
    unsigned char roots[ASGN_FORS_TREES * ASGN_N];
    uint32_t fors_tree_addr[8] = {0};
    struct fors_gen_leaf_info fors_info = {0};
    uint32_t *fors_leaf_addr = fors_info.leaf_addrx;
    uint32_t fors_pk_addr[8] = {0};
    uint32_t idx_offset;
    unsigned int i;

    copy_keypair_addr(fors_tree_addr, fors_addr);
    copy_keypair_addr(fors_leaf_addr, fors_addr);

    copy_keypair_addr(fors_pk_addr, fors_addr);
    set_type(fors_pk_addr, ASGN_ADDR_TYPE_FORSPK);

    message_to_indices(indices, m);

    for (i = 0; i < ASGN_FORS_TREES; i++) {
        idx_offset = i * (1 << ASGN_FORS_HEIGHT);

        set_tree_height(fors_tree_addr, 0);
        set_tree_index(fors_tree_addr, indices[i] + idx_offset);
        set_type(fors_tree_addr, ASGN_ADDR_TYPE_FORSPRF);

        /* Include the secret key part that produces the selected leaf node. */
        fors_gen_sk(sig, ctx, fors_tree_addr);
        set_type(fors_tree_addr, ASGN_ADDR_TYPE_FORSTREE);
        sig += ASGN_N;

        /* Compute the authentication path for this leaf node. */
        treehashx1(roots + i*ASGN_N, sig, ctx,
                 indices[i], idx_offset, ASGN_FORS_HEIGHT, fors_gen_leafx1,
                 fors_tree_addr, &fors_info);

        sig += ASGN_N * ASGN_FORS_HEIGHT;
    }

    /* Hash horizontally across all tree roots to derive the public key. */
    thash(pk, roots, ASGN_FORS_TREES, ctx, fors_pk_addr);
}

/**
 * Derives the FORS public key from a signature.
 * This can be used for verification by comparing to a known public key, or to
 * subsequently verify a signature on the derived public key. The latter is the
 * typical use-case when used as an FTS below an OTS in a hypertree.
 * Assumes m contains at least ASGN_FORS_HEIGHT * ASGN_FORS_TREES bits.
 */
void fors_pk_from_sig(unsigned char *pk,
                      const unsigned char *sig, const unsigned char *m,
                      const ascon_sign_ctx* ctx,
                      const uint32_t fors_addr[8])
{
    uint32_t indices[ASGN_FORS_TREES];
    unsigned char roots[ASGN_FORS_TREES * ASGN_N];
    unsigned char leaf[ASGN_N];
    uint32_t fors_tree_addr[8] = {0};
    uint32_t fors_pk_addr[8] = {0};
    uint32_t idx_offset;
    unsigned int i;

    copy_keypair_addr(fors_tree_addr, fors_addr);
    copy_keypair_addr(fors_pk_addr, fors_addr);

    set_type(fors_tree_addr, ASGN_ADDR_TYPE_FORSTREE);
    set_type(fors_pk_addr, ASGN_ADDR_TYPE_FORSPK);

    message_to_indices(indices, m);

    for (i = 0; i < ASGN_FORS_TREES; i++) {
        idx_offset = i * (1 << ASGN_FORS_HEIGHT);

        set_tree_height(fors_tree_addr, 0);
        set_tree_index(fors_tree_addr, indices[i] + idx_offset);

        /* Derive the leaf from the included secret key part. */
        fors_sk_to_leaf(leaf, sig, ctx, fors_tree_addr);
        sig += ASGN_N;

        /* Derive the corresponding root node of this tree. */
        compute_root(roots + i*ASGN_N, leaf, indices[i], idx_offset,
                     sig, ASGN_FORS_HEIGHT, ctx, fors_tree_addr);
        sig += ASGN_N * ASGN_FORS_HEIGHT;
    }

    /* Hash horizontally across all tree roots to derive the public key. */
    thash(pk, roots, ASGN_FORS_TREES, ctx, fors_pk_addr);
}
