// Modifiy from source found here: https://gist.github.com/leto/1a984b1daa710ed55464c093640d6769

#include <cstdint>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <string>
#include <termios.h>
#include <unistd.h>

#include <sodium.h>

#include "StreamUtils.h"

const std::string BOLD{"\033[1m"};
const std::string ENDC{"\033[0m"};
const std::string RED{"\033[31m"};
const std::string YELLOW{"\033[33m"};
const std::string PURPLE{"\033[35m"};
const std::string LBLUE{"\033[1;34m"};
const std::string UP_ONE = "\033[1A";
const uint32_t CHUNK_SIZE{4096};

#define ERROR_MSG       BOLD << RED << "ERROR: " << ENDC


void get_password(char* password)
{
    static struct termios old_terminal;
    static struct termios new_terminal;

    // get settings of the actual terminal
    tcgetattr(STDIN_FILENO, &old_terminal);

    // do not echo the characters
    new_terminal = old_terminal;
    new_terminal.c_lflag &= ~(ECHO);

    // set this as the new terminal options
    tcsetattr(STDIN_FILENO, TCSANOW, &new_terminal);

    // get the password
    // the user can add chars and delete if he puts it wrong
    // the input process is done when he hits the enter
    // the \n is stored, we replace it with \0
    if (fgets(password, BUFSIZ, stdin) == NULL)
    {
        password[0] = '\0';
    }
    else
    {
        password[strlen(password) - 1] = '\0';
    }

    // go back to the old settings
    tcsetattr(STDIN_FILENO, TCSANOW, &old_terminal);
}

static int encrypt(
    const char *target_file,
    const char *source_file,
    const uint8_t key[crypto_secretstream_xchacha20poly1305_KEYBYTES],
    const uint8_t salt[crypto_pwhash_SALTBYTES])
{
    // NOTE: the salt value is not used directly in encrpytion but is written
    // to the first XYZ bytes of the output file

    uint8_t buf_in[CHUNK_SIZE]{};
    uint8_t buf_out[CHUNK_SIZE + crypto_secretstream_xchacha20poly1305_ABYTES]{};
    uint8_t header[crypto_secretstream_xchacha20poly1305_HEADERBYTES]{};
    crypto_secretstream_xchacha20poly1305_state st;

    FILE* fp_s = std::fopen(source_file, "rb");
    FILE* fp_t = std::fopen(target_file, "wb");

    // First write out salt bytes to the output file
    std::fwrite(salt, 1, crypto_pwhash_SALTBYTES, fp_t);

    // Next write out header
    crypto_secretstream_xchacha20poly1305_init_push(&st, header, key);
    std::fwrite(header, 1, sizeof(header), fp_t);

    // Iterate over the input file a chunk at a time, encrypt each, and then
    // write it to the output file
    size_t rlen;
    long long unsigned out_len;
    int eof;
    uint8_t tag;
    do
    {
        rlen = std::fread(buf_in, 1, sizeof(buf_in), fp_s);
        eof = std::feof(fp_s);
        tag = eof ? crypto_secretstream_xchacha20poly1305_TAG_FINAL : 0;
        crypto_secretstream_xchacha20poly1305_push(&st, buf_out, &out_len, buf_in,
                                                   rlen,nullptr, 0, tag);
        std::fwrite(buf_out, 1, (size_t)out_len, fp_t);
    }
    while (! eof);

    std::fclose(fp_t);
    std::fclose(fp_s);

    return 0;
}

static int decrypt(
    const char *target_file,
    const char *source_file,
    const uint8_t key[crypto_secretstream_xchacha20poly1305_KEYBYTES],
    uint32_t saltLength)
{
    // NOTE: the "saltLength" parameter is used to skip past the salt value at
    // the beginning of the file

    int ret = -1;
    uint8_t buf_in[CHUNK_SIZE + crypto_secretstream_xchacha20poly1305_ABYTES]{};
    uint8_t buf_out[CHUNK_SIZE]{};
    uint8_t header[crypto_secretstream_xchacha20poly1305_HEADERBYTES]{};
    crypto_secretstream_xchacha20poly1305_state st;

    FILE* fp_s = std::fopen(source_file, "rb");
    FILE* fp_t = std::fopen(target_file, "wb");

    // Skip past the salt value at the beginning of the file
    std::fseek(fp_s, saltLength, SEEK_SET);

    std::fread(header, 1, sizeof header, fp_s);

    if (crypto_secretstream_xchacha20poly1305_init_pull(&st, header, key) != 0)
    {
        std::cerr << ERROR_MSG << "during decrypt: incomplete header" << std::endl;
        goto ret;
    }

    size_t rlen;
    long long unsigned out_len;
    int eof;
    uint8_t tag;
    do
    {
        rlen = fread(buf_in, 1, sizeof buf_in, fp_s);
        eof = feof(fp_s);

        int rc = crypto_secretstream_xchacha20poly1305_pull(&st, buf_out,
                                                            &out_len, &tag,
                                                            buf_in, rlen,
                                                            nullptr, 0);
        if (rc != 0)
        {
            std::cerr << ERROR_MSG << "during decrypt: corrupted chunk" << std::endl;
            goto ret;
        }

        if (tag == crypto_secretstream_xchacha20poly1305_TAG_FINAL && ! eof)
        {
            std::cerr << ERROR_MSG << "during decrypt: premature end of file" << std::endl;
            goto ret;
        }
        std::fwrite(buf_out, 1, (size_t) out_len, fp_t);
    }
    while (! eof);

    ret = 0;

ret:
    std::fclose(fp_t);
    std::fclose(fp_s);
    return ret;
}

int main(int argc, char** argv)
{
    const uint32_t saltLen{crypto_pwhash_SALTBYTES};
    const uint32_t keyLen{crypto_secretstream_xchacha20poly1305_KEYBYTES};
    char passwordBuffer[BUFSIZ]{};

    if (sodium_init() != 0)
    {
        return 1;
    }

    uint8_t key[keyLen]{};
#if 1
    uint8_t salt[saltLen]{};
#else
    uint8_t salt[saltLen]{0x1A, 0x19, 0x2F, 0x62, 0x65, 0x9A,
                          0x9C, 0x64, 0x99, 0xA7, 0xD1, 0x3B,
                          0x02, 0xC4, 0xA1, 0xC9};
#endif

    // DEBUGGERY
//    std::cout << "SALT: " << HexBytesInlineBE(salt, saltLen) << std::endl;
//    std::cout << "KEY: " << HexBytesInlineBE(key, keyLen) << std::endl;

    if (argc < 2)
    {
        std::cerr << ERROR_MSG << "please provide correct number of arguments!"
                  << std::endl;
        return -1;
    }


    if (std::strlen(argv[1]) == 3 && std::strncmp(argv[1], "enc", 3) == 0)
    {
        // Encrypt the file

        // TODO: confirm input file exists

        // Generate random salt
        randombytes_buf(salt, saltLen);

        std::cout << "Password: ";
        get_password(passwordBuffer);
        std::cout << std::endl;

        // Generate key from password and salt
        if (crypto_pwhash(key, sizeof key, passwordBuffer,
                          std::strlen(passwordBuffer), salt,
                          crypto_pwhash_OPSLIMIT_INTERACTIVE,
                          crypto_pwhash_MEMLIMIT_INTERACTIVE,
                          crypto_pwhash_ALG_ARGON2ID13) != 0)
        {
            std::cerr << ERROR_MSG << "crypto_pwhash - out of memory" << std::endl;
            return -1;
        }

        // Explicitly zeroize the password buffer
        explicit_bzero(passwordBuffer, BUFSIZ);

        // Create a name for the cipher text file by appending ".enc" to the
        // name of the plain text file
        std::string plainTextFile(argv[2]);
        std::string cipherTextFile(plainTextFile + ".enc");

        std::cout << "PT: " << plainTextFile << std::endl;
        std::cout << "CT: " << cipherTextFile << std::endl;

        if (encrypt(cipherTextFile.c_str(),
                    plainTextFile.c_str(),
                    key, salt) != 0)
        {
            return 1;
        }
    }
    else if (std::strlen(argv[1]) == 3 && std::strncmp(argv[1], "dec", 3) == 0)
    {
        // Decrypt the file

        // TODO: confirm input file exists

        // Create a name for the cipher text file by appending ".enc" to the
        // name of the plain text file
        std::string cipherTextFile(argv[2]);
        std::string plainTextFile(cipherTextFile + ".dec");

        std::cout << "CT: " << cipherTextFile << std::endl;
        std::cout << "PT: " << plainTextFile << std::endl;

        std::cout << "Password: ";
        get_password(passwordBuffer);
        std::cout << std::endl;

        // Read in salt value
        FILE* ctFile = std::fopen(cipherTextFile.c_str(), "rb");
        std::fread(salt, 1, saltLen, ctFile);
        std::fclose(ctFile);

        // Generate key from password and salt
        if (crypto_pwhash(key, sizeof key, passwordBuffer,
                          std::strlen(passwordBuffer), salt,
                          crypto_pwhash_OPSLIMIT_INTERACTIVE,
                          crypto_pwhash_MEMLIMIT_INTERACTIVE,
                          crypto_pwhash_ALG_ARGON2ID13) != 0)
        {
            std::cerr << ERROR_MSG << "crypto_pwhash - out of memory" << std::endl;
            return -1;
        }

        // Explicitly zeroize the password buffer
        explicit_bzero(passwordBuffer, BUFSIZ);

        if (decrypt(plainTextFile.c_str(),
                    cipherTextFile.c_str(),
                    key, saltLen) != 0)
        {
            return 1;
        }
    }
    else
    {
        std::cerr << ERROR_MSG << "invalid mode \"" << argv[1] << "\"!" << std::endl;
        return 2;
    }

    return 0;
}