/*******************************************************************************
 * Copyright (c) 2026, Jan Koester jan.koester@gmx.net
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 * Neither the name of the <organization> nor the
 *      names of its contributors may be used to endorse or promote products
 *      derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *******************************************************************************/

#include "certgen.h"
#include "../crypto/rsa.h"
#include "../crypto/rsa_pkcs1_sha256.h"
#include "../crypto/sha.h"
#include "../crypto/base64.h"

#include <ctime>
#include <cstring>
#include <fstream>
#include <algorithm>

namespace netplus {

// ═════════════════════════════════════════════════════════════
//  ASN.1 DER builder helpers
// ═════════════════════════════════════════════════════════════

static void derEncodeLength(std::vector<uint8_t> &out, size_t len) {
    if (len < 0x80) {
        out.push_back(static_cast<uint8_t>(len));
    } else if (len <= 0xFF) {
        out.push_back(0x81);
        out.push_back(static_cast<uint8_t>(len));
    } else if (len <= 0xFFFF) {
        out.push_back(0x82);
        out.push_back(static_cast<uint8_t>(len >> 8));
        out.push_back(static_cast<uint8_t>(len & 0xFF));
    } else {
        out.push_back(0x83);
        out.push_back(static_cast<uint8_t>(len >> 16));
        out.push_back(static_cast<uint8_t>((len >> 8) & 0xFF));
        out.push_back(static_cast<uint8_t>(len & 0xFF));
    }
}

static std::vector<uint8_t> derTag(uint8_t tag, const std::vector<uint8_t> &content) {
    std::vector<uint8_t> result;
    result.push_back(tag);
    derEncodeLength(result, content.size());
    result.insert(result.end(), content.begin(), content.end());
    return result;
}

static std::vector<uint8_t> derSequence(const std::vector<uint8_t> &content) {
    return derTag(0x30, content);
}

static std::vector<uint8_t> derSet(const std::vector<uint8_t> &content) {
    return derTag(0x31, content);
}

static std::vector<uint8_t> derInteger(const std::vector<uint8_t> &bytes) {
    std::vector<uint8_t> val = bytes;
    // Strip leading zero bytes, but keep at least 1
    while (val.size() > 1 && val[0] == 0x00 && (val[1] & 0x80) == 0)
        val.erase(val.begin());
    // Prepend 0x00 if high bit is set (positive integer)
    if (!val.empty() && (val[0] & 0x80))
        val.insert(val.begin(), 0x00);
    return derTag(0x02, val);
}

static std::vector<uint8_t> derIntegerSmall(int64_t val) {
    std::vector<uint8_t> bytes;
    if (val == 0) {
        bytes.push_back(0);
    } else {
        bool neg = val < 0;
        uint64_t uval = neg ? static_cast<uint64_t>(-val) : static_cast<uint64_t>(val);
        while (uval > 0) {
            bytes.insert(bytes.begin(), static_cast<uint8_t>(uval & 0xFF));
            uval >>= 8;
        }
        if (!neg && (bytes[0] & 0x80))
            bytes.insert(bytes.begin(), 0x00);
    }
    return derTag(0x02, bytes);
}

static std::vector<uint8_t> derOid(const std::vector<uint8_t> &oidBytes) {
    return derTag(0x06, oidBytes);
}

static std::vector<uint8_t> derNull() {
    return {0x05, 0x00};
}

static std::vector<uint8_t> derBitString(const std::vector<uint8_t> &bits) {
    std::vector<uint8_t> content;
    content.push_back(0x00); // unused bits = 0
    content.insert(content.end(), bits.begin(), bits.end());
    return derTag(0x03, content);
}

static std::vector<uint8_t> derUtf8String(const std::string &s) {
    std::vector<uint8_t> data(s.begin(), s.end());
    return derTag(0x0C, data);
}

static std::vector<uint8_t> derPrintableString(const std::string &s) {
    std::vector<uint8_t> data(s.begin(), s.end());
    return derTag(0x13, data);
}

static std::vector<uint8_t> derUtcTime(time_t t) {
    struct tm utc;
#ifdef _WIN32
    gmtime_s(&utc, &t);
#else
    gmtime_r(&t, &utc);
#endif
    char buf[64];
    snprintf(buf, sizeof(buf), "%02d%02d%02d%02d%02d%02dZ",
             utc.tm_year % 100, utc.tm_mon + 1, utc.tm_mday,
             utc.tm_hour, utc.tm_min, utc.tm_sec);
    std::vector<uint8_t> data(buf, buf + 13);
    return derTag(0x17, data);
}

static std::vector<uint8_t> derExplicit(uint8_t tagNum, const std::vector<uint8_t> &content) {
    uint8_t tag = 0xA0 | (tagNum & 0x1F);
    return derTag(tag, content);
}

// Concatenate multiple DER elements into a single buffer
static void derAppend(std::vector<uint8_t> &out, const std::vector<uint8_t> &elem) {
    out.insert(out.end(), elem.begin(), elem.end());
}

// ═════════════════════════════════════════════════════════════
//  Well-known OIDs
// ═════════════════════════════════════════════════════════════

// sha256WithRSAEncryption (1.2.840.113549.1.1.11)
static const uint8_t OID_SHA256_RSA[] = {0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x0B};

// rsaEncryption (1.2.840.113549.1.1.1)
static const uint8_t OID_RSA_ENC[] = {0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x01};

// id-at-commonName (2.5.4.3)
static const uint8_t OID_CN[] = {0x55, 0x04, 0x03};

// id-at-organizationName (2.5.4.10)
static const uint8_t OID_ORG[] = {0x55, 0x04, 0x0A};

// id-at-countryName (2.5.4.6)
static const uint8_t OID_COUNTRY[] = {0x55, 0x04, 0x06};

// ═════════════════════════════════════════════════════════════
//  RSA key DER export helpers
// ═════════════════════════════════════════════════════════════

static std::vector<uint8_t> bigIntBytes(const rsa::bigInt &bi) {
    size_t bits = bi.bitLength();
    size_t byteLen = (bits + 7) / 8;
    if (byteLen == 0) byteLen = 1;
    return rsa::bigIntToBytesBE(bi, byteLen);
}

// Build PKCS#1 RSAPrivateKey DER:
// RSAPrivateKey ::= SEQUENCE {
//   version           INTEGER (0),
//   modulus           INTEGER,
//   publicExponent    INTEGER,
//   privateExponent   INTEGER,
//   prime1            INTEGER,
//   prime2            INTEGER,
//   exponent1         INTEGER,  -- d mod (p-1)
//   exponent2         INTEGER,  -- d mod (q-1)
//   coefficient       INTEGER   -- q^-1 mod p
// }
static std::vector<uint8_t> buildRsaPrivateKeyDer(const rsa::bigInt &n,
                                                   const rsa::bigInt &e,
                                                   const rsa::bigInt &d,
                                                   const rsa::bigInt &p,
                                                   const rsa::bigInt &q) {
    // Compute CRT parameters
    rsa::bigInt one(1U, 1);
    rsa::bigInt pm1(p.capacity), qm1(q.capacity);
    rsa::bigInt dp_q, dp_r, dq_q, dq_r;
    rsa::subtract(p, one, pm1);    // p - 1
    rsa::subtract(q, one, qm1);    // q - 1
    rsa::divide(d, pm1, dp_q, dp_r);   // dp = d mod (p-1)
    rsa::divide(d, qm1, dq_q, dq_r);   // dq = d mod (q-1)
    rsa::bigInt qInv = rsa::modInverse(q, p);  // qInv = q^-1 mod p

    std::vector<uint8_t> body;
    derAppend(body, derIntegerSmall(0));                    // version = 0
    derAppend(body, derInteger(bigIntBytes(n)));            // modulus
    derAppend(body, derInteger(bigIntBytes(e)));            // publicExponent
    derAppend(body, derInteger(bigIntBytes(d)));            // privateExponent
    derAppend(body, derInteger(bigIntBytes(p)));            // prime1
    derAppend(body, derInteger(bigIntBytes(q)));            // prime2
    derAppend(body, derInteger(bigIntBytes(dp_r)));         // exponent1 (dp)
    derAppend(body, derInteger(bigIntBytes(dq_r)));         // exponent2 (dq)
    derAppend(body, derInteger(bigIntBytes(qInv)));         // coefficient
    return derSequence(body);
}

// Build SubjectPublicKeyInfo DER for RSA:
// SubjectPublicKeyInfo ::= SEQUENCE {
//   algorithm AlgorithmIdentifier,
//   subjectPublicKey BIT STRING containing RSAPublicKey
// }
// RSAPublicKey ::= SEQUENCE { modulus INTEGER, publicExponent INTEGER }
static std::vector<uint8_t> buildRsaPublicKeyInfo(const rsa::bigInt &n,
                                                   const rsa::bigInt &e) {
    // RSAPublicKey SEQUENCE
    std::vector<uint8_t> pubKeyBody;
    derAppend(pubKeyBody, derInteger(bigIntBytes(n)));
    derAppend(pubKeyBody, derInteger(bigIntBytes(e)));
    auto rsaPubKey = derSequence(pubKeyBody);

    // AlgorithmIdentifier SEQUENCE { OID rsaEncryption, NULL }
    std::vector<uint8_t> algBody;
    derAppend(algBody, derOid({OID_RSA_ENC, OID_RSA_ENC + sizeof(OID_RSA_ENC)}));
    derAppend(algBody, derNull());
    auto algId = derSequence(algBody);

    // SubjectPublicKeyInfo
    std::vector<uint8_t> spkiBody;
    derAppend(spkiBody, algId);
    derAppend(spkiBody, derBitString(rsaPubKey));
    return derSequence(spkiBody);
}

// ═════════════════════════════════════════════════════════════
//  X.509 self-signed certificate builder
// ═════════════════════════════════════════════════════════════

// Build a DN (distinguished name) from config
static std::vector<uint8_t> buildDN(const CertGenConfig &cfg) {
    std::vector<uint8_t> dnBody;

    // CN
    if (!cfg.commonName.empty()) {
        std::vector<uint8_t> atv;
        derAppend(atv, derOid({OID_CN, OID_CN + sizeof(OID_CN)}));
        derAppend(atv, derUtf8String(cfg.commonName));
        derAppend(dnBody, derSet(derSequence(atv)));
    }

    // O
    if (!cfg.organization.empty()) {
        std::vector<uint8_t> atv;
        derAppend(atv, derOid({OID_ORG, OID_ORG + sizeof(OID_ORG)}));
        derAppend(atv, derUtf8String(cfg.organization));
        derAppend(dnBody, derSet(derSequence(atv)));
    }

    // C
    if (!cfg.country.empty()) {
        std::vector<uint8_t> atv;
        derAppend(atv, derOid({OID_COUNTRY, OID_COUNTRY + sizeof(OID_COUNTRY)}));
        derAppend(atv, derPrintableString(cfg.country));
        derAppend(dnBody, derSet(derSequence(atv)));
    }

    return derSequence(dnBody);
}

// AlgorithmIdentifier for sha256WithRSAEncryption
static std::vector<uint8_t> signatureAlgId() {
    std::vector<uint8_t> body;
    derAppend(body, derOid({OID_SHA256_RSA, OID_SHA256_RSA + sizeof(OID_SHA256_RSA)}));
    derAppend(body, derNull());
    return derSequence(body);
}

// ═════════════════════════════════════════════════════════════
//  PEM encoding
// ═════════════════════════════════════════════════════════════

static std::string derToPem(const std::vector<uint8_t> &der, const std::string &label) {
    // Base64 encode the DER data
    base64 b64;
    std::string rawStr(reinterpret_cast<const char *>(der.data()), der.size());
    b64 << rawStr;

    // Get the encoded data from the base64 object
    const auto &enc = b64.encoded();
    std::string encoded(enc.begin(), enc.end());

    // Format as PEM with 64-char lines
    std::string pem = "-----BEGIN " + label + "-----\n";
    for (size_t i = 0; i < encoded.size(); i += 64) {
        size_t len = std::min<size_t>(64, encoded.size() - i);
        pem.append(encoded, i, len);
        pem += '\n';
    }
    pem += "-----END " + label + "-----\n";
    return pem;
}

// ═════════════════════════════════════════════════════════════
//  Public API
// ═════════════════════════════════════════════════════════════

bool generateSelfSignedCert(const CertGenConfig &cfg, CertKeyPair &out) {
    // 1. Generate RSA key pair
    rsa key;
    key.generateKeys(cfg.rsaBits);

    // Access private members through friend (certgen is friended in rsa.h)
    const auto &n = key.getN();
    const auto &e = key.getE();
    const auto &d = key.getD();

    if (n.isZero() || e.isZero() || d.isZero())
        return false;

    const auto &p = key.getP();
    const auto &q = key.getQ();

    // 2. Build key DER
    out.keyDer = buildRsaPrivateKeyDer(n, e, d, p, q);

    // 3. Build TBSCertificate
    auto dn = buildDN(cfg);
    auto spki = buildRsaPublicKeyInfo(n, e);

    time_t now = time(nullptr);
    time_t expire = now + static_cast<time_t>(cfg.validDays) * 86400;

    // Serial number: use first 8 bytes of SHA-256(current time + pubkey)
    std::vector<uint8_t> serialInput;
    {
        auto tBytes = reinterpret_cast<const uint8_t *>(&now);
        serialInput.insert(serialInput.end(), tBytes, tBytes + sizeof(now));
        auto nBytes = bigIntBytes(n);
        serialInput.insert(serialInput.end(), nBytes.begin(), nBytes.end());
    }
    auto serialHash = sha256_hash(serialInput);
    serialHash.resize(8);
    serialHash[0] &= 0x7F; // ensure positive

    // Validity SEQUENCE { notBefore UTCTime, notAfter UTCTime }
    std::vector<uint8_t> validityBody;
    derAppend(validityBody, derUtcTime(now));
    derAppend(validityBody, derUtcTime(expire));
    auto validity = derSequence(validityBody);

    // TBSCertificate SEQUENCE
    std::vector<uint8_t> tbsBody;
    derAppend(tbsBody, derExplicit(0, derIntegerSmall(2)));  // version v3 (explicit [0])
    derAppend(tbsBody, derInteger(serialHash));              // serialNumber
    derAppend(tbsBody, signatureAlgId());                    // signature algorithm
    derAppend(tbsBody, dn);                                  // issuer (self-signed = subject)
    derAppend(tbsBody, validity);                            // validity
    derAppend(tbsBody, dn);                                  // subject
    derAppend(tbsBody, spki);                                // subjectPublicKeyInfo

    auto tbsCert = derSequence(tbsBody);

    // 4. Sign TBSCertificate with RSASSA-PKCS1-v1_5 SHA-256
    auto signature = rsa_pkcs1_sha256::sign(key, tbsCert);
    if (signature.empty())
        return false;

    // 5. Build full Certificate SEQUENCE
    std::vector<uint8_t> certBody;
    derAppend(certBody, tbsCert);
    derAppend(certBody, signatureAlgId());
    derAppend(certBody, derBitString(signature));

    out.certDer = derSequence(certBody);

    // 6. PEM encode
    out.certPem = derToPem(out.certDer, "CERTIFICATE");
    out.keyPem  = derToPem(out.keyDer, "RSA PRIVATE KEY");

    return true;
}

bool writePemFile(const std::string &path, const std::string &pem) {
    std::ofstream ofs(path, std::ios::binary);
    if (!ofs) return false;
    ofs.write(pem.data(), static_cast<std::streamsize>(pem.size()));
    return ofs.good();
}

bool writeDerFile(const std::string &path, const std::vector<uint8_t> &der) {
    std::ofstream ofs(path, std::ios::binary);
    if (!ofs) return false;
    ofs.write(reinterpret_cast<const char *>(der.data()),
              static_cast<std::streamsize>(der.size()));
    return ofs.good();
}

} // namespace netplus
