#include "network/network.h"
#include "client/client.h"
#include "server/server.h"
#include "tlibc/tlibc.h"
#include "tlibc/base64.h"

#define _DEFAULT_CONFIG_PATH_CLIENT "tcp-chat-client.config"
#define _DEFAULT_CONFIG_PATH_SERVER "tcp-chat-server.config"

typedef enum ProgramMode {
    ClientMode,
    ServerMode,
    RsaGenStdin,
    RsaGenRandom,
    RandomBytes,
    RandomBytesBase64,
} ProgramMode;

#define arg_is(LITERAL) str_equals(arg_str, STR(LITERAL))

int main(const int argc, cstr const* argv){
    Deferral(32);

    try_fatal_void(tlibc_init());
    Defer(tlibc_deinit());
    try_fatal_void(network_init());
    Defer(network_deinit());

    if(br_prng_seeder_system(NULL) == NULL){
        printfe("Can't get system random seeder. Bearssl is compiled incorrectly.");
        return 1;
    }

    ProgramMode mode = ClientMode;
    cstr config_path = NULL;
    u32 size_arg = 0;

    for(int argi = 1; argi < argc; argi++){
        str arg_str = str_from_cstr(argv[argi]);
        if(arg_is("-h") || arg_is("--help")){
            printf(
                "USAGE:\n"
                "no arguments                   Interactive client mode.\n"
                "-h, --help                     Show this message.\n"
                "-l, --listen                   Start server.\n"
                "--config [path]                Load config from specified path.\n"
                "    Default path for config is '" _DEFAULT_CONFIG_PATH_CLIENT "' or '" _DEFAULT_CONFIG_PATH_SERVER "'\n"
                "--rsa-gen-stdin [size]         Generate RSA private and public keys based on stdin data (64Kb max).\n"
                "    size: 2048 / 3072 (default) / 4096\n"
                "    Usage: `cat somefile | tcp-chat --gen-rsa-stdin`\n"
                "--rsa-gen-random [size]        Generate random RSA private and public keys.\n"
                "    size: 2048 / 3072 (default) / 4096\n"
                "--random-bytes [size]          Generate random bytes.\n"
                "    size: any number (default=32)\n"
                "--random-bytes-base64 [size]   Generate random bytes and print them in base64 encoding.\n"
                "    size: any number (default=32)\n"
            );
            Return 0;
        }
        if(arg_is("-l") || arg_is("--listen")){
            if(mode != ClientMode){
                printf("program mode is set already\n");
                Return 1;
            }
            mode = ServerMode;
        }
        else if(arg_is("--config")){
            if(++argi >= argc){
                printfe("ERROR: no config path specified\n");
                Return 1;
            }
            config_path = argv[argi];
        }
        else if(arg_is("--rsa-gen-stdin")){
            if(mode != ClientMode){
                printf("program mode is set already\n");
                Return 1;
            }

            mode = RsaGenStdin;
            if(++argi >= argc){
                size_arg = RSA_DEFAULT_KEY_SIZE;
            }
            else if(sscanf(argv[argi], "%u", &size_arg) != 1){
                printfe("ERROR: no key size specified\n");
            }
        }
        
        else if(arg_is("--rsa-gen-random")){
            if(mode != ClientMode){
                printf("program mode is set already\n");
                Return 1;
            }

            mode = RsaGenRandom;
            if(++argi >= argc){
                size_arg = RSA_DEFAULT_KEY_SIZE;
            }
            else if(sscanf(argv[argi], "%u", &size_arg) != 1){
                printfe("ERROR: no key size specified\n");
            }
        }
        else if(arg_is("--random-bytes")){
            if(mode != ClientMode){
                printf("program mode is set already\n");
                Return 1;
            }

            mode = RandomBytes;
            if(++argi >= argc){
                size_arg = 32;
            }
            else if(sscanf(argv[argi], "%u", &size_arg) != 1){
                printfe("ERROR: no size specified\n");
            }
        }
        else if(arg_is("--random-bytes-base64")){
            if(mode != ClientMode){
                printf("program mode is set already\n");
                Return 1;
            }

            mode = RandomBytesBase64;
            if(++argi >= argc){
                size_arg = 32;
            }
            else if(sscanf(argv[argi], "%u", &size_arg) != 1){
                printfe("ERROR: no size specified\n");
            }
        }
        else {
            printfe("ERROR: unknown argument '%s'\n"
                "Use '-h' to see list of avaliable arguments\n",
                argv[argi]);
            Return 1;
        }
    }
    
    switch(mode){
        case ClientMode: {
            if(!config_path)
                config_path = _DEFAULT_CONFIG_PATH_CLIENT;
            
            try_fatal(Client* client, p, Client_createFromConfig(config_path));
            Defer(Client_free(client));
            try_fatal_void(Client_run(client));
            break;
        }

        case ServerMode: {
            if(!config_path)
                config_path = _DEFAULT_CONFIG_PATH_SERVER;
            
            try_fatal(Server* server, p, Server_createFromConfig(config_path));
            Defer(Server_free(server));
            try_fatal_void(Server_run(server));
            break;
        }

        case RsaGenStdin: {
            printfe("reading stdin...\n");
            Array(u8) input_buf = Array_alloc_size(64*1024);
            Defer(free(input_buf.data));
            br_hmac_drbg_context rng = { .vtable = &br_hmac_drbg_vtable };
            br_hmac_drbg_init(&rng, &br_sha256_vtable, NULL, 0);
            i64 read_n = 0;
            do {
                read_n = fread(input_buf.data, 1, input_buf.size, stdin);
                if(read_n < 0){
                    printfe("ERROR: no input\n");
                    Return 1;
                }
                // put bytes to rng as seed
                br_hmac_drbg_update(&rng, input_buf.data, read_n);
            } while(read_n == input_buf.size);
            printfe("generating RSA key pair based on stdin...\n");
            br_rsa_private_key sk;
            br_rsa_public_key pk;
            try_fatal_void(RSA_generateKeyPair(size_arg, &sk, &pk, &rng.vtable));
            Defer(
                RSA_destroyPrivateKey(&sk);
                RSA_destroyPublicKey(&pk);
            );
            
            str sk_str = RSA_serializePrivateKey_base64(&sk);
            printf("rsa_private_key = %s\n", sk_str.data);
            free(sk_str.data);
            
            str pk_str = RSA_serializePublicKey_base64(&pk);
            printf("\nrsa_public_key = %s\n", pk_str.data);
            free(pk_str.data);
            break;
        }
        
        case RsaGenRandom: {
            printfe("generating random RSA key pair...\n");
            br_rsa_private_key sk;
            br_rsa_public_key pk;
            try_fatal_void(RSA_generateKeyPairFromSystemRandom(size_arg, &sk, &pk));
            Defer(
                RSA_destroyPrivateKey(&sk);
                RSA_destroyPublicKey(&pk);
            );

            str sk_str = RSA_serializePrivateKey_base64(&sk);
            printf("rsa_private_key = %s\n", sk_str.data);
            free(sk_str.data);
            
            str pk_str = RSA_serializePublicKey_base64(&pk);
            printf("\nrsa_public_key = %s\n", pk_str.data);
            free(pk_str.data);
            break;
        }
        
        case RandomBytes: {
            printfe("generating random bytes...\n");
            br_hmac_drbg_context rng = { .vtable = &br_hmac_drbg_vtable };
            rng_init_sha256_seedFromSystem(&rng.vtable);
            Array(u8) random_buf = Array_alloc_size(1024);
            u32 full_buffers_n = size_arg / random_buf.size;
            u32 remaining_n = size_arg % random_buf.size;
            while(full_buffers_n > 0){
                full_buffers_n--;
                br_hmac_drbg_generate(&rng, random_buf.data, random_buf.size);
                fwrite(random_buf.data, 1, random_buf.size, stdout);
            }
            
            br_hmac_drbg_generate(&rng, random_buf.data, remaining_n);
            fwrite(random_buf.data, 1, remaining_n, stdout);
            break;
        }

        case RandomBytesBase64: {
            printfe("generating random bytes...\n");
            br_hmac_drbg_context rng = { .vtable = &br_hmac_drbg_vtable };
            rng_init_sha256_seedFromSystem(&rng.vtable);
            Array(u8) random_buf = Array_alloc_size(1024);
            Array(u8) base64_buf = Array_alloc_size(base64_encodedSize(random_buf.size));
            u32 full_buffers_n = size_arg / random_buf.size;
            u32 remaining_n = size_arg % random_buf.size;
            u32 enc_size = 0;
            while(full_buffers_n > 0){
                full_buffers_n--;
                br_hmac_drbg_generate(&rng, random_buf.data, random_buf.size);
                enc_size = base64_encode(random_buf.data, random_buf.size, base64_buf.data);
                fwrite(base64_buf.data, 1, enc_size, stdout);
            }
            
            br_hmac_drbg_generate(&rng, random_buf.data, remaining_n);
            enc_size = base64_encode(random_buf.data, remaining_n, base64_buf.data);
            fwrite(base64_buf.data, 1, enc_size, stdout);
            break;
        }
        
        default:
            printfe("ERROR: invalid program mode %i\n", mode);
            Return 1;
    }

    Return 0;
}