#include <getopt.h>
#include <stdlib.h>
#include <stdio.h>

#include "ciphers.h"

#define MAX_FILE_BUFFER_SIZE    (1024 * 1024 * 500) // 500MiB file

int read_from_file(char* filepath, unsigned char* buff, unsigned int buff_len) {
    FILE* fp = NULL;
    if ((fp = fopen(filepath, "rb")) == NULL) return 0;
    fseek(fp, 0L, SEEK_END);
    int sz = ftell(fp);
    fseek(fp, 0L, SEEK_SET);
    fread(buff, sz, 1, fp);
    fclose(fp);
    return sz;
}

int write_to_file(char* filepath, unsigned char* buff, unsigned int buff_len) {
    FILE* fp = NULL;
    if ((fp = fopen(filepath, "wb")) == NULL) return 0;
    int ret = fwrite(buff, buff_len, 1, fp);
    fclose(fp);
    return ret;
}

void usage(char* prog_name) {
    printf("Usage: %s [-d | --dec] --cipher=cipher --in=inputfile --out=outputfile --passfile=password_file [--aadfile=aad_file] [--format=format] [--helpq]\n", prog_name);
}

int main (int argc, char** argv) {
    int                 c;
    char                enc                         = 1;
    char*               inputFilePath;
    char*               outputFilePath;
    char*               passFilePath;
    char*               aadFilePath;
    unsigned char*      aad                         = NULL;
    unsigned int        aad_len;
    unsigned char*      plaintext                   = NULL;
    int                 plaintext_len;
    unsigned char*      ciphertext                  = NULL;
    int                 ciphertext_len;
    unsigned char       password[512];
    int                 password_len;
    unsigned char       tag[TAG_LEN]                = {0};
    const EVP_CIPHER*   cipher                      = NULL;
    int                 AEAD                        = 0;
    timer               duration;
    unsigned long int   (*get_duration)(timer* t)   = &get_milliseconds;
    char*               format                      = "ms";
    int                 ret                         = EXIT_SUCCESS;

    while (1) {
        int option_index = 0;
        static struct option long_options[] = {
            {"cipher",      required_argument, 0, 'c' },
            {"in",          required_argument, 0, 'i' },
            {"out",         required_argument, 0, 'o' },
            {"dec",         no_argument,       0, 'd' },
            {"aadfile",     required_argument, 0, 'a' },
            {"passfile",    required_argument, 0, 'p' },
            {"format",      required_argument, 0, 'f' },
            {"help",        no_argument,       0, 'h' },
            {0,             0,                 0,  0  }
        };

        // enc/dec, infile, outfile, passfile, format, help
        if ((c = getopt_long(argc, argv, "dc:i:o:p:a:f:h", long_options, &option_index)) == -1)
            break;

        switch (c) {
            case 'h':
                usage(argv[0]);
                return ret;
            case 'c':
                if ((cipher = EVP_get_cipherbyname(optarg)) == NULL) {
                    printf("cipher selected %s is not supported\n", optarg);
                    return -1;
                }
                int opt_len = strlen(optarg);
                if (!memcmp(optarg + opt_len - 3, "gcm", 3))
                    AEAD = 1;
                break;
            case 'i':
                inputFilePath = optarg;
                break;
            case 'o':
                outputFilePath = optarg;
                break;
            case 'd':
                enc = 0;
                break;
            case 'p':
                passFilePath = optarg;
                break;
            case 'a':
                aadFilePath = optarg;
                break;
            case 'f':
                if (!memcmp(optarg, "milli", 5)) {
                    format = "ms";
                    get_duration = &get_milliseconds;
                } else if (!memcmp(optarg, "micro", 5)) {
                    format = "micros";
                    get_duration = &get_microseconds;
                } else if (!memcmp(optarg, "nano", 4)) {
                    format = "ns";
                    get_duration = &get_nanoseconds;
                }
                break;
            case '?':
                usage(argv[0]);
                return EXIT_FAILURE;
            default:
                break;
        }
    }

    if ((plaintext = malloc(MAX_FILE_BUFFER_SIZE * sizeof(unsigned char))) == NULL) {
        ret = EXIT_FAILURE;
        goto plaintext_alloc;
    }

    if ((ciphertext = malloc(MAX_FILE_BUFFER_SIZE * sizeof(unsigned char))) == NULL) {
        ret = EXIT_FAILURE;
        goto ciphertext_alloc;
    }

    if (AEAD) {
        if ((aad = malloc(MAX_FILE_BUFFER_SIZE * sizeof(unsigned char))) == NULL) {
            ret = EXIT_FAILURE;
            goto aad_alloc;
        }
        if (!(aad_len = read_from_file(aadFilePath, aad, MAX_FILE_BUFFER_SIZE))) {
            fprintf(stderr, "Unable to open aad file\n");
            ret = EXIT_FAILURE;
            goto error;
        }
    }

    if (enc) {
        if (!(plaintext_len = read_from_file(inputFilePath, plaintext, MAX_FILE_BUFFER_SIZE))) {
            fprintf(stderr, "Unable to open plaintext file\n");
            ret = EXIT_FAILURE;
            goto error;
        }

        if (!(password_len = read_from_file(passFilePath, password, 512))) {
            fprintf(stderr, "Unable to open password file\n");
            ret = EXIT_FAILURE;
            goto error;
        }

        if (AEAD) {
            if ((ciphertext_len = gcm_encrypt(cipher, plaintext, plaintext_len,
                        password, password_len, aad, aad_len, ciphertext, tag, &duration)) <= 0) {
                            fprintf(stderr, "wrong encrypting\n");
                            ret = EXIT_FAILURE;
                            goto error;
                        }
            //print_hexbuf("C", ciphertext, ciphertext_len);
            memcpy(ciphertext + ciphertext_len, tag, TAG_LEN); // copy tag to ciphertext end
            ciphertext_len += TAG_LEN;
        } else
            if ((ciphertext_len = encrypt(cipher, plaintext, plaintext_len, password, password_len, ciphertext, &duration)) <= 0) {
                            fprintf(stderr, "wrong encrypting\n");
                            ret = EXIT_FAILURE;
                            goto error;
                        }

        printf("encryption took %lu %s\n", get_duration(&duration), format);

        if (!write_to_file(outputFilePath, ciphertext, ciphertext_len)) {
            fprintf(stderr, "Unable to open output file\n");
            ret = EXIT_FAILURE;
            goto error;
        }
    } else {
        if (!(ciphertext_len = read_from_file(inputFilePath, ciphertext, MAX_FILE_BUFFER_SIZE))) {
            fprintf(stderr, "Unable to open ciphertext file\n");
            ret = EXIT_FAILURE;
            goto error;
        }

        if (!(password_len = read_from_file(passFilePath, password, 512))) {
            fprintf(stderr, "Unable to open password file\n");
            ret = EXIT_FAILURE;
            goto error;
        }

        if (AEAD) {
            ciphertext_len -= TAG_LEN;
            memcpy(tag, ciphertext + ciphertext_len, TAG_LEN); // copy ciphertext end to tag
            // check authenticated decryption
            if ((plaintext_len = gcm_decrypt(cipher, ciphertext, ciphertext_len,
                            password, password_len, aad, aad_len, tag, plaintext, &duration)) <= 0) {
                fprintf(stderr, "wrong decrypting\n");
                ret = EXIT_FAILURE;
                goto error;
            }
        } else
            plaintext_len = decrypt(cipher, ciphertext, ciphertext_len, password, password_len, plaintext, &duration);

        printf("decryption took %lu %s\n", get_duration(&duration), format);

        if (!write_to_file(outputFilePath, plaintext, plaintext_len)) {
            fprintf(stderr, "Unable to open output file\n");
            ret = EXIT_FAILURE;
            goto error;
        }
    }

error:
    free(ciphertext);
ciphertext_alloc:
    free(plaintext);
plaintext_alloc:
    if (aad != NULL) free(aad);
aad_alloc:
    return ret;
}