#include "common/OidcCrypto.h"

#include <openssl/bn.h>
#include <openssl/crypto.h>
#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <openssl/rsa.h>

#include <array>
#include <memory>
#include <stdexcept>
#include <vector>

namespace visione::oidc
{
namespace
{
Json::Value parseJsonText(const std::string &text)
{
    Json::CharReaderBuilder builder;
    builder["collectComments"] = false;
    std::string errors;
    Json::Value value;
    std::unique_ptr<Json::CharReader> reader(builder.newCharReader());
    if (!reader->parse(text.data(), text.data() + text.size(), &value, &errors))
    {
        throw std::runtime_error("invalid_json");
    }
    return value;
}

using EvpKeyPtr = std::unique_ptr<EVP_PKEY, decltype(&EVP_PKEY_free)>;
using BnPtr = std::unique_ptr<BIGNUM, decltype(&BN_free)>;
using RsaPtr = std::unique_ptr<RSA, decltype(&RSA_free)>;
using PkeyCtxPtr = std::unique_ptr<EVP_PKEY_CTX, decltype(&EVP_PKEY_CTX_free)>;
using CipherCtxPtr =
    std::unique_ptr<EVP_CIPHER_CTX, decltype(&EVP_CIPHER_CTX_free)>;

BnPtr bnFromBase64Url(const std::string &value)
{
    const auto decoded = base64UrlDecode(value);
    BIGNUM *raw = BN_bin2bn(reinterpret_cast<const unsigned char *>(decoded.data()),
                            static_cast<int>(decoded.size()),
                            nullptr);
    if (!raw)
    {
        throw std::runtime_error("invalid_jwk_bn");
    }
    return {raw, &BN_free};
}

EvpKeyPtr keyFromJwk(const Jwk &jwk, bool includePrivate)
{
    if (jwk.kty != "RSA")
    {
        throw std::runtime_error("unsupported_jwk_type");
    }

    RsaPtr rsa(RSA_new(), &RSA_free);
    if (!rsa)
    {
        throw std::runtime_error("rsa_new_failed");
    }

    auto n = bnFromBase64Url(jwk.n);
    auto e = bnFromBase64Url(jwk.e);
    BIGNUM *nRaw = n.release();
    BIGNUM *eRaw = e.release();
    if (RSA_set0_key(rsa.get(), nRaw, eRaw, nullptr) != 1)
    {
        BN_free(nRaw);
        BN_free(eRaw);
        throw std::runtime_error("rsa_set_public_failed");
    }

    if (includePrivate)
    {
        if (jwk.d.empty() || jwk.p.empty() || jwk.q.empty())
        {
            throw std::runtime_error("missing_private_jwk_fields");
        }

        auto d = bnFromBase64Url(jwk.d);
        auto p = bnFromBase64Url(jwk.p);
        auto q = bnFromBase64Url(jwk.q);
        BIGNUM *dRaw = d.release();
        if (RSA_set0_key(rsa.get(), nullptr, nullptr, dRaw) != 1)
        {
            BN_free(dRaw);
            throw std::runtime_error("rsa_set_private_failed");
        }

        BIGNUM *pRaw = p.release();
        BIGNUM *qRaw = q.release();
        if (RSA_set0_factors(rsa.get(), pRaw, qRaw) != 1)
        {
            BN_free(pRaw);
            BN_free(qRaw);
            throw std::runtime_error("rsa_set_factors_failed");
        }
    }

    EVP_PKEY *rawPkey = EVP_PKEY_new();
    if (!rawPkey)
    {
        throw std::runtime_error("evp_pkey_new_failed");
    }
    EvpKeyPtr pkey(rawPkey, &EVP_PKEY_free);
    if (EVP_PKEY_assign_RSA(pkey.get(), rsa.release()) != 1)
    {
        throw std::runtime_error("evp_assign_rsa_failed");
    }
    return pkey;
}

std::vector<std::string> split(const std::string &text, char separator)
{
    std::vector<std::string> parts;
    std::string current;
    for (char ch : text)
    {
        if (ch == separator)
        {
            parts.push_back(current);
            current.clear();
            continue;
        }
        current.push_back(ch);
    }
    parts.push_back(current);
    return parts;
}

std::string rsaOaepDecrypt(const std::string &ciphertext,
                           const Jwk &jwk,
                           const std::string &algorithm)
{
    auto pkey = keyFromJwk(jwk, true);
    PkeyCtxPtr ctx(EVP_PKEY_CTX_new(pkey.get(), nullptr), &EVP_PKEY_CTX_free);
    if (!ctx)
    {
        throw std::runtime_error("jwe_pkey_ctx_new_failed");
    }
    if (EVP_PKEY_decrypt_init(ctx.get()) != 1)
    {
        throw std::runtime_error("jwe_pkey_decrypt_init_failed");
    }
    if (EVP_PKEY_CTX_set_rsa_padding(ctx.get(), RSA_PKCS1_OAEP_PADDING) != 1)
    {
        throw std::runtime_error("jwe_rsa_padding_failed");
    }

    const EVP_MD *md = EVP_sha1();
    if (algorithm == "RSA-OAEP-256")
    {
        md = EVP_sha256();
    }
    else if (algorithm != "RSA-OAEP")
    {
        throw std::runtime_error("unsupported_jwe_alg");
    }
    if (EVP_PKEY_CTX_set_rsa_oaep_md(ctx.get(), md) != 1 ||
        EVP_PKEY_CTX_set_rsa_mgf1_md(ctx.get(), md) != 1)
    {
        throw std::runtime_error("jwe_rsa_oaep_md_failed");
    }

    size_t outLen = 0;
    if (EVP_PKEY_decrypt(ctx.get(),
                         nullptr,
                         &outLen,
                         reinterpret_cast<const unsigned char *>(ciphertext.data()),
                         ciphertext.size()) != 1)
    {
        throw std::runtime_error("jwe_cek_size_failed");
    }

    std::string plaintext(outLen, '\0');
    if (EVP_PKEY_decrypt(
            ctx.get(),
            reinterpret_cast<unsigned char *>(plaintext.data()),
            &outLen,
            reinterpret_cast<const unsigned char *>(ciphertext.data()),
            ciphertext.size()) != 1)
    {
        throw std::runtime_error("jwe_cek_decrypt_failed");
    }
    plaintext.resize(outLen);
    return plaintext;
}

std::array<unsigned char, 8> bigEndianBitLength(uint64_t lenBytes)
{
    const auto lenBits = lenBytes * 8;
    return {
        static_cast<unsigned char>((lenBits >> 56) & 0xff),
        static_cast<unsigned char>((lenBits >> 48) & 0xff),
        static_cast<unsigned char>((lenBits >> 40) & 0xff),
        static_cast<unsigned char>((lenBits >> 32) & 0xff),
        static_cast<unsigned char>((lenBits >> 24) & 0xff),
        static_cast<unsigned char>((lenBits >> 16) & 0xff),
        static_cast<unsigned char>((lenBits >> 8) & 0xff),
        static_cast<unsigned char>(lenBits & 0xff),
    };
}

void requireExact(bool ok, const char *error)
{
    if (!ok)
    {
        throw std::runtime_error(error);
    }
}

std::string decryptCbcHmac(const JweParts &parts,
                           const std::string &cek,
                           size_t expectedCekSize,
                           size_t ivSize,
                           size_t tagSize,
                           const EVP_CIPHER *cipher,
                           const EVP_MD *md)
{
    if (cek.size() != expectedCekSize)
    {
        throw std::runtime_error("invalid_jwe_cek_size");
    }

    const auto iv = base64UrlDecode(parts.encodedIv);
    const auto ciphertext = base64UrlDecode(parts.encodedCiphertext);
    const auto tag = base64UrlDecode(parts.encodedTag);
    if (iv.size() != ivSize || tag.size() != tagSize)
    {
        throw std::runtime_error("invalid_jwe_parts");
    }

    const auto macKeySize = expectedCekSize / 2;
    const auto encKeySize = expectedCekSize / 2;
    const std::string macKey = cek.substr(0, macKeySize);
    const std::string encKey = cek.substr(macKeySize, encKeySize);
    const auto al = bigEndianBitLength(parts.encodedHeader.size());

    std::string macInput;
    macInput.reserve(parts.encodedHeader.size() + iv.size() + ciphertext.size() +
                     al.size());
    macInput.append(parts.encodedHeader);
    macInput.append(iv);
    macInput.append(ciphertext);
    macInput.append(reinterpret_cast<const char *>(al.data()), al.size());

    unsigned char hmac[EVP_MAX_MD_SIZE];
    unsigned int hmacLen = 0;
    if (!HMAC(md,
              macKey.data(),
              static_cast<int>(macKey.size()),
              reinterpret_cast<const unsigned char *>(macInput.data()),
              macInput.size(),
              hmac,
              &hmacLen))
    {
        throw std::runtime_error("jwe_hmac_failed");
    }
    if (hmacLen < tag.size() ||
        CRYPTO_memcmp(hmac, tag.data(), tag.size()) != 0)
    {
        throw std::runtime_error("jwe_auth_tag_mismatch");
    }

    CipherCtxPtr ctx(EVP_CIPHER_CTX_new(), &EVP_CIPHER_CTX_free);
    if (!ctx)
    {
        throw std::runtime_error("jwe_cipher_ctx_new_failed");
    }
    requireExact(EVP_DecryptInit_ex(
                     ctx.get(),
                     cipher,
                     nullptr,
                     reinterpret_cast<const unsigned char *>(encKey.data()),
                     reinterpret_cast<const unsigned char *>(iv.data())) == 1,
                 "jwe_decrypt_init_failed");

    std::string plaintext(ciphertext.size() + EVP_MAX_BLOCK_LENGTH, '\0');
    int outLen1 = 0;
    requireExact(EVP_DecryptUpdate(
                     ctx.get(),
                     reinterpret_cast<unsigned char *>(plaintext.data()),
                     &outLen1,
                     reinterpret_cast<const unsigned char *>(ciphertext.data()),
                     ciphertext.size()) == 1,
                 "jwe_decrypt_update_failed");

    int outLen2 = 0;
    requireExact(EVP_DecryptFinal_ex(
                     ctx.get(),
                     reinterpret_cast<unsigned char *>(plaintext.data()) + outLen1,
                     &outLen2) == 1,
                 "jwe_decrypt_final_failed");
    plaintext.resize(outLen1 + outLen2);
    return plaintext;
}
}  // namespace

JweParts parseJwe(const std::string &jwe)
{
    const auto parts = split(jwe, '.');
    if (parts.size() != 5)
    {
        throw std::runtime_error("invalid_jwe");
    }

    JweParts parsed;
    parsed.encodedHeader = parts[0];
    parsed.encodedEncryptedKey = parts[1];
    parsed.encodedIv = parts[2];
    parsed.encodedCiphertext = parts[3];
    parsed.encodedTag = parts[4];
    parsed.header = parseJsonText(base64UrlDecode(parts[0]));
    return parsed;
}

std::string decryptJwe(const std::string &jwe, const Jwk &jwk)
{
    const auto parsed = parseJwe(jwe);
    const auto alg = parsed.header.get("alg", "").asString();
    const auto enc = parsed.header.get("enc", "").asString();

    const auto cek =
        rsaOaepDecrypt(base64UrlDecode(parsed.encodedEncryptedKey), jwk, alg);
    if (enc == "A256CBC-HS512")
    {
        return decryptCbcHmac(parsed,
                              cek,
                              64,
                              16,
                              32,
                              EVP_aes_256_cbc(),
                              EVP_sha512());
    }
    if (enc == "A128CBC-HS256")
    {
        return decryptCbcHmac(parsed,
                              cek,
                              32,
                              16,
                              16,
                              EVP_aes_128_cbc(),
                              EVP_sha256());
    }
    throw std::runtime_error("unsupported_jwe_enc");
}
}  // namespace visione::oidc