#include "WedgeSecurity.h"
#include <openssl/evp.h>
#include <openssl/rand.h>
#include <QByteArray>

static const int ITERATIONS = 9999;
static const int KEY_LENGTH = 32;
static const int SALT_LENGTH = 16;
static const int IV_LENGTH = 12;
static const QString PREFIX = "wgd:";

bool WedgeSecurity::isEncrypted(const QString& text) {
    return text.startsWith(PREFIX);
}

QString WedgeSecurity::encrypt(const QString& plainText, const QString& passPhrase) {
    QByteArray salt(SALT_LENGTH, 0);
    RAND_bytes((unsigned char*)salt.data(), SALT_LENGTH);

    QByteArray iv(IV_LENGTH, 0);
    RAND_bytes((unsigned char*)iv.data(), IV_LENGTH);

    QByteArray key(KEY_LENGTH, 0);
    QByteArray phrase = passPhrase.toUtf8();
    PKCS5_PBKDF2_HMAC(phrase.constData(), phrase.length(),
                      (const unsigned char*)salt.constData(), salt.length(),
                      ITERATIONS, EVP_sha256(),
                      KEY_LENGTH, (unsigned char*)key.data());

    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr);
    EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, IV_LENGTH, nullptr);
    EVP_EncryptInit_ex(ctx, nullptr, nullptr, (unsigned char*)key.data(), (unsigned char*)iv.data());

    QByteArray pText = plainText.toUtf8();
    QByteArray cText(pText.length() + EVP_MAX_BLOCK_LENGTH, 0);
    int len = 0, ciphertext_len = 0;

    EVP_EncryptUpdate(ctx, (unsigned char*)cText.data(), &len, (const unsigned char*)pText.constData(), pText.length());
    ciphertext_len = len;

    EVP_EncryptFinal_ex(ctx, (unsigned char*)cText.data() + len, &len);
    ciphertext_len += len;
    cText.resize(ciphertext_len);

    QByteArray tag(16, 0);
    EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, 16, tag.data());
    EVP_CIPHER_CTX_free(ctx);

    QByteArray combined = salt + iv + cText + tag;
    return PREFIX + QString::fromLatin1(combined.toBase64());
}

QString WedgeSecurity::decrypt(const QString& cipherTextBase64, const QString& passPhrase) {
    if (!isEncrypted(cipherTextBase64)) return QString();

    QString cleanInput = cipherTextBase64.mid(PREFIX.length());
    QByteArray combined = QByteArray::fromBase64(cleanInput.toLatin1());
    if (combined.length() < SALT_LENGTH + IV_LENGTH + 16) return QString();

    QByteArray salt = combined.left(SALT_LENGTH);
    QByteArray iv = combined.mid(SALT_LENGTH, IV_LENGTH);
    QByteArray encrypted = combined.mid(SALT_LENGTH + IV_LENGTH, combined.length() - SALT_LENGTH - IV_LENGTH - 16);
    QByteArray tag = combined.right(16);

    QByteArray key(KEY_LENGTH, 0);
    QByteArray phrase = passPhrase.toUtf8();
    PKCS5_PBKDF2_HMAC(phrase.constData(), phrase.length(),
                      (const unsigned char*)salt.constData(), salt.length(),
                      ITERATIONS, EVP_sha256(),
                      KEY_LENGTH, (unsigned char*)key.data());

    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr);
    EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, IV_LENGTH, nullptr);
    EVP_DecryptInit_ex(ctx, nullptr, nullptr, (unsigned char*)key.data(), (unsigned char*)iv.data());

    QByteArray pText(encrypted.length() + EVP_MAX_BLOCK_LENGTH, 0);
    int len = 0, plaintext_len = 0;

    EVP_DecryptUpdate(ctx, (unsigned char*)pText.data(), &len, (const unsigned char*)encrypted.constData(), encrypted.length());
    plaintext_len = len;

    EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, 16, tag.data());
    int ret = EVP_DecryptFinal_ex(ctx, (unsigned char*)pText.data() + len, &len);
    EVP_CIPHER_CTX_free(ctx);

    if (ret > 0) {
        plaintext_len += len;
        pText.resize(plaintext_len);
        return QString::fromUtf8(pText);
    }
    return QString();
}