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

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

typedef enum ProgramMode {
    Client,
    Server,
    RsaGenStdin,
    RsaGenRandom,
} ProgramMode;

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

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

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

    ProgramMode mode = Client;
    cstr server_endpoint_cstr = NULL;
    cstr config_path = NULL;
    u32 key_size = 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 [addr:port]       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"
            );
            Return 0;
        }
        if(arg_is("-l") || arg_is("--listen")){
            if(mode != Client){
                printf("program mode is set already\n");
                Return 1;
            }

            mode = Server;
            if(++argi >= argc){
                printfe("ERROR: no endpoint specified\n");
                Return 1;
            }
            server_endpoint_cstr = argv[argi];
        }
        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 != Client){
                printf("program mode is set already\n");
                Return 1;
            }

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

            mode = RsaGenRandom;
            if(++argi >= argc){
                key_size = RSA_DEFAULT_KEY_SIZE;
            }
            else if(sscanf(argv[argi], "%u", &key_size) != 1){
                printfe("ERROR: no key size specified\n");
            }
        }
        else {
            printfe("ERROR: unknown argument '%s'\n"
                "Use '-h' to see list of avaliable arguments\n",
                argv[argi]);
            Return 1;
        }
    }

    try_fatal_void(network_init());
    Defer(network_deinit());
    
    switch(mode){
        case Client: {
            if(config_path == NULL)
                config_path = _DEFAULT_CONFIG_PATH_CLIENT;
            try_fatal_void(client_run());
            break;
        }

        case Server: {
            if(config_path == NULL)
                config_path = _DEFAULT_CONFIG_PATH_SERVER;
            try_fatal_void(server_run(server_endpoint_cstr, config_path));
            break;
        }

        case RsaGenStdin: {
                size_t input_max_size = 64*1024;
                char* input_buf = malloc(input_max_size);
                Defer(free(input_buf));
                size_t read_n = fread(input_buf, 1, input_max_size, stdin);
                if(read_n == 0){
                    printfe("ERROR: no input\n");
                    Return 1;
                }
                str input_str = str_construct(input_buf, read_n, false);
                printfe("generating RSA key pair based on stdin...\n");
                br_rsa_private_key sk;
                br_rsa_public_key pk;
                try_fatal_void(RSA_generateKeyPairFromPassword(key_size, &sk, &pk, input_str));
                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(key_size, &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;
        
        default:
            printfe("ERROR: invalid program mode %i\n", mode);
            Return 1;
    }

    Return 0;
}