#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/blst.h"
#include "debugprint.h"
#include "fstoken.h"

char* token_path;

void print_help(char* name){
    printf("FemtoStar Credit Token Manager (ctm)\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("ctm 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 targets you have keys for - use \"%s list verbose\" to also display paths to the keys\n", name, name);
    printf("%s keygen [targ] - create a new target keypair [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 keypair [targ] using token format [tfs] and the system true randomness source\n", name);
    printf("%s keygen [targ] [tfs] [ikm] - create a new target keypair [targ] using token format specifier [tfs] and 32-byte hexadecimal seed [ikm]\n", name);
    printf("%s keydump [targ] - dump Public and, if available, Secret Keys for target [targ]\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 req [targ] - generate a token request for target [targ]\n", name);
}

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

    path_len = strlen(token_path) + strlen(target) + 17;

    *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, target_descriptor* target){ // pk/sk/idbits/hasbits pointers can be NULL if you don't want to read those
    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(target_descriptor), 1, targ_file);
    fclose(targ_file);

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

int keydump(char* targ_name){
    int targ_status;
    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);
            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);
            print_bytes("Public Key: ", target.pk, 96);
            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 keyrepair(char* target){
    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;
    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(targ), 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;
    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 directory at token path.\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);
    }
}

// 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 ctm.\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");
        }

        // 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;
            }

            bendian_from_hex_string(ikm, argv[4], 64);
        }

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

        return keydump(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]);
    }*/
}