#include "common/OidcJwt.h"

#include <drogon/utils/Utilities.h>

#include <jwt-cpp/traits/open-source-parsers-jsoncpp/defaults.h>
#include <openssl/bn.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/rsa.h>
#include <openssl/sha.h>

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

namespace visione::oidc
{
namespace
{
using BnPtr = std::unique_ptr<BIGNUM, decltype(&BN_free)>;
using RsaPtr = std::unique_ptr<RSA, decltype(&RSA_free)>;
using EvpKeyPtr = std::unique_ptr<EVP_PKEY, decltype(&EVP_PKEY_free)>;
using BioPtr = std::unique_ptr<BIO, decltype(&BIO_free)>;
using BnCtxPtr = std::unique_ptr<BN_CTX, decltype(&BN_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");
        }

        const BIGNUM *rsaD = nullptr;
        const BIGNUM *rsaP = nullptr;
        const BIGNUM *rsaQ = nullptr;
        const BIGNUM *rsaE = nullptr;
        RSA_get0_key(rsa.get(), nullptr, &rsaE, &rsaD);
        RSA_get0_factors(rsa.get(), &rsaP, &rsaQ);

        BnCtxPtr bnCtx(BN_CTX_new(), &BN_CTX_free);
        if (!bnCtx)
        {
            throw std::runtime_error("bn_ctx_new_failed");
        }

        BIGNUM *dp = BN_new();
        BIGNUM *dq = BN_new();
        BIGNUM *qi = BN_new();
        BIGNUM *pMinus1 = BN_dup(rsaP);
        BIGNUM *qMinus1 = BN_dup(rsaQ);
        if (!dp || !dq || !qi || !pMinus1 || !qMinus1)
        {
            BN_free(dp);
            BN_free(dq);
            BN_free(qi);
            BN_free(pMinus1);
            BN_free(qMinus1);
            throw std::runtime_error("bn_alloc_failed");
        }

        BN_sub_word(pMinus1, 1);
        BN_sub_word(qMinus1, 1);
        if (BN_mod(dp, rsaD, pMinus1, bnCtx.get()) != 1 ||
            BN_mod(dq, rsaD, qMinus1, bnCtx.get()) != 1 ||
            BN_mod_inverse(qi, rsaQ, rsaP, bnCtx.get()) == nullptr)
        {
            BN_free(dp);
            BN_free(dq);
            BN_free(qi);
            BN_free(pMinus1);
            BN_free(qMinus1);
            throw std::runtime_error("rsa_crt_params_failed");
        }
        BN_free(pMinus1);
        BN_free(qMinus1);

        if (RSA_set0_crt_params(rsa.get(), dp, dq, qi) != 1)
        {
            BN_free(dp);
            BN_free(dq);
            BN_free(qi);
            throw std::runtime_error("rsa_set_crt_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::string pemFromJwk(const Jwk &jwk, bool includePrivate)
{
    auto pkey = keyFromJwk(jwk, includePrivate);
    BioPtr bio(BIO_new(BIO_s_mem()), &BIO_free);
    if (!bio)
    {
        throw std::runtime_error("bio_new_failed");
    }

    const int ok = includePrivate
                       ? PEM_write_bio_PrivateKey(
                             bio.get(), pkey.get(), nullptr, nullptr, 0, nullptr, nullptr)
                       : PEM_write_bio_PUBKEY(bio.get(), pkey.get());
    if (ok != 1)
    {
        throw std::runtime_error("pem_write_failed");
    }

    BUF_MEM *buffer = nullptr;
    BIO_get_mem_ptr(bio.get(), &buffer);
    if (!buffer || !buffer->data || buffer->length == 0)
    {
        throw std::runtime_error("pem_buffer_failed");
    }
    return {buffer->data, buffer->length};
}

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;
}

template <typename Builder>
void setClaims(Builder &builder,
               const Json::Value &value,
               bool header)
{
    if (!value.isObject())
    {
        throw std::runtime_error("invalid_jwt_claims_json");
    }

    for (const auto &name : value.getMemberNames())
    {
        if (header)
        {
            builder.set_header_claim(name, value[name]);
        }
        else
        {
            builder.set_payload_claim(name, value[name]);
        }
    }
}
}  // namespace

Jwk parseJwk(const Json::Value &value)
{
    Jwk jwk;
    jwk.kty = value.get("kty", "").asString();
    jwk.use = value.get("use", "").asString();
    jwk.alg = value.get("alg", "").asString();
    jwk.kid = value.get("kid", "").asString();
    jwk.n = value.get("n", "").asString();
    jwk.e = value.get("e", "").asString();
    jwk.d = value.get("d", "").asString();
    jwk.p = value.get("p", "").asString();
    jwk.q = value.get("q", "").asString();
    return jwk;
}

Json::Value jwkToPublicJson(const Jwk &jwk)
{
    Json::Value value;
    value["kty"] = jwk.kty;
    value["use"] = jwk.use;
    if (!jwk.alg.empty())
    {
        value["alg"] = jwk.alg;
    }
    value["kid"] = jwk.kid;
    value["n"] = jwk.n;
    value["e"] = jwk.e;
    return value;
}

std::string base64UrlEncode(const std::string &value)
{
    return drogon::utils::base64EncodeUnpadded(value, true);
}

std::string base64UrlDecode(const std::string &value)
{
    std::string padded = value;
    while (padded.size() % 4 != 0)
    {
        padded.push_back('=');
    }
    for (char &ch : padded)
    {
        if (ch == '-')
        {
            ch = '+';
        }
        else if (ch == '_')
        {
            ch = '/';
        }
    }
    return drogon::utils::base64Decode(padded);
}

std::string randomUrlSafeString(size_t sizeBytes)
{
    return base64UrlEncode(drogon::utils::genRandomString(sizeBytes));
}

std::string signJwt(const Json::Value &header,
                    const Json::Value &payload,
                    const Jwk &jwk)
{
    auto builder = jwt::create();
    setClaims(builder, header, true);
    setClaims(builder, payload, false);
    return builder.sign(
        jwt::algorithm::rs256(pemFromJwk(jwk, false), pemFromJwk(jwk, true)));
}

bool verifyJwtSignature(const std::string &jwtText, const Jwk &jwk)
{
    try
    {
        auto decoded = jwt::decode(jwtText);
        auto verifier = jwt::verify().allow_algorithm(
            jwt::algorithm::rs256(pemFromJwk(jwk, false)));
        verifier.verify(decoded);
        return true;
    }
    catch (const std::exception &)
    {
        return false;
    }
}

JwtParts parseJwt(const std::string &jwtText)
{
    const auto decoded = jwt::decode(jwtText);

    JwtParts parsed;
    parsed.encodedHeader = decoded.get_header_base64();
    parsed.encodedPayload = decoded.get_payload_base64();
    parsed.encodedSignature = decoded.get_signature_base64();
    parsed.header = decoded.get_header_json();
    parsed.payload = decoded.get_payload_json();
    return parsed;
}

Json::Value parseJsonObjectJwtPayload(const std::string &jwtText)
{
    const auto parsed = parseJwt(jwtText);
    if (!parsed.payload.isObject())
    {
        throw std::runtime_error("invalid_jwt_payload");
    }
    return parsed.payload;
}

std::optional<Jwk> findJwkByKid(const std::vector<Jwk> &jwks,
                                const std::string &kid)
{
    for (const auto &jwk : jwks)
    {
        if (jwk.kid == kid)
        {
            return jwk;
        }
    }
    return std::nullopt;
}

std::string sha256Base64Url(const std::string &value)
{
    unsigned char digest[SHA256_DIGEST_LENGTH];
    SHA256(reinterpret_cast<const unsigned char *>(value.data()),
           value.size(),
           digest);
    return drogon::utils::base64EncodeUnpadded(digest,
                                               SHA256_DIGEST_LENGTH,
                                               true);
}

std::string leftHashBase64Url(const std::string &value)
{
    const auto digest = base64UrlDecode(sha256Base64Url(value));
    return base64UrlEncode(digest.substr(0, digest.size() / 2));
}
}  // namespace visione::oidc