#include "rsa_pkcs1_sha256.h"
#include "sha.h"
#include <cstring>

namespace netplus {

// DER encoding of DigestInfo for SHA-256 (RFC 8017 Section 9.2, Note 1)
// SEQUENCE { SEQUENCE { OID sha256, NULL }, OCTET STRING hash }
static const uint8_t SHA256_DIGEST_INFO_PREFIX[] = {
    0x30, 0x31,             // SEQUENCE (49 bytes)
    0x30, 0x0d,             //   SEQUENCE (13 bytes)
    0x06, 0x09,             //     OID (9 bytes) = 2.16.840.1.101.3.4.2.1 (SHA-256)
    0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01,
    0x05, 0x00,             //     NULL
    0x04, 0x20              //   OCTET STRING (32 bytes = SHA-256 hash follows)
};

// EMSA-PKCS1-v1_5-ENCODE (RFC 8017 Section 9.2)
std::vector<uint8_t> rsa_pkcs1_sha256::emsa_pkcs1_v15_encode(
    const std::vector<uint8_t>& mHash, size_t emLen)
{
    // T = DigestInfo prefix || Hash
    const size_t tLen = sizeof(SHA256_DIGEST_INFO_PREFIX) + mHash.size();

    if (emLen < tLen + 11)
        return {};

    // EM = 0x00 || 0x01 || PS || 0x00 || T
    // PS = 0xFF bytes, length = emLen - tLen - 3
    std::vector<uint8_t> EM(emLen);
    size_t idx = 0;

    EM[idx++] = 0x00;
    EM[idx++] = 0x01;

    size_t psLen = emLen - tLen - 3;
    std::memset(&EM[idx], 0xFF, psLen);
    idx += psLen;

    EM[idx++] = 0x00;

    std::memcpy(&EM[idx], SHA256_DIGEST_INFO_PREFIX, sizeof(SHA256_DIGEST_INFO_PREFIX));
    idx += sizeof(SHA256_DIGEST_INFO_PREFIX);

    std::memcpy(&EM[idx], mHash.data(), mHash.size());

    return EM;
}

// RSASSA-PKCS1-v1_5-SIGN (RFC 8017 Section 8.2.1)
std::vector<uint8_t> rsa_pkcs1_sha256::sign(netplus::rsa& key,
                                            const std::vector<uint8_t>& message)
{
    // Step 1: Hash the message with SHA-256
    std::vector<uint8_t> mHash = sha256_hash(message);
    if (mHash.size() != 32)
        return {};

    // Step 2: Encode to EMSA-PKCS1-v1_5
    size_t modBits = key.n.bitLength();
    size_t k = (modBits + 7) / 8;  // RSA modulus length in bytes

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

    // Step 3: RSA signature primitive: s = m^d mod n
    rsa::bigInt m = rsa::bytesToBigIntBE(EM);
    rsa::bigInt s = rsa::modPow(m, key.d, key.n);

    // Step 4: Convert to byte string
    std::vector<uint8_t> out = rsa::bigIntToBytesBE(s, k);
    return out;
}

} // namespace netplus
