#ifndef _AES
#define _AES

#include <iostream>
#include <stdio.h>
#include <string>
#include <sstream>
#include <string.h>
#include <algorithm>
#include <typeinfo>
#include <regex>

#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/aes.h>
#include <openssl/err.h>
#include <openssl/rand.h>

#define FAILURE -1
#define SUCCESS 0

#define AES_KEYLEN 256
#define AES_ROUNDS 6

class aes {
protected:
    unsigned char* aesKey;
    unsigned char* aesIV;
    aes& clone(const aes& _e);
    void init_all();
    void clear_all();
    void setExtension(std::string _ex);
private:
    std::string filename;
    std::string generateNewFilename();
    std::string extension = ".new";
    
public:
    aes() {};
    aes(std::string _filename);
    std::string writeFile(unsigned char *file, size_t fileLength);
    int readFile(unsigned char **file);

    std::string PrintAesKey();
    void setAesIV(unsigned char* _aesIV);
    void setAesKey(unsigned char* _aesKey);
    unsigned char* getAesKey();
    unsigned char* getAesIV();

    virtual ~aes() = default;
};

aes::aes(std::string _filename) {
    filename = _filename;
    aesKey = (unsigned char*)malloc(AES_KEYLEN/8);
    aesIV = (unsigned char*)malloc(AES_KEYLEN/8);
}

aes& aes::clone(const aes& _e) {
    filename = _e.filename;
    
    setAesKey(_e.aesKey);
    setAesIV(_e.aesIV);

    return *this;
}

void aes::init_all() {

    unsigned char *aesPass = (unsigned char*)malloc(AES_KEYLEN/8);
    unsigned char *aesSalt = (unsigned char*)malloc(8);

    if(aesKey == NULL || aesIV == NULL || aesPass == NULL || aesSalt == NULL) {
        exit(FAILURE);
    }

    // #define USE_PBKDF
    #ifdef USE_PBKDF
        std::cerr << "utilisation de USE_PBKDF" << std::endl;
        // Get some random data to use as the AES pass and salt
        if(RAND_bytes(aesPass, AES_KEYLEN/8) == 0) {
            exit(FAILURE);
        }

        if(RAND_bytes(aesSalt, 8) == 0) {
            exit(FAILURE);
        }

        if(EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha256(), aesSalt, aesPass, AES_KEYLEN/8, AES_ROUNDS, aesKey, aesIV) == 0) {
            exit(FAILURE);
        }
    #else
        if(RAND_bytes(aesKey, AES_KEYLEN/8) == 0) {
            exit(FAILURE);
        }

        if(RAND_bytes(aesIV, AES_KEYLEN/8) == 0) {
            exit(FAILURE);
        }
    #endif

    free(aesPass);
    free(aesSalt);
}

// peut être déporté le buffer lut avec methode pour travaillé dessus
int aes::readFile(unsigned char** file) {
    FILE *fd = fopen(filename.c_str(), "rb");
    if(fd == NULL) {
        fprintf(stderr, "Failed to open file: %s\n", strerror(errno));
        exit(1);
    }

    // Determine size of the file
    fseek(fd, 0, SEEK_END);
    size_t fileLength = ftell(fd);
    fseek(fd, 0, SEEK_SET);

    // Allocate space for the file
    *file = (unsigned char*)malloc(fileLength);
    if(*file == NULL) {
        fprintf(stderr, "Failed to allocate memory\n");
        exit(1);
    }
    
    // Read the file into the buffer
    size_t bytesRead = fread(*file, 1, fileLength, fd);

    if(bytesRead != fileLength) {
        fprintf(stderr, "Error reading file\n");
        exit(1);
    }

    fclose(fd);

    return fileLength;
}

std::string aes::generateNewFilename() {
    std::regex e("\\b(.*[(\\.txt)]*)\\.[^(txt)].*");
    filename = std::regex_replace(filename, e,"$1");
    return filename+extension;
}

std::string aes::writeFile(unsigned char *file, size_t fileLength) {    
    std::string newFile = generateNewFilename();
    FILE *fd = fopen(newFile.c_str(), "wb");
    if(fd == NULL) {
        fprintf(stderr, "Failed to open file: %s\n", strerror(errno));
        exit(1);
    }

    size_t bytesWritten = fwrite(file, 1, fileLength, fd);

    if(bytesWritten != fileLength) {
        fprintf(stderr, "Failed to write file\n");
        exit(1);
    }

    fclose(fd);
    return newFile;
}

void aes::setExtension(std::string _ex) {
    extension = _ex;
}

void aes::setAesKey(unsigned char* _aesKey) {
    if(aesKey == NULL)
        aesKey = (unsigned char*)malloc(AES_KEYLEN/8);
    memcpy(aesKey , _aesKey, AES_KEYLEN/8);
}

void aes::setAesIV(unsigned char* _aesIV) {
    if(aesIV == NULL)
        aesIV = (unsigned char*)malloc(AES_KEYLEN/8);    
    memcpy(aesIV , _aesIV, AES_KEYLEN/8);
}


unsigned char* aes::getAesKey() {
    unsigned char* res = (unsigned char*)malloc(AES_KEYLEN/8);    
    memcpy(res, aesKey, AES_KEYLEN/8);
    return res;
}

unsigned char* aes::getAesIV() {
    unsigned char* res = (unsigned char*)malloc(AES_KEYLEN/8);    
    memcpy(res, aesIV, AES_KEYLEN/8);
    return res;
}

std::string aes::PrintAesKey() {
    std::stringstream ss;
    std::string res;
    
    res = "aesKey : ";
    for (int i = 0; i < 32; i++) {
        ss << std::hex;
        ss >> res;
        res += std::string(1, (char)aesKey[i]+'\n');
    }
    return res;
}

void aes::clear_all() {
    free(aesIV);
    free(aesKey);
}

#endif