#include <iostream>
#include <sstream>
#include <limits>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <sqlite3.h>
#include <gcrypt.h>

using namespace std;


string int_to_str(int num) {
    stringstream ss;

    ss << num;

    return ss.str();
}

//                                          ///
//          ENCRYPTION / DECRYPTION         ///
//                                          ///

void hash(string input, string* output){
    size_t txtLength = input.length();
    char * hash = (char*)malloc(sizeof(char) * 33);
    char * textBuffer = (char*)malloc(sizeof(char) * (txtLength+1));
    memset(hash, 0, 33);

    gcry_md_hash_buffer(
        GCRY_MD_SHA256, // gcry_cipher_hd_t
        hash,    // void *
        input.c_str(),    // const void *
        txtLength);   // size_t

    *output = hash;
    free(hash);
    free(textBuffer);
}

void salsa(string message, string passphrase, bool encrypt, string *output){

    #define GCRY_CIPHER GCRY_CIPHER_AES128   // Pick the cipher here

    gcry_error_t     gcryError;
    gcry_cipher_hd_t gcryCipherHd;
    size_t           index;

    size_t keyLength = gcry_cipher_get_algo_keylen(GCRY_CIPHER);
    size_t blkLength = gcry_cipher_get_algo_blklen(GCRY_CIPHER);
    char * txtBuffer = (char *)message.c_str();
    size_t txtLength = message.length(); // string plus termination
    char * encBuffer = (char *)malloc(txtLength);
    char * outBuffer = (char *)malloc(txtLength);
    char * aesSymKey = (char *)passphrase.c_str(); // 16 bytes
    string iv = "SalsaManSalsaMan";
    char * iniVector = (char *)iv.c_str();

    gcry_cipher_open(
        &gcryCipherHd, // gcry_cipher_hd_t *
        GCRY_CIPHER_SALSA20,   // int
        GCRY_CIPHER_MODE_STREAM,   // int
        0);            // unsigned int


    gcry_cipher_setkey(gcryCipherHd, aesSymKey, keyLength);

    gcry_cipher_setiv(gcryCipherHd, iniVector, 8);

    if (encrypt){
        gcry_cipher_encrypt(
            gcryCipherHd, // gcry_cipher_hd_t
            encBuffer,    // void *
            txtLength,    // size_t
            txtBuffer,    // const void *
            txtLength);   // size_t
    }


    gcry_cipher_setiv(gcryCipherHd, iniVector, 8);


    if (!encrypt){
        gcry_cipher_decrypt(
            gcryCipherHd, // gcry_cipher_hd_t
            outBuffer,    // void *
            txtLength,    // size_t
            txtBuffer,    // const void *
            txtLength);   // size_t
    }


    if (encrypt){
        *output = encBuffer;
    }
    else {
        *output = outBuffer;
    }

    // clean up after ourselves
    gcry_cipher_close(gcryCipherHd);
    free(encBuffer);
    free(outBuffer);
}



//                                          ///
//              SEND MESSAGES               ///
//                                          ///

static int messageSendSuccessCallback(void *data, int argc, char **argv, char **azColName){
    bool *messageSendSuccess = reinterpret_cast<bool*>(data);

    *messageSendSuccess = ((argv[0][0] - '0') > 0);

    return 0;
}

void sendMessage(string sender, string recipient, string message, string passphrase, sqlite3 *db){
    char *zErrMsg = 0;
    int rc;
    string sql;
    const char* data = "";
    bool messageSendSuccess = false;

    string passphraseHash;
    hash(passphrase, &passphraseHash);
    string messageEncryption;
    salsa(message, passphraseHash, true, &messageEncryption);

    sql = "SELECT COUNT(*) FROM USERS WHERE USERNAME=\"" + recipient + "\";";
    rc = sqlite3_exec(db, sql.c_str(), messageSendSuccessCallback, &messageSendSuccess, &zErrMsg);

    if( rc != SQLITE_OK ) {
        sqlite3_free(zErrMsg);
    }

    if (messageSendSuccess){
        sql = "INSERT INTO MESSAGES(SENDER, RECIPIENT, MESSAGE, PASSPHRASE) VALUES(\"" + sender + "\", \"" + recipient + "\", \"" + messageEncryption + "\", \"" + passphraseHash + "\");";

        rc = sqlite3_exec(db, sql.c_str(), NULL, (void*)data, &zErrMsg);

        if( rc != SQLITE_OK ) {
            sqlite3_free(zErrMsg);
        }

        cout << "Message sent successfully" << endl << endl;
    }
    else {
        cout << "Could not send message, inputted recipient does not exist" << endl << endl;
    }
}


//                                          ///
//              READ MESSAGES               ///
//                                          ///

static int messageDisplayCallback(void *data, int argc, char **argv, char **azColName){
    string message = argv[2];
    string passphrase = argv[3];

    string answer;
    salsa(message, passphrase, false, &answer);

    cout << "  _____________________________________________________________" << endl;
    printf("  |\n  | Id: %s | From : %s\n  |____________________________________________________________\n  |\n  | Message: %s\n", argv[0], argv[1], answer.c_str());
    cout << "  |____________________________________________________________";
    printf("\n");
    return 0;
}

static int messageDisplaySuccessCallback(void *data, int argc, char **argv, char **azColName){
    bool *messageDisplaySuccess = reinterpret_cast<bool*>(data);

    *messageDisplaySuccess = ((argv[0][0] - '0') > 0);

    return 0;
}

void displayMessage(int id, string passphrase, sqlite3 *db){
    char *zErrMsg = 0;
    int rc;
    string sql;
    const char* data = "";
    bool messageDisplaySuccess = false;

    string passphraseHash;
    hash(passphrase, &passphraseHash);

    sql = "SELECT COUNT(*) FROM MESSAGES WHERE ID=" + int_to_str(id) + " AND PASSPHRASE=\"" + passphraseHash + "\";";
    rc = sqlite3_exec(db, sql.c_str(), messageDisplaySuccessCallback, &messageDisplaySuccess, &zErrMsg);

    if( rc != SQLITE_OK ) {
        sqlite3_free(zErrMsg);
    }

    if (messageDisplaySuccess){
        sql = "SELECT ID, SENDER, MESSAGE, PASSPHRASE FROM MESSAGES WHERE ID=" + int_to_str(id) + " AND PASSPHRASE=\"" + passphraseHash + "\";";

        rc = sqlite3_exec(db, sql.c_str(), messageDisplayCallback, (void*)data, &zErrMsg);

        if( rc != SQLITE_OK ) {
            sqlite3_free(zErrMsg);
        }

    }
    else {
        cout << "Could not display message with inputted id and passphrase" << endl;
    }
}

//                                          ///
//          DISPLAY MESSAGE LIST            ///
//                                          ///

static int messageListCallback(void *data, int argc, char **argv, char **azColName){
    printf("    Id: %s | From : %s", argv[0], argv[1]);

    printf("\n");
    return 0;
}

void displayMessageList(string username, sqlite3 *db){
    char *zErrMsg = 0;
    int rc;
    string sql;
    const char* data = "";

    sql = "SELECT ID, SENDER FROM MESSAGES WHERE RECIPIENT=\"" + username + "\" ;";

    rc = sqlite3_exec(db, sql.c_str(), messageListCallback, (void*)data, &zErrMsg);

    if( rc != SQLITE_OK ) {
        sqlite3_free(zErrMsg);
    }
}

//                                          ///
//              COUNT MESSAGES              ///
//                                          ///

static int messageCountCallback(void *data, int argc, char **argv, char **azColName){
    cout << (argv[0][0]);

    return 0;
}


void countMessages(string username, sqlite3 *db){
    cout << "You have ";

    char *zErrMsg = 0;
    int rc;
    string sql;
    const char* data = "";

    sql = "SELECT COUNT(*) FROM MESSAGES WHERE RECIPIENT=\"" + username + "\";";

    rc = sqlite3_exec(db, sql.c_str(), messageCountCallback, (void*)data, &zErrMsg);

    if( rc != SQLITE_OK ) {
        sqlite3_free(zErrMsg);
    }

    cout << " messages:";
}

//                                          ///
//                  LOGIN                   ///
//                                          ///

static int loginCallback(void *data, int argc, char **argv, char **azColName){
    bool *loginSuccess = reinterpret_cast<bool*>(data);

    *loginSuccess = ((argv[0][0] - '0') > 0);

    return 0;
}

bool loginAccount(string username, string password,sqlite3 *db){
    cout << "\nLogging in..." << endl;

    char *zErrMsg = 0;
    int rc;
    string sql;
    bool loginSuccess = false;

    string hashValue;
    hash(password, &hashValue);

    sql = "SELECT COUNT(*) FROM USERS WHERE USERNAME=\"" + username + "\" AND PASSWORD=\"" + hashValue + "\";";

    rc = sqlite3_exec(db, sql.c_str(), loginCallback, &loginSuccess, &zErrMsg);

    if( rc != SQLITE_OK ) {
        sqlite3_free(zErrMsg);
    }


    if (loginSuccess){
        return true;
    }
    else {
        return false;
    }
}

//                                          ///
//              REGISTER ACCOUNT            ///
//                                          ///

static int callback(void *data, int argc, char **argv, char **azColName){
    int i;

    for(i = 0; i<argc; i++){
        printf("%s = %s\n", azColName[i], argv[i] ? argv[i] : "NULL");
    }

    printf("\n");
    return 0;
}

bool registerAccount(string username, string password,sqlite3 *db){
    char *zErrMsg = 0;
    int rc;
    string sql;

    string hashValue;
    hash(password, &hashValue);

    sql = "INSERT INTO USERS(USERNAME, PASSWORD) VALUES(\"" + username + "\", \"" + hashValue + "\");";

    rc = sqlite3_exec(db, sql.c_str(), NULL, NULL, &zErrMsg);


    if( rc != SQLITE_OK ) {
        sqlite3_free(zErrMsg);
        return false;
    } else {
        return true;
    }

    free(zErrMsg);
}

//                                          ///
//                  MAIN                    ///
//                                          ///

int main(){
    sqlite3 *db;
    char *zErrMsg = 0;
    int rc;
    string sql;

    /* Open database */
    rc = sqlite3_open("mail.db", &db);

    if( rc ) {
        fprintf(stderr, "Can't open database: %s\n", sqlite3_errmsg(db));
        return(0);
    }


    // /* Create SQL statement */
    // sql = "SELECT USERNAME from USERS";

    // /* Execute SQL statement */
    // rc = sqlite3_exec(db, sql.c_str(), callback, NULL, &zErrMsg);

    // if( rc != SQLITE_OK ) {
    //     fprintf(stderr, "SQL error: %s\n", zErrMsg);
    //     sqlite3_free(zErrMsg);
    // }

    cout << endl;

    int homeChoice = -1;
    int registerLoginChoice = -1;
    int mainChoice = -1;
    int messageChoice = -1;
    string passphrase = "";
    string username = "";
    string password = "";
    string recipient = "";
    string message = "";
    string sharedPassphrase = "";
    bool loginAccountSuccess = false;
    bool registerAccountSuccess = false;
    bool readMessageSuccess = false;

    cout << "Welcome to Gee-Mail!" << endl << endl;
    cout << "Enter (0) to sign-in or (1) to register: ";
    cin >> homeChoice;

    if (!cin){
		cin.clear();
		cin.ignore(numeric_limits<streamsize>::max(), '\n');
	}

    while (!(homeChoice == 0 || homeChoice ==  1)){
        cout << "Enter (0) to sign-in or (1) to register: ";
        cin >> homeChoice;
        if (!cin){
			cin.clear();
			cin.ignore(numeric_limits<streamsize>::max(), '\n');
		}
    }

    if (homeChoice == 0) {
        cout << "\nSign into your account:\n\n";

        while (!loginAccountSuccess){
            cout << "Enter your username: ";
            cin >> username;
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }
            cout << "Enter your password: ";
            cin >> password;
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }

            loginAccountSuccess = loginAccount(username, password, db);

            if (!loginAccountSuccess){
                cout << "\nCannot login - username and password combination not recognized" << endl;
                cout << "Please try again with a different username and password combination" << endl << endl;
            }
        }
    }
    else if (homeChoice == 1){
        cout << "\nCreate your account:\n\n";

        while (!registerAccountSuccess){
            cout << "Enter your username: ";
            cin >> username;
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }
            cout << "Enter your password: ";
            cin >> password;
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }

            registerAccountSuccess = registerAccount(username, password, db);

            if (!registerAccountSuccess){
                cout << "\nCannot create account with username: " << username << " - username is already taken!" << endl;
                cout << "Please try again with a different account name:" << endl << endl;
            }
        }

        cout << endl << "Would you like to sign in to your new account? (0) = no, (1) = yes : ";

        cin >> registerLoginChoice;
        if (!cin){
		    cin.clear();
    	    cin.ignore(numeric_limits<streamsize>::max(), '\n');
	    }

        while (!(registerLoginChoice == 0 || registerLoginChoice ==  1)){
            cout << "Enter (0) to exit or (1) to sign in: ";
            cin >> registerLoginChoice;
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }
        }

        if (registerLoginChoice == 1){
            loginAccount(username, password, db);
        }
        else {
            cout << endl << "Goodbye!" << endl << endl;
            return 0;
        }
    }

    cout << endl << "Succesfully logged in" << endl << endl;

    while(mainChoice != 2) {
        mainChoice = -1;
        messageChoice = -1;
        passphrase = "";
        recipient = "";
        message = "";
        sharedPassphrase = "";

        countMessages(username, db);

        cout << endl;

        displayMessageList(username, db);

        cout << endl;

        cout << "Enter (0) to read a message, (1) to send a message, or (2) to quit: ";
        cin >> mainChoice;
        if (!cin){
			cin.clear();
			cin.ignore(numeric_limits<streamsize>::max(), '\n');
		}

        while (!(mainChoice == 0 || mainChoice ==  1 || mainChoice == 2)){
            cout << endl << "Enter (0) to read a message, (1) to send a message, or (2) to quit: ";
            cin >> mainChoice;
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }
        }

        cout << endl;

        if (mainChoice == 0){
            cout << "Enter the Id of the message you'd like to read: ";
            cin >> messageChoice;
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }
            cout << "Enter the shared passphrase with the sender: ";
            cin >> passphrase;
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }
            cout << endl;

            displayMessage(messageChoice, passphrase, db);

            cout << endl;
        }
        else if (mainChoice == 1){
            string test;
            getline(cin, test);

            cout << "Enter the username of your desired recipient: ";
            getline(cin, recipient);
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }

            cout <<  "Enter your message: ";
            getline(cin, message);
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }

            cout << "Enter your shared passphrase: ";
            getline(cin, sharedPassphrase);
            if (!cin){
			    cin.clear();
			    cin.ignore(numeric_limits<streamsize>::max(), '\n');
		    }

            cout << endl;

            sendMessage(username, recipient, message, sharedPassphrase, db);

        }
    }

    cout << "Goodbye!" << endl << endl;

    sqlite3_close(db);

    return 0;
};