#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <dirent.h>
#include <unistd.h>
#include <sys/random.h>
#include <stdbool.h>
#include <assert.h>
#include "blst/bindings/blst.h"
#include "debugprint.h"
#include "fstoken.h"
#include "base64.h"

char* token_path;

void print_help(char* name){
    printf("FemtoStar Token Utility (ftu)\n");
    printf("This tool can be used to generate and process Credit Tokens for use with the FemtoStar Protocol.\n\n");
    
    printf("Warning: This tool lets you do insecure or broken things! Be careful with it.\n");
    printf("ftu is still under development! Do not assume it is secure or complete yet.\n");
    printf("In particular, note that keys and tokens are currently stored unencrypted on-disk.\n\n");
    
    printf("%s help - display this help\n", name);
    printf("%s path - display your token path\n", name);
    printf("%s list - list available targets - use \"%s list verbose\" to also display format specifiers and paths to target files\n", name, name);
    printf("%s keygen [targ] - create a new target [targ] using the default token format (128/256) and the system true randomness source\n", name);
    printf("%s keygen [targ] [tfs] - create a new target[targ] using token format [tfs] and the system true randomness source\n", name);
    printf("%s keygen [targ] [tfs] [ikm] - create a new target [targ] using token format specifier [tfs] and 32-byte hexadecimal seed [ikm]\n", name);
    printf("%s info [targ] - dump all available target information for target [targ], including keys\n", name);
    //printf("%s keyrepair [targ] - regenerate a missing Public Key for a target [targ] for which a Secret Key is available\n", name);
    printf("%s strip [targin] [targout] - duplicate target [targin] as [targout], removing Secret Key (for distribution as a Public-Key-only target)\n", name);
    printf("%s ireq [targ] - interactively generate a Token Request Code for target [targ] using a random Token ID and Blinding Factor\n", name);
    printf("%s ireq [targ] [id] - interactively generate a Token Request Code for target [targ] using a set Token ID [id] and a random Blinding Factor\n", name);
    printf("%s ireq [targ] [id] [factor] - interactively generate a Token Request Code for target [targ] using a set Token ID [id] and Blinding Factor [factor]\n", name);
    printf("%s sign [targ] [req] - generate a signature for Token Request Code [req], on target [targ]\n", name);
    printf("%s bsign [targ] - bulk-sign token requests, accepting Token Request Codes on stdin and outputting signatures to stdout\n", name);
    printf("%s verify [targ] [token] - verify token [token] against target [targ], as base64 - returns 0 on verified\n", name);
    printf("%s bverify [targ] - bulk-verify tokens against target [targ], accepting base64 tokens on stdin and outputting 0 (verified) or 1 to stdout\n", name);
    printf("%s inspect [targ] [token] - inpect a token [token] for target [targ], to see its token ID and signature\n", name);
    printf("%s shorten [targ] [token] - shorten a token [token] for target [targ], generating a symmetrically-verifiable short token\n", name);
}

int get_targ_path(char* target, char** targ_path){
    int path_len;

    path_len = strlen(token_path) + strlen(target) + 17; //string length of "/targets/", ".target", and null terminator

    *targ_path = malloc(path_len);
    if(*targ_path == NULL) return 1;

    strcpy(*targ_path, token_path);
    strcat(*targ_path, "/targets/");
    strcat(*targ_path, target);
    strcat(*targ_path, ".target");

    return 0;
}

int get_target(char* targ_name, fstoken_target_descriptor* target){
    FILE *targ_file;
    char *targ_path;

    get_targ_path(targ_name, &targ_path);
    
    targ_file = fopen(targ_path, "r");
    if(!targ_file){
        return 0;
    }
    fread(target, sizeof(fstoken_target_descriptor), 1, targ_file);
    fclose(targ_file);

    // 0 = no keys (bad target), 1 = PK only, 2 = PK+SK
    return target->sk_available + 1;
}

// Fill a buffer dest of size dest_size with exactly (dest_size - 2) characters and a null terminator, returning 1 if the user supplies any other number
int get_fixed_length(char* dest, size_t dest_size){
    dest[dest_size - 2] = '\0'; // Initialize this place, we're about to need to check it even if fgets never gets here
    fgets(dest, dest_size, stdin);
    if(dest[dest_size - 2] != '\n') return 1; //If there isn't a newline in the second-last character (i.e. the user specified too many) return 1
    
    dest[dest_size - 2] = '\0'; // Cut it down to length and get rid of the newline we're now sure is in the right place

    return 0;
}

int decode_base64_token(fstoken_target_descriptor* target, fstoken_token* token, char* token_b64){
    int token_length, target_token_length;

    target_token_length = fstoken_target_token_length(target);
    token_length = base64_binlength(token_b64);

    if(token_length == -1){
        fprintf(stderr, "Invalid base64 provided\n");
        return 3; // should be an enum
    }

    if(token_length != target_token_length){
        fprintf(stderr, "Provided token (%i bytes) of incorrect length for specified target (%i bytes).\n", token_length, target_token_length);
        return 4;
    }

    base64_decode_oneshot(token_b64, strlen(token_b64), (char*) token);
    
    return 0;
}

int inspect_token(char* targ_name, char* token_b64){
    fstoken_target_descriptor target;
    fstoken_token token;
    int targ_idbytes;

    if(!get_target(targ_name, &target)){
        printf("Invalid target %s\n", targ_name);
        return -1;
    }

    if(decode_base64_token(&target, &token, token_b64)){
        return 2;
    }

    if(fstoken_verify_token(&token, &target) == 0){
        printf("Token is valid for target %s.\n", targ_name);
    }else{
        printf("Token is not valid for target!\n");
        return 1;
    }

    targ_idbytes = (target.idbits + 7) / 8;

    printf("\n");
    print_bytes("Signature (hex): ", token.signature, 48, ' ');
    printf("Signature (base64): ");
    print_base64(token.signature, 48);
    printf("\n");
    print_bytes("Token ID (hex): ", token.token_id, targ_idbytes, ' ');
    printf("Token ID (base64): ");
    print_base64(token.token_id, targ_idbytes);

    return 0;
}

// Receives tokens on stdin and provides 0 (valid), 1 (invalid), or 2/3/4 (various bad input) to stdout, with additional messages on stderr in case of error
int verify_bulk(char* targ_name){
    char token_b64[((((HASHBYTES_MAX + IDBYTES_MAX) + 2) / 3) * 4) + 2]; // maximum possible base64 length + 2 characters for get_fixed_length
    fstoken_target_descriptor target;
    fstoken_token token;
    int target_token_b64_length, result;

    if(!get_target(targ_name, &target)){
        fprintf(stderr, "Invalid target %s\n", targ_name);
        return -1;
    }

    target_token_b64_length = base64_base64length(fstoken_target_token_length(&target));

    printf("Target token length %i\n", target_token_b64_length);

    while(1){
        if(get_fixed_length(token_b64, target_token_b64_length + 2)){ // +2 because of get_fixed_length, I should probably change that to be not stupid
            printf("2");
            fprintf(stderr, "Provided token base64 incorrect length for target (should be %i characters).\n", target_token_b64_length);
            continue;
        }

        result = decode_base64_token(&target, &token, token_b64);

        if(result){
            printf("%i", result);
            continue;
        }

        if(fstoken_verify_token(&token, &target) == 0){
            printf("0");
        }else{
            printf("1");
        }
    }
}

int verify_single(char* targ_name, char* token_b64){
    fstoken_target_descriptor target;
    fstoken_token token;

    if(!get_target(targ_name, &target)){
        printf("Invalid target %s\n", targ_name);
        return -1;
    }

    decode_base64_token(&target, &token, token_b64);

    if(fstoken_verify_token(&token, &target) == 0){
        printf("Verified!\n");
        return 0;
    }else{
        printf("Not verified!\n");
        return 1;
    }
}

int sign_bulk(char* targ_name){
    char req_b64[66];
    fstoken_request req;
    fstoken_target_descriptor targ;
    fstoken_signature sig;

    if(get_target(targ_name, &targ) != 2){
        printf("Target or private key unavailable. Are you sure you have a private key for this target?\n");
        return 1;
    }

    fprintf(stderr, "Bulk signing started. Enter token requests\n");

    while(1){
        if(get_fixed_length(req_b64, sizeof(req_b64))){
            fprintf(stderr, "Invalid request length\n");
            continue;
        };

        base64_decode_oneshot((char*) req_b64, 64, (char*) req.id_blind);
        
        //debug_print_bytes("Signature hex: ", req.id_blind, 48);

        if(fstoken_sign_request(&req, &sig, &targ)){
            fprintf(stderr, "Signing failed. Is the token request valid?\n");
            continue;
        };

        print_base64(sig.signature, 48);
    }

    return 0;
}

int sign_single(char* targ_name, char* req_b64){
    fstoken_request req;
    fstoken_target_descriptor targ;
    fstoken_signature sig;

    base64_decode_oneshot((char*) req_b64, 64, (char*) req.id_blind);
    
    //debug_print_bytes("Signature hex: ", req.id_blind, 48);

    if(get_target(targ_name, &targ) != 2){
        printf("Target or private key unavailable. Are you sure you have a private key for this target?\n");
        return 1;
    }

    if(fstoken_sign_request(&req, &sig, &targ)){
        printf("Signing failed. Is the token request valid?\n");
        return 1;
    };

    print_base64(sig.signature, 48);

    return 0;
}

/*int verify_token(char* targ_name, byte* token_bytes){
    fstoken_target_descriptor target;
    fstoken_token token;
    int token_length;

    if(!get_target(targ_name, &target)){
        printf("Invalid target %s\n", targ_name);
        return 1;
    }

    //base64_decode_oneshot(token_bytes, )
}*/

int interactive_request(char* targ_name, byte* token_id, byte* blinding_factor){
    fstoken_request req;
    fstoken_request_descriptor reqdesc;
    fstoken_target_descriptor target;
    fstoken_signature sig;
    fstoken_token token;
    char req_base64_encoded[65];
    char sig_base64[66];
    int token_length;

    if(!get_target(targ_name, &target)){
        printf("Invalid target %s\n", targ_name);
        return 1;
    }

    fstoken_gen_request(&req, &reqdesc, &target, token_id, blinding_factor);

    base64_encode_oneshot((char*) req.id_blind, 48, req_base64_encoded);

    printf("A Token Request Code has been generated. Provide this code to the signer. Once you have received a signature, enter it below to generate your token.\n");
    printf("Token Request Code: %s\n", req_base64_encoded);
    printf("Enter your signature: ");
    
    if(get_fixed_length(sig_base64, sizeof(sig_base64))){
        printf("Invalid signature length\n");
        return 1;
    };
    
    base64_decode_oneshot(sig_base64, 64, (char*) sig.signature);

    //debug_print_bytes("Signature as hex: ", sig.signature, 48);

    if(fstoken_gen_token(&reqdesc, &sig, &token)){
        printf("Failed to generate token.\n");
        return 1;
    };

    token_length = 48 + (target.idbits + 7) / 8;
    
    //debug_print_bytes("Token Unblinded Signature: ", token.signature, 48);
    //debug_print_bytes("TOKEN: ", &token, token_length);

    printf("\nToken: ");
    print_base64(&token, token_length);

    //token.token_id[3] = '\0';

    switch(fstoken_verify_token(&token, &target)){
        case 0:
            printf("Verified!\n");
            break;
        case 1:
            printf("Signature does not match Token ID for target. Did the signer sign for the wrong target or sign the wrong request?\n");
            return 1;
        case 2:
            printf("Invalid signature (not in G1 group - either something has gone very wrong or the signer is attempting to do something malicious)\n");
            return 1;
    }

    return 0;
}

int targinfo(char* targ_name){
    int targ_status;
    fstoken_target_descriptor target;
    
    targ_status = get_target(targ_name, &target);

    switch(targ_status){
        case 0:
            printf("No target file found - target unknown.\n");
            return 1;
        case 1:
            printf("Public Key available - can verify and request for this target\n");
            print_bytes("Public Key: ", target.pk, 96, '\0');
            break;
        case 2:
            printf("Secret Key and Public Key available - can verify, request, and sign for this target.\n");
            print_bytes("Secret Key: ", target.sk, 32, '\0');
            print_bytes("Public Key: ", target.pk, 96, '\0');
            break;
    }

    printf("Token Format Specifier: %i/%i (%i ID bits, %i hash bits)\n", target.idbits, target.hashbits, target.idbits, target.hashbits);

    return 0;
}

int striptarg(char* targ_name_in, char* targ_name_out){
    FILE *targ_out_file;
    fstoken_target_descriptor targ_out;
    char* targ_out_path;

    if(!get_target(targ_name_in, &targ_out)){
        printf("Unknown input target. Exiting.\n");
        return 1;
    }

    if(targ_out.sk_available == false){
        printf("Target %s already lacks a Secret Key - only targets with Secret Keys can be stripped.\n", targ_name_in);
        return 2;
    }

    memset(targ_out.sk, 0x00, 32);
    targ_out.sk_available = false;

    get_targ_path(targ_name_out, &targ_out_path);

    targ_out_file = fopen(targ_out_path, "w");
    if(!targ_out_file){
        printf("Could not open output target file. Exiting.\n");
        return 1;
    }
    fwrite(&targ_out, sizeof(fstoken_target_descriptor), 1, targ_out_file);
    fclose(targ_out_file);

    printf("Public-Key-only target %s succesfully created from original target %s.\n", targ_name_in, targ_name_out);

    return 0;
}

/*int keyrepair(char* target){ // this is probably no longer necessary with single-file targets
    FILE *key_file;
    byte sk[32];
    byte pk[96];
    char* sk_path;
    char* pk_path;
    char* tfs_path;
    int key_status;
    
    key_status = get_keys(target, sk, NULL, NULL, NULL);

    if(key_status != 2){
        printf("This target does not refer to a keypair with only a Secret Key available. Exiting.\n");
        return 1;
    }

    printf("Regenerating Public Key from Private Key for broken keypair %s\n", target);

    fstoken_get_pk_from_sk(sk, pk);
    debug_print_bytes("Regenerated Public Key: ", pk, 96);

    get_key_paths(target, &sk_path, &pk_path, &tfs_path);

    key_file = fopen(pk_path, "w");
    if(!key_file){
        printf("Could not open Public Key file. Exiting.\n");
        return 1;
    }
    fwrite(pk, 96, 1, key_file);
    fclose(key_file);

    printf("Saved to %s\n", pk_path);

    return 0;
}*/

int keygen(char* target, byte* ikm, int idbits, int hashbits){
    char *targ_path;
    FILE *targ_file;
    fstoken_target_descriptor targ;

    debug_print_bytes("IKM: ", ikm, 32);

    fstoken_keygen(ikm, targ.sk, targ.pk);

    debug_print_bytes("Secret Key: ", targ.sk, 32);
    debug_print_bytes("Public Key: ", targ.pk, 96);

    targ.idbits = idbits;
    targ.hashbits = hashbits;
    targ.target_descriptor_version = 0x00;
    targ.sk_available = true;

    debug_print_bytes("Target Struct: ", &targ, sizeof(targ));

    if(get_targ_path(target, &targ_path)) return 1;

    targ_file = fopen(targ_path, "w");
    if(!targ_file){
        printf("Could not open target file. Exiting.\n");
        return 1;
    }
    fwrite(&targ, sizeof(fstoken_target_descriptor), 1, targ_file);
    fclose(targ_file);

    return 0;
}

void print_path(){
    printf("Token Path (from FEMTOSTAR_TOKEN_PATH environment variable): %s\n", token_path);
}

bool string_endswith(const char *str, const char *suffix){
    if (!str || !suffix)
        return 0;
    size_t lenstr = strlen(str);
    size_t lensuffix = strlen(suffix);
    if (lensuffix >  lenstr)
        return 0;
    return strncmp(str + lenstr - lensuffix, suffix, lensuffix) == 0;
}

// Believe it or not, this function used to be WORSE, and substantially so. It could still do to be cleaned up.
int list_targets(bool verbose){
    printf("Listing all targets - you have secret keys for, and can issue tokens for, targets marked with (*)\n\n");
    int n, targname_len;
    struct dirent **files;
    char *targdir_path, *targ_path, *targ_name;
    fstoken_target_descriptor targ;

    targdir_path = malloc(strlen(token_path) + 9);
    if(targdir_path == NULL) return 1;

    strcpy(targdir_path, token_path);
    strcat(targdir_path, "/targets");

    #ifndef __INTELLISENSE__ // VSCodium doesn't know where alphasort is and highlights an error
    n = scandir(targdir_path, &files, NULL, alphasort);
    #endif

    if(n == -1){
        fprintf(stderr, "Could not list targets directory in token path. Does it exist?\n");
        exit(1);
    }

    for(int i=0;i<n;i++){
        if(string_endswith(files[i]->d_name, ".target")){
            targname_len = strlen(files[i]->d_name);

            targ_name = malloc(targname_len + 1);
            if(targ_name == NULL) return 1;

            strcpy(targ_name, files[i]->d_name);            
            targ_name[targname_len - 7] = '\0';

            printf("%s", targ_name);

            get_target(targ_name, &targ);
            if(targ.sk_available) printf(" (*)");

            if(verbose){
                get_targ_path(targ_name, &targ_path);
                printf(" (Path: %s, Token Format Specifier: %i/%i)", targ_path, targ.idbits, targ.hashbits);
            }

            printf("\n");
            free(targ_name);
        }
    }
    free(targdir_path);

    return 0;
}

void bendian_from_hex_string(byte* bendian, char* string, int length){
    char byte[2];
    for(int i=0; i<length; i++){
        memcpy(byte, &string[i*2], 2);
        bendian[i] = strtol(byte, 0, 16);
    }
}

void lendian_from_hex_string(byte* lendian, char* string, int length){
    char byte[2];
    for(int i = (length-1); i >= 0; i--){
        memcpy(byte, &string[i*2], 2);
        lendian[i] = strtol(byte, 0, 16);
    }
}

// mostly boring command line parsing
int main(int argc, char *argv[]){
    token_path = getenv("FEMTOSTAR_TOKEN_PATH");

    if(!token_path){
        fprintf(stderr, "The environment variable FEMTOSTAR_TOKEN_PATH does not exist! Please set it before using ftu.\n");
        exit(1);
    }

    if(argc < 2){
        fprintf(stderr, "Provide at least one argument. Try \"%s help\" for more information.\n", argv[0]);
        return 1;
    }else if(strcmp(argv[1], "help") == 0){
        print_help(argv[0]);
        return 0;
    }else if(strcmp(argv[1], "path") == 0){
        print_path();
        return 0;
    }else if(strcmp(argv[1], "list") == 0){
        return list_targets(argc > 2 && strcmp(argv[2], "verbose") == 0); // i don't know if this is cursed or genius
    }else if(strcmp(argv[1], "keygen") == 0){
        byte ikm[32];
        int idbits, hashbits;

        if(argc > 5){
            printf("Too many arguments. Exiting.\n");
            return(1);
        }

        // Make sure there's a target name
        if(argc < 3){
            fprintf(stderr, "A target name must be provided, e.g. %s keygen [targ]\n", argv[0]);
            return(1);
        }

        // Default behaviour for if only target name is provided: default TFS, random IKM. Otherwise, validate and use provided.
        if(argc < 4){
            idbits = IDBITS_DEFAULT;
            hashbits = HASHBITS_DEFAULT;           
        }else{
            sscanf(argv[3], "%i/%i", &idbits, &hashbits);
            if(idbits < 1 || idbits > IDBITS_MAX){
                printf("Invalid Token Format Specifier: number of ID bits must be between 1 and 256 inclusive\n");
                return 1;
            }
            if(hashbits < 1 || hashbits > HASHBITS_MAX){
                printf("Invalid Token Format Specifier: number of hash bits must be between 1 and 256 inclusive\n");
                return 1;
            }
        }
        
        // If no IKM is provided, use the system true randomness source
        if(argc < 5){
            getrandom(ikm, 32, GRND_RANDOM);
        }else{
            if(strlen(argv[4]) != 64){
                fprintf(stderr, "If providing IKM, it must be 32 bytes (64 hexadecimal digits)\n");
                return 1;
            }

            lendian_from_hex_string(ikm, argv[4], 32);
        }

        return keygen(argv[2], ikm, idbits, hashbits);
    }else if(strcmp(argv[1], "info") == 0){
        // Make sure there's a target name
        if(argc < 3){
            fprintf(stderr, "A target name must be provided, e.g. %s info [targ]\n", argv[0]);
            return(1);
        }

        return targinfo(argv[2]);
    }/*else if(strcmp(argv[1], "keyrepair") == 0){
        // Make sure there's a target name
        if(argc < 3){
            fprintf(stderr, "A target name must be provided, e.g. %s keyrepair [targ]\n", argv[0]);
            return(1);
        }

        return keyrepair(argv[2]);
    }*/else if(strcmp(argv[1], "strip") == 0){
        // Make sure there's a target name
        if(argc < 4){
            fprintf(stderr, "Two target names must be provided, e.g. %s strip [targin] [targout]\n", argv[0]);
            return(1);
        }

        return striptarg(argv[2], argv[3]);
    }else if(strcmp(argv[1], "ireq") == 0){
        byte blinding_factor[32];
        byte token_id[IDBYTES_MAX];

        // Make sure there's a target name
        if(argc < 3){
            fprintf(stderr, "A target name must be provided, e.g. %s req [targ]\n", argv[0]);
            return(1);
        }

        if(argc > 5){
            printf("Too many arguments. Exiting.\n");
            return(1);
        }

        if(argc < 4){
            getrandom(token_id, IDBYTES_MAX, GRND_NONBLOCK);
        }else{
            if(strlen(argv[3]) > (IDBYTES_MAX * 2)){
                fprintf(stderr, "If providing a Token ID, it must be %i bits (%i hexadecimal digits) or less \n", IDBITS_MAX, IDBYTES_MAX * 2);
                return 1;
            }

            lendian_from_hex_string(token_id, argv[3], (strlen(argv[3]) + 1) / 2);
        }
        
        // If no IKM is provided, use the system true randomness source
        if(argc < 5){
            getrandom(blinding_factor, 32, GRND_NONBLOCK);
        }else{
            if(strlen(argv[4]) != 64){
                fprintf(stderr, "If providing a blinding factor, it must be 32 bytes (64 hexadecimal digits)\n");
                return 1;
            }

            lendian_from_hex_string(blinding_factor, argv[4], 32);
        }

        return interactive_request(argv[2], token_id, blinding_factor);
    }else if(strcmp(argv[1], "sign") == 0){
        if(argc != 4){
            fprintf(stderr, "Please provide a target [targ] and a Token Request Code [req]. See help.\n");
            return 1;
        }

        if(strlen(argv[3]) != 64){
            fprintf(stderr, "Token request code must be 48 bytes (64 base64 characters).\n");
            return 1;
        }

        return sign_single(argv[2], argv[3]);
    }else if(strcmp(argv[1], "bsign") == 0){
        if(argc != 3){
            fprintf(stderr, "Please provide a target [targ]. See help.\n");
            return 1;
        }

        return sign_bulk(argv[2]);
    }else if(strcmp(argv[1], "verify") == 0){
        if(argc != 4){
            fprintf(stderr, "Please provide a target [targ] and a token [token]. See help.\n");
            return 1;
        }

        if(strlen(argv[3]) < 68){
            fprintf(stderr, "Token must be at least 49 bytes (68 base64 characters).\n");
            return 1;
        }

        return(verify_single(argv[2], argv[3]));
    }else if(strcmp(argv[1], "bverify") == 0){
        if(argc != 3){
            fprintf(stderr, "Please provide a target [targ]. See help.\n");
            return 1;
        }

        verify_bulk(argv[2]);
    }else if(strcmp(argv[1], "inspect") == 0){
        if(argc != 4){
            fprintf(stderr, "Please provide a target [targ] and a token [token]. See help.\n");
            return 1;
        }

        if(strlen(argv[3]) < 68){
            fprintf(stderr, "Token must be at least 49 bytes (68 base64 characters).\n");
            return 1;
        }

        return(inspect_token(argv[2], argv[3]));
    }
}