#include "fstoken.h"

void fstoken_keygen(byte* ikm, byte* sk_byte, byte* pk_byte){
    blst_scalar sk;
    blst_p2 pk;
    blst_p2_affine pk_affine;

    blst_keygen(&sk, ikm, 32, 0, 0); // generate a secret key from IKM
    blst_bendian_from_scalar(sk_byte, &sk); // convert it to 32 big-endian bytes in sk_byte to return

    blst_sk_to_pk_in_g2(&pk, &sk); // get a public key from the secret key
    blst_p2_to_affine(&pk_affine, &pk); // convert it to an affine point, which is what most uses of the public key use
    blst_p2_affine_compress(pk_byte, &pk_affine); // compress it to 96 bytes in pk_byte to return
}

void fstoken_get_pk_from_sk(byte* sk_byte, byte* pk_byte){
    blst_p2 pk;
    blst_p2_affine pk_affine;
    blst_scalar sk;

    blst_scalar_from_bendian(&sk, sk_byte); //convert the SK to a blst_scalar
    blst_sk_to_pk_in_g2(&pk, &sk); // get a PK from it
    blst_p2_to_affine(&pk_affine, &pk); // convert the PK to affine
    blst_p2_affine_compress(pk_byte, &pk_affine); // compress the PK
}

void fstoken_gen_request(fstoken_request* req, fstoken_request_descriptor* reqdesc, fstoken_target_descriptor* target, byte* token_id, byte* blinding_factor){
    byte debug_print_buf[256]; //get rid of this
    
    int targ_idbytes;
    blst_p1 hash, blinded_id;
    blst_scalar blinding_r;

    targ_idbytes = (target->idbits + 7)/8;

    //printf("IDBYTES_MAX: %i, IDBYTES for target: %i\n", IDBYTES_MAX, targ_idbytes);

    memcpy(&reqdesc->blinding_factor, blinding_factor, 32);

    bit_truncated_copy(token_id, reqdesc->token_id, IDBYTES_MAX, target->idbits); //cut down the supplied token ID to match the target's IDBITS value

    //debug_print_bytes("Working with ID: ", reqdesc->token_id, IDBYTES_MAX);
    //debug_print_bytes("Working with Blinding Factor: ", blinding_factor, 32);

    // Hash the Token ID to a BLST P1
    blst_hash_to_g1(&hash, token_id, targ_idbytes, FSTOKEN_DST, strlen((char *) FSTOKEN_DST), target->pk, 96);

    // Get the blinding factor as a BLST scalar
    blst_scalar_from_bendian(&blinding_r, blinding_factor);

    blst_p1_serialize(debug_print_buf, &hash);
    //debug_print_bytes("HASH: ", debug_print_buf, 96);

    // Sign the hash as if the blinding factor was a private key
    blst_sign_pk_in_g2(&blinded_id, &hash, &blinding_r);

    // Compress the resulting signature and put it in the signature request
    blst_p1_compress(req->id_blind, &blinded_id);

    //debug_print_bytes("Blinded and compressed for wire: ", req->id_blind, 48);
}

int fstoken_sign_request(fstoken_request* req, fstoken_signature* sig, fstoken_target_descriptor* targ){
    blst_scalar sk;
    blst_p1 req_id, signature;
    blst_p1_affine req_id_affine;

    // get the secret key as a scalar
    blst_scalar_from_bendian(&sk, targ->sk);

    // Deserialize the request - it's already a serialized P1 point, we don't need to (literally) rehash it
    blst_p1_uncompress(&req_id_affine, req->id_blind);

    // i do not know why deserializing always gives you affine points
    blst_p1_from_affine(&req_id, &req_id_affine);

    // Return 1 if the point isn't in G1, this is basically a sanity check
    if(!blst_p1_in_g1(&req_id)) return 1;

    // sign with it
    blst_sign_pk_in_g2(&signature, &req_id, &sk);

    // Compress and print the signature
    blst_p1_compress(sig->signature, &signature);

    //debug_print_bytes("Compressed signature:", sig->signature, 48);

    return 0;
}

int fstoken_gen_token(fstoken_request_descriptor* req, fstoken_signature* sig, fstoken_token* token){
    blst_p1 returned_signature, unblinded_signature;
    blst_p1_affine returned_signature_affine;
    blst_scalar blinding_r, inverse_r;

    // We now have the signature back. returned_signature is a blst_p1_affine because this is pk_in_g2.
    blst_p1_uncompress(&returned_signature_affine, sig->signature);

    // Convert the decompressed returned signature from an affine to a P1
    blst_p1_from_affine(&returned_signature, &returned_signature_affine);

    // Confirm the signature point is in the G1 group
    if(!blst_p1_in_g1(&returned_signature)) return 1;

    // Retreive the blinding factor and take its inverse
    blst_scalar_from_bendian(&blinding_r, req->blinding_factor);
    blst_sk_inverse(&inverse_r, &blinding_r);

    //print_scalar("Inverse R: ", &inverse_r);

    // Unblind the signature using the inverse of the blinding factor
    blst_sign_pk_in_g2(&unblinded_signature, &returned_signature, &inverse_r);

    memcpy(token->token_id, req->token_id, IDBYTES_MAX);
    blst_p1_compress(token->signature, &unblinded_signature);

    return 0;
}

// 0 = valid token, 1 = signature check failed, 2 = signature point invalid
int fstoken_verify_token(fstoken_token* token, fstoken_target_descriptor* targ){
    int targ_idbytes;
    blst_p1_affine sig;
    blst_p2_affine pk;

    blst_p1_uncompress(&sig, token->signature);
    blst_p2_uncompress(&pk, targ->pk);

    BLST_ERROR returned;

    if(!blst_p1_affine_in_g1(&sig)) return 2;

    targ_idbytes = (targ->idbits + 7) / 8;

    returned = blst_core_verify_pk_in_g2(&pk, &sig, 1, token->token_id, targ_idbytes, FSTOKEN_DST, strlen((char *) FSTOKEN_DST), targ->pk, 96);

    return (returned != BLST_SUCCESS);
}

void bit_truncated_copy(byte* src, byte* dest, size_t destlen, int bits){
    uint8_t bitmask;
    int bytes;

    bytes = (bits + 7) / 8;

    memset(dest, 0, destlen); //clear out the destination
    memcpy(dest, src, bytes); //copy over as many bytes as are needed for the number of bits

    //bitmask out the high (fsp is little-endian) bits of the destination if bits is not a round number of bytes
    bitmask = 0xFF >> (bits % 8);
    //printf("mask: 0x%0x\n", bitmask);
    dest[bytes - 1] = dest[bytes - 1] & bitmask;
}

// Helper function to get the length (in bytes) of an (unshortened) token for the target
int fstoken_target_token_length(fstoken_target_descriptor* target){
    return 48 + (target->idbits + 7) / 8;
}