#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_key_paths(char* target, char** sk_path, char** pk_path, char** tfs_path){
    int key_path_len;

    key_path_len = strlen(token_path) + strlen(target) + 13;

    *sk_path = malloc(key_path_len);
    if(*sk_path == NULL) return 1;

    *pk_path = malloc(key_path_len);
    if(*pk_path == NULL) return 1;

    *tfs_path = malloc(key_path_len + 1);
    if(*tfs_path == NULL) return 1;

    strcpy(*sk_path, token_path);
    strcat(*sk_path, "/targets/");
    strcat(*sk_path, target);
    strcpy(*pk_path, *sk_path);
    strcpy(*tfs_path, *sk_path);
    strcat(*sk_path, ".sk");
    strcat(*pk_path, ".pk");
    strcat(*tfs_path, ".tfs");

    return 0;
}

int get_keys(char* target, byte* sk, byte* pk, int* idbits, int* hashbits){ // pk/sk/idbits/hasbits pointers can be NULL if you don't want to read those
    FILE *targ_file;
    char *sk_path;
    char *pk_path;
    char *tfs_path;
    bool sk_available, pk_available, tfs_available;
    int idbits_buf, hashbits_buf;

    get_key_paths(target, &sk_path, &pk_path, &tfs_path);
    
    sk_available = (access(sk_path, R_OK) == 0);
    pk_available = (access(pk_path, R_OK) == 0);
    tfs_available = (access(tfs_path, R_OK) == 0);

    if(sk_available && sk != NULL){
        targ_file = fopen(sk_path, "r");
        if(!targ_file){
            printf("Could not open Secret Key file. Exiting.\n");
            return 1;
        }
        fread(sk, 32, 1, targ_file);
        fclose(targ_file);
    }

    if(pk_available && pk != NULL){
        targ_file = fopen(pk_path, "r");
        if(!targ_file){
            printf("Could not open Public Key file. Exiting.\n");
            return 1;
        }
        fread(pk, 96, 1, targ_file);
        fclose(targ_file);
    }

    if(idbits != NULL || hashbits != NULL){
        if(tfs_available){
            targ_file = fopen(tfs_path, "r");
            if(!targ_file){
                printf("Could not open Token Format Specifier file. Exiting.\n");
                return 1;
            }
            fscanf(targ_file, "%i/%i", &idbits_buf, &hashbits_buf);
            fclose(targ_file);

            if(idbits != NULL) *idbits = idbits_buf;
            if(hashbits != NULL) *hashbits = hashbits_buf;
        }else{
            printf("WARNING: Token Format Specifier not set, this is a broken state. Using default (128/256) - please add a .tfs file for the target\n");
            if(idbits != NULL) *idbits = IDBITS_DEFAULT;
            if(hashbits != NULL) *hashbits = HASHBITS_DEFAULT;
        }
    }

    // 0 = no keys (bad target), 1 = PK only, 2 = SK only (broken state), 3 = PK+SK (can sign)
    return (2 * sk_available) + pk_available;
}

int keydump(char* target){
    byte sk[32];
    byte pk[96];
    int key_status;
    int idbits, hashbits;
    
    key_status = get_keys(target, sk, pk, &idbits, &hashbits);

    switch(key_status){
        case 0:
            printf("No keys found - target unknown.\n");
            break;
        case 1:
            printf("Public Key available - can verify and request for this target\n");
            print_bytes("Public Key: ", pk, 96);
            break;
        case 2:
            printf("Secret Key ONLY available - this is a broken state, please keyrepair this keypair (see help)\n");
            print_bytes("Secret Key: ", sk, 32);
            break;
        case 3:
            printf("Secret Key and Public Key available - can verify, request, and sign for this target.\n");
            print_bytes("Secret Key: ", sk, 32);
            print_bytes("Public Key: ", pk, 96);
            break;
    }

    printf("Token Format Specifier: %i/%i (%i ID bits, %i hash bits)\n", idbits, hashbits, idbits, 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 *sk_path;
    char *pk_path;
    char* tfs_path;
    FILE *targ_file;
    byte sk_byte[32];
    byte pk_byte[96];

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

    fstoken_keygen(ikm, sk_byte, pk_byte);

    debug_print_bytes("Secret Key: ", sk_byte, 32);
    debug_print_bytes("Public Key: ", pk_byte, 96);

    if(get_key_paths(target, &sk_path, &pk_path, &tfs_path)) return 1;

    printf("Writing Secret Key to %s\n", sk_path);

    targ_file = fopen(sk_path, "w");
    if(!targ_file){
        printf("Could not open Secret Key file. Exiting.\n");
        return 1;
    }
    fwrite(sk_byte, 32, 1, targ_file);
    fclose(targ_file);

    printf("Writing Public Key to %s\n", pk_path);

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

    printf("Writing Token Format Specifier to %s\n", tfs_path);

    targ_file = fopen(tfs_path, "w");
    if(!targ_file){
        printf("Could not open Token Format Specifier file. Exiting.\n");
        return 1;
    }
    fprintf(targ_file, "%i/%i", idbits, hashbits);
    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;
}

// This function is awful because strings in C. It should probably be improved.
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, keyname_len;
    struct dirent **files;
    char *keydir_path, *key_path, *key_name;
    bool sk_available;

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

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

    #ifndef __INTELLISENSE__ // VSCodium doesn't know where alphasort is and highlights an error
    n = scandir(keydir_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, ".pk")){
            keyname_len = strlen(files[i]->d_name);

            key_name = malloc(keyname_len + 1);
            if(key_name == NULL) return 1;

            strcpy(key_name, files[i]->d_name);            
            key_name[keyname_len - 3] = '\0';

            printf("%s", key_name);
            
            key_path = malloc(strlen(token_path) + 9 + strlen(files[i]->d_name));
            if(key_path == NULL) return 1;

            strcpy(key_path, token_path);
            strcat(key_path, "/targets/");
            strcat(key_path, files[i]->d_name);

            if(verbose) printf(" (PK: %s", key_path);

            key_path[strlen(key_path) - 2] = 's';

            if(access(key_path, R_OK) == 0){
                sk_available = true;

                if(verbose) printf(", SK: %s", key_path);
            }else{
                sk_available = false;
            }

            if(verbose) printf(")");
            if(sk_available) printf(" (*)");

            printf("\n");
            free(key_path);
            free(key_name);
        }
    }
    free(keydir_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]);
    }
}