#include <openssl/evp.h>
#include <openssl/rand.h>
#include <openssl/sha.h>
#include <stdexcept>
#include <string>
#include <vector>
#include <cstring>

constexpr int IV_LEN = 12;   // 96-bit recommended for GCM
constexpr int KEY_LEN = 32;   // 256-bit key
constexpr int SALT_LEN = 16;   // for PBKDF2
constexpr int TAG_LEN = 16;   // GCM tag length

using ByteVec = std::vector<unsigned char>;

// Generate random bytes
ByteVec rb(int bytes) {
    ByteVec buf(bytes);
    if (RAND_bytes(buf.data(), bytes) != 1)
        throw std::runtime_error("RAND_bytes failed");
    return buf;
}

// AES-256-GCM encryption with raw key
ByteVec ec_wky(const std::string& plain, const ByteVec& key) {
    if (key.size() != KEY_LEN)
        throw std::runtime_error("Key must be 32 bytes");

    ByteVec iv = rb(IV_LEN);
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) throw std::runtime_error("EVP_CIPHER_CTX_new failed");

    ByteVec ciphertext(plain.size() + EVP_MAX_BLOCK_LENGTH);
    int len = 0, ciphertext_len = 0;

    if (EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr) != 1)
        throw std::runtime_error("EncryptInit failed");

    EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, IV_LEN, nullptr);

    if (EVP_EncryptInit_ex(ctx, nullptr, nullptr, key.data(), iv.data()) != 1)
        throw std::runtime_error("EncryptInit key/iv failed");

    if (EVP_EncryptUpdate(ctx, ciphertext.data(), &len,
        reinterpret_cast<const unsigned char*>(plain.data()), plain.size()) != 1)
        throw std::runtime_error("EncryptUpdate failed");
    ciphertext_len = len;

    if (EVP_EncryptFinal_ex(ctx, ciphertext.data() + len, &len) != 1)
        throw std::runtime_error("EncryptFinal failed");
    ciphertext_len += len;

    ByteVec tag(TAG_LEN);
    EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, TAG_LEN, tag.data());

    EVP_CIPHER_CTX_free(ctx);

    ciphertext.resize(ciphertext_len);

    ByteVec output;
    output.reserve(IV_LEN + ciphertext.size() + TAG_LEN);
    output.insert(output.end(), iv.begin(), iv.end());
    output.insert(output.end(), ciphertext.begin(), ciphertext.end());
    output.insert(output.end(), tag.begin(), tag.end());

    return output;
}

// AES-256-GCM decryption with raw key
std::string dec_wky(const ByteVec& data, const ByteVec& key) {
    if (key.size() != KEY_LEN)
        throw std::runtime_error("Key must be 32 bytes");
    if (data.size() < IV_LEN + TAG_LEN)
        throw std::runtime_error("Data too short");

    const unsigned char* iv = data.data();
    const unsigned char* tag = data.data() + data.size() - TAG_LEN;
    const unsigned char* ciphertext = data.data() + IV_LEN;
    size_t ciphertext_len = data.size() - IV_LEN - TAG_LEN;

    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) throw std::runtime_error("EVP_CIPHER_CTX_new failed");

    if (EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr) != 1)
        throw std::runtime_error("DecryptInit failed");

    EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, IV_LEN, nullptr);
    if (EVP_DecryptInit_ex(ctx, nullptr, nullptr, key.data(), iv) != 1)
        throw std::runtime_error("DecryptInit key/iv failed");

    ByteVec plain(ciphertext_len + EVP_MAX_BLOCK_LENGTH);
    int len = 0, plain_len = 0;

    if (EVP_DecryptUpdate(ctx, plain.data(), &len, ciphertext, ciphertext_len) != 1)
        throw std::runtime_error("DecryptUpdate failed");
    plain_len = len;

    EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, TAG_LEN, (void*)tag);

    int ret = EVP_DecryptFinal_ex(ctx, plain.data() + len, &len);
    EVP_CIPHER_CTX_free(ctx);

    if (ret <= 0)
        throw std::runtime_error("Decryption failed (auth tag mismatch)");

    plain_len += len;
    plain.resize(plain_len);

    return std::string(reinterpret_cast<char*>(plain.data()), plain.size());
}

// PBKDF2 derive key
struct DerivedKey {
    ByteVec key;
    ByteVec salt;
};

DerivedKey dervk_p(const std::string& password, const ByteVec* saltOpt = nullptr, int iterations = 200000) {
    ByteVec salt = saltOpt ? *saltOpt : rb(SALT_LEN);
    ByteVec key(KEY_LEN);
    if (PKCS5_PBKDF2_HMAC(password.c_str(), password.size(),
        salt.data(), salt.size(),
        iterations, EVP_sha512(), KEY_LEN, key.data()) != 1)
        throw std::runtime_error("PBKDF2 failed");
    return { key, salt };
}

// Encrypt using password
ByteVec ec_p(const std::string& plain, const std::string& password) {
    DerivedKey dk = dervk_p(password);
    ByteVec payload = ec_wky(plain, dk.key);
    ByteVec result;
    result.reserve(SALT_LEN + payload.size());
    result.insert(result.end(), dk.salt.begin(), dk.salt.end());
    result.insert(result.end(), payload.begin(), payload.end());
    return result;
}

// Decrypt using password
std::string dec_p(const ByteVec& data, const std::string& password, int iterations = 200000) {
    if (data.size() < SALT_LEN + IV_LEN + TAG_LEN)
        throw std::runtime_error("Data too short");

    ByteVec salt(data.begin(), data.begin() + SALT_LEN);
    ByteVec payload(data.begin() + SALT_LEN, data.end());
    DerivedKey dk = dervk_p(password, &salt, iterations);
    return dec_wky(payload, dk.key);
}