#include "rsa_pss_sha256.h"
#include "sha.h"
#include "random.h"
#include <cstring>

namespace netplus {

static inline void u32be(uint32_t x, uint8_t out[4]) {
    out[0] = (x >> 24) & 0xFF;
    out[1] = (x >> 16) & 0xFF;
    out[2] = (x >> 8) & 0xFF;
    out[3] = x & 0xFF;
}

std::vector<uint8_t> rsa_pss_sha256::sha256(const std::vector<uint8_t>& in) {
    return sha256_hash(in);
}

// MGF1(SHA256)
std::vector<uint8_t> rsa_pss_sha256::mgf1_sha256(const std::vector<uint8_t>& seed,
                                                 size_t outLen)
{
    std::vector<uint8_t> out;
    out.reserve(outLen);

    uint32_t counter = 0;
    while (out.size() < outLen) {
        uint8_t c[4];
        u32be(counter++, c);

        std::vector<uint8_t> tmp(seed);
        tmp.insert(tmp.end(), c, c + 4);

        std::vector<uint8_t> h = sha256(tmp);

        size_t need = outLen - out.size();
        size_t take = (need < h.size()) ? need : h.size();
        out.insert(out.end(), h.begin(), h.begin() + take);
    }
    return out;
}

// EMSA-PSS-ENCODE (RFC 8017)
std::vector<uint8_t> rsa_pss_sha256::emsa_pss_encode(const std::vector<uint8_t>& mHash,
                                                     size_t emBits)
{
    const size_t hLen = 32;
    const size_t sLen = 32; // TLS 1.3 üblicherweise 32 Byte Salt bei SHA256

    size_t emLen = (emBits + 7) / 8;
    if (emLen < hLen + sLen + 2)
        return {};

    // salt random
    std::vector<uint8_t> salt(sLen);
    fillRandom(salt);

    // M' = 0x00..00 (8) || mHash || salt
    std::vector<uint8_t> Mprime(8, 0x00);
    Mprime.insert(Mprime.end(), mHash.begin(), mHash.end());
    Mprime.insert(Mprime.end(), salt.begin(), salt.end());

    std::vector<uint8_t> H = sha256(Mprime);

    // DB = PS || 0x01 || salt
    size_t psLen = emLen - sLen - hLen - 2;
    std::vector<uint8_t> DB(psLen, 0x00);
    DB.push_back(0x01);
    DB.insert(DB.end(), salt.begin(), salt.end());

    // dbMask = MGF1(H, emLen - hLen - 1)
    std::vector<uint8_t> dbMask = mgf1_sha256(H, emLen - hLen - 1);

    // maskedDB = DB XOR dbMask
    std::vector<uint8_t> maskedDB(DB.size());
    for (size_t i = 0; i < DB.size(); i++)
        maskedDB[i] = DB[i] ^ dbMask[i];

    // leftmost bits = 0
    int unusedBits = int(8 * emLen - emBits);
    if (unusedBits > 0)
        maskedDB[0] &= (0xFF >> unusedBits);

    // EM = maskedDB || H || 0xBC
    std::vector<uint8_t> EM;
    EM.reserve(emLen);
    EM.insert(EM.end(), maskedDB.begin(), maskedDB.end());
    EM.insert(EM.end(), H.begin(), H.end());
    EM.push_back(0xBC);

    return EM;
}

// SIGN: RSASSA-PSS
std::vector<uint8_t> rsa_pss_sha256::sign(netplus::rsa& key,
                                         const std::vector<uint8_t>& message)
{
    // hash message
    std::vector<uint8_t> mHash = sha256(message);
    if (mHash.size() != 32)
        return {};

    // emBits = modBits-1
    size_t modBits = key.n.bitLength();
    size_t emBits = modBits - 1;

    std::vector<uint8_t> EM = emsa_pss_encode(mHash, emBits);
    if (EM.empty())
        return {};

    netplus::rsa::bigInt m = rsa::bytesToBigIntBE(EM);

    // s = m^d mod n  (raw RSA)
    netplus::rsa::bigInt s = netplus::rsa::modPow(m, key.d, key.n);

    size_t k = (modBits + 7) / 8;
    std::vector<uint8_t> out = rsa::bigIntToBytesBE(s, k);
    return out;
}

} // namespace netplus

