#include "rsa_pss_sha384.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_sha384::sha384(const std::vector<uint8_t>& in) {
    return sha384_hash(in);
}

// MGF1(SHA384)
std::vector<uint8_t> rsa_pss_sha384::mgf1_sha384(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 = sha384(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) with SHA-384
std::vector<uint8_t> rsa_pss_sha384::emsa_pss_encode(const std::vector<uint8_t>& mHash,
                                                     size_t emBits)
{
    const size_t hLen = 48;  // SHA-384 produces 48-byte hashes
    const size_t sLen = 48;  // Salt length for SHA-384

    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 = sha384(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_sha384(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 with SHA-384
std::vector<uint8_t> rsa_pss_sha384::sign(netplus::rsa& key,
                                         const std::vector<uint8_t>& message)
{
    // hash message with SHA-384
    std::vector<uint8_t> mHash = sha384(message);
    if (mHash.size() != 48)
        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;
}

// VERIFY: RSASSA-PSS with SHA-384
bool rsa_pss_sha384::verify(const netplus::rsa& key,
                           const std::vector<uint8_t>& signature,
                           const std::vector<uint8_t>& message)
{
    // hash message with SHA-384
    std::vector<uint8_t> mHash = sha384(message);
    if (mHash.size() != 48)
        return false;

    // Decrypt signature to get EM
    size_t modBits = key.n.bitLength();
    size_t emLen = (modBits + 7) / 8;

    // Convert signature to bigInt and do raw RSA verify: m = s^e mod n
    netplus::rsa::bigInt s = rsa::bytesToBigIntBE(signature);
    netplus::rsa::bigInt m = netplus::rsa::modPow(s, key.e, key.n);

    std::vector<uint8_t> EM = rsa::bigIntToBytesBE(m, emLen);
    
    if (EM.empty() || EM.back() != 0xBC)
        return false;

    // Extract H and verify
    const size_t hLen = 48;
    const size_t sLen = 48;
    size_t psLen = emLen - sLen - hLen - 2;

    // Extract maskedDB and H
    std::vector<uint8_t> maskedDB(EM.begin(), EM.begin() + psLen + 1 + sLen);
    std::vector<uint8_t> H(EM.begin() + psLen + 1 + sLen, EM.end() - 1);

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

    // Verify DB format
    if (DB[psLen] != 0x01)
        return false;

    std::vector<uint8_t> salt(DB.begin() + psLen + 1, DB.end());

    // Verify H
    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_expected = sha384(Mprime);

    return H == H_expected;
}

} // namespace netplus
