#include "aes_gcm.hpp"
#include <openssl/evp.h>
#include <openssl/rand.h>
#include <openssl/err.h>
#include <nlohmann/json.hpp>
#include <cstring>
#include <stdexcept>

namespace smartbotic::crypto {

using namespace common;

// Constants
constexpr size_t AES_256_KEY_SIZE = 32;
constexpr size_t GCM_IV_SIZE = 12;
constexpr size_t GCM_TAG_SIZE = 16;

// Base64 encoding table
static const char base64_chars[] =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    "abcdefghijklmnopqrstuvwxyz"
    "0123456789+/";

std::string base64::encode(const std::vector<uint8_t>& data) {
    return encode(data.data(), data.size());
}

std::string base64::encode(const uint8_t* data, size_t len) {
    std::string result;
    result.reserve(((len + 2) / 3) * 4);

    for (size_t i = 0; i < len; i += 3) {
        uint32_t octet_a = i < len ? data[i] : 0;
        uint32_t octet_b = i + 1 < len ? data[i + 1] : 0;
        uint32_t octet_c = i + 2 < len ? data[i + 2] : 0;

        uint32_t triple = (octet_a << 16) | (octet_b << 8) | octet_c;

        result += base64_chars[(triple >> 18) & 0x3F];
        result += base64_chars[(triple >> 12) & 0x3F];
        result += (i + 1 < len) ? base64_chars[(triple >> 6) & 0x3F] : '=';
        result += (i + 2 < len) ? base64_chars[triple & 0x3F] : '=';
    }

    return result;
}

Result<std::vector<uint8_t>> base64::decode(const std::string& encoded) {
    if (encoded.empty()) {
        return std::vector<uint8_t>{};
    }

    // Build decoding table
    static int decode_table[256] = {-1};
    static bool initialized = false;
    if (!initialized) {
        for (int i = 0; i < 256; ++i) decode_table[i] = -1;
        for (int i = 0; i < 64; ++i) {
            decode_table[static_cast<unsigned char>(base64_chars[i])] = i;
        }
        initialized = true;
    }

    size_t len = encoded.length();
    if (len % 4 != 0) {
        return Error(ErrorCode::InvalidArgument, "Invalid base64 length");
    }

    // Calculate output size
    size_t out_len = (len / 4) * 3;
    if (len >= 1 && encoded[len - 1] == '=') out_len--;
    if (len >= 2 && encoded[len - 2] == '=') out_len--;

    std::vector<uint8_t> result(out_len);

    size_t j = 0;
    for (size_t i = 0; i < len; i += 4) {
        int a = decode_table[static_cast<unsigned char>(encoded[i])];
        int b = decode_table[static_cast<unsigned char>(encoded[i + 1])];
        int c = encoded[i + 2] == '=' ? 0 : decode_table[static_cast<unsigned char>(encoded[i + 2])];
        int d = encoded[i + 3] == '=' ? 0 : decode_table[static_cast<unsigned char>(encoded[i + 3])];

        if (a == -1 || b == -1 || (encoded[i + 2] != '=' && c == -1) || (encoded[i + 3] != '=' && d == -1)) {
            return Error(ErrorCode::InvalidArgument, "Invalid base64 character");
        }

        uint32_t triple = (a << 18) | (b << 12) | (c << 6) | d;

        if (j < out_len) result[j++] = (triple >> 16) & 0xFF;
        if (j < out_len) result[j++] = (triple >> 8) & 0xFF;
        if (j < out_len) result[j++] = triple & 0xFF;
    }

    return result;
}

std::string EncryptedData::toJson() const {
    nlohmann::json j;
    j["ciphertext"] = base64::encode(ciphertext);
    j["iv"] = base64::encode(iv);
    j["auth_tag"] = base64::encode(auth_tag);
    return j.dump();
}

Result<EncryptedData> EncryptedData::fromJson(const std::string& json) {
    try {
        auto j = nlohmann::json::parse(json);

        EncryptedData data;

        auto ct_result = base64::decode(j.value("ciphertext", ""));
        if (ct_result.failed()) return ct_result.error();
        data.ciphertext = ct_result.value();

        auto iv_result = base64::decode(j.value("iv", ""));
        if (iv_result.failed()) return iv_result.error();
        data.iv = iv_result.value();

        auto tag_result = base64::decode(j.value("auth_tag", ""));
        if (tag_result.failed()) return tag_result.error();
        data.auth_tag = tag_result.value();

        return data;
    } catch (const std::exception& e) {
        return Error(ErrorCode::InvalidArgument, "Failed to parse encrypted data JSON: " + std::string(e.what()));
    }
}

std::string EncryptedData::toBase64() const {
    // Format: iv || auth_tag || ciphertext, all concatenated and base64 encoded
    std::vector<uint8_t> combined;
    combined.reserve(iv.size() + auth_tag.size() + ciphertext.size());
    combined.insert(combined.end(), iv.begin(), iv.end());
    combined.insert(combined.end(), auth_tag.begin(), auth_tag.end());
    combined.insert(combined.end(), ciphertext.begin(), ciphertext.end());
    return base64::encode(combined);
}

Result<EncryptedData> EncryptedData::fromBase64(const std::string& encoded) {
    auto decode_result = base64::decode(encoded);
    if (decode_result.failed()) {
        return decode_result.error();
    }

    const auto& data = decode_result.value();
    if (data.size() < GCM_IV_SIZE + GCM_TAG_SIZE) {
        return Error(ErrorCode::InvalidArgument, "Encrypted data too short");
    }

    EncryptedData result;
    result.iv.assign(data.begin(), data.begin() + GCM_IV_SIZE);
    result.auth_tag.assign(data.begin() + GCM_IV_SIZE, data.begin() + GCM_IV_SIZE + GCM_TAG_SIZE);
    result.ciphertext.assign(data.begin() + GCM_IV_SIZE + GCM_TAG_SIZE, data.end());

    return result;
}

AesGcm::AesGcm(const AesGcmConfig& config) : config_(config) {
    if (config_.master_key.empty()) {
        throw std::invalid_argument("Master key cannot be empty");
    }

    // Pre-derive key if not using per-encryption salt
    if (!config_.use_key_derivation) {
        // Use key directly (must be 32 bytes)
        if (config_.master_key.size() < AES_256_KEY_SIZE) {
            // Pad with zeros if too short (not recommended for production)
            derived_key_.resize(AES_256_KEY_SIZE, 0);
            std::memcpy(derived_key_.data(), config_.master_key.data(),
                       std::min(config_.master_key.size(), AES_256_KEY_SIZE));
        } else {
            derived_key_.assign(config_.master_key.begin(),
                               config_.master_key.begin() + AES_256_KEY_SIZE);
        }
    }
}

std::vector<uint8_t> AesGcm::deriveKey(const std::vector<uint8_t>& salt) {
    std::vector<uint8_t> key(AES_256_KEY_SIZE);

    if (PKCS5_PBKDF2_HMAC(
            config_.master_key.c_str(),
            static_cast<int>(config_.master_key.length()),
            salt.data(),
            static_cast<int>(salt.size()),
            config_.pbkdf2_iterations,
            EVP_sha256(),
            static_cast<int>(key.size()),
            key.data()) != 1) {
        throw std::runtime_error("PBKDF2 key derivation failed");
    }

    return key;
}

std::vector<uint8_t> AesGcm::generateRandomIv() {
    std::vector<uint8_t> iv(GCM_IV_SIZE);
    if (RAND_bytes(iv.data(), static_cast<int>(iv.size())) != 1) {
        throw std::runtime_error("Failed to generate random IV");
    }
    return iv;
}

std::string AesGcm::generateRandomKey() {
    std::vector<uint8_t> key(AES_256_KEY_SIZE);
    if (RAND_bytes(key.data(), static_cast<int>(key.size())) != 1) {
        throw std::runtime_error("Failed to generate random key");
    }
    return base64::encode(key);
}

Result<EncryptedData> AesGcm::encrypt(const std::string& plaintext) {
    return encrypt(std::vector<uint8_t>(plaintext.begin(), plaintext.end()));
}

Result<EncryptedData> AesGcm::encrypt(const std::vector<uint8_t>& plaintext) {
    EncryptedData result;

    // Generate random IV
    result.iv = generateRandomIv();

    // Get or derive key
    std::vector<uint8_t> key;
    if (config_.use_key_derivation) {
        // Use IV as salt for key derivation (provides unique key per encryption)
        key = deriveKey(result.iv);
    } else {
        key = derived_key_;
    }

    // Create cipher context
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        return Error(ErrorCode::Internal, "Failed to create cipher context");
    }

    // Cleanup helper
    struct CtxGuard {
        EVP_CIPHER_CTX* ctx;
        ~CtxGuard() { EVP_CIPHER_CTX_free(ctx); }
    } guard{ctx};

    // Initialize encryption
    if (EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr) != 1) {
        return Error(ErrorCode::Internal, "Failed to initialize encryption");
    }

    // Set IV length
    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, GCM_IV_SIZE, nullptr) != 1) {
        return Error(ErrorCode::Internal, "Failed to set IV length");
    }

    // Set key and IV
    if (EVP_EncryptInit_ex(ctx, nullptr, nullptr, key.data(), result.iv.data()) != 1) {
        return Error(ErrorCode::Internal, "Failed to set key and IV");
    }

    // Encrypt
    result.ciphertext.resize(plaintext.size() + EVP_MAX_BLOCK_LENGTH);
    int out_len = 0;
    int total_len = 0;

    if (EVP_EncryptUpdate(ctx, result.ciphertext.data(), &out_len,
                          plaintext.data(), static_cast<int>(plaintext.size())) != 1) {
        return Error(ErrorCode::Internal, "Encryption failed");
    }
    total_len = out_len;

    // Finalize
    if (EVP_EncryptFinal_ex(ctx, result.ciphertext.data() + total_len, &out_len) != 1) {
        return Error(ErrorCode::Internal, "Encryption finalization failed");
    }
    total_len += out_len;
    result.ciphertext.resize(total_len);

    // Get authentication tag
    result.auth_tag.resize(GCM_TAG_SIZE);
    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, GCM_TAG_SIZE, result.auth_tag.data()) != 1) {
        return Error(ErrorCode::Internal, "Failed to get authentication tag");
    }

    return result;
}

Result<std::string> AesGcm::decrypt(const EncryptedData& data) {
    auto result = decryptBytes(data);
    if (result.failed()) {
        return result.error();
    }
    return std::string(result.value().begin(), result.value().end());
}

Result<std::vector<uint8_t>> AesGcm::decryptBytes(const EncryptedData& data) {
    if (data.iv.size() != GCM_IV_SIZE) {
        return Error(ErrorCode::InvalidArgument, "Invalid IV size");
    }
    if (data.auth_tag.size() != GCM_TAG_SIZE) {
        return Error(ErrorCode::InvalidArgument, "Invalid auth tag size");
    }

    // Get or derive key
    std::vector<uint8_t> key;
    if (config_.use_key_derivation) {
        key = deriveKey(data.iv);
    } else {
        key = derived_key_;
    }

    // Create cipher context
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        return Error(ErrorCode::Internal, "Failed to create cipher context");
    }

    struct CtxGuard {
        EVP_CIPHER_CTX* ctx;
        ~CtxGuard() { EVP_CIPHER_CTX_free(ctx); }
    } guard{ctx};

    // Initialize decryption
    if (EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr) != 1) {
        return Error(ErrorCode::Internal, "Failed to initialize decryption");
    }

    // Set IV length
    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, GCM_IV_SIZE, nullptr) != 1) {
        return Error(ErrorCode::Internal, "Failed to set IV length");
    }

    // Set key and IV
    if (EVP_DecryptInit_ex(ctx, nullptr, nullptr, key.data(), data.iv.data()) != 1) {
        return Error(ErrorCode::Internal, "Failed to set key and IV");
    }

    // Set authentication tag
    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, GCM_TAG_SIZE,
                            const_cast<uint8_t*>(data.auth_tag.data())) != 1) {
        return Error(ErrorCode::Internal, "Failed to set authentication tag");
    }

    // Decrypt
    std::vector<uint8_t> plaintext(data.ciphertext.size() + EVP_MAX_BLOCK_LENGTH);
    int out_len = 0;
    int total_len = 0;

    if (EVP_DecryptUpdate(ctx, plaintext.data(), &out_len,
                          data.ciphertext.data(), static_cast<int>(data.ciphertext.size())) != 1) {
        return Error(ErrorCode::Internal, "Decryption failed");
    }
    total_len = out_len;

    // Finalize and verify tag
    if (EVP_DecryptFinal_ex(ctx, plaintext.data() + total_len, &out_len) != 1) {
        // Tag verification failed
        return Error(ErrorCode::InvalidArgument, "Authentication failed - data may be corrupted or tampered");
    }
    total_len += out_len;
    plaintext.resize(total_len);

    return plaintext;
}

Result<std::string> AesGcm::encryptToBase64(const std::string& plaintext) {
    auto result = encrypt(plaintext);
    if (result.failed()) {
        return result.error();
    }
    return result.value().toBase64();
}

Result<std::string> AesGcm::decryptFromBase64(const std::string& encoded) {
    auto data_result = EncryptedData::fromBase64(encoded);
    if (data_result.failed()) {
        return data_result.error();
    }
    return decrypt(data_result.value());
}

} // namespace smartbotic::crypto