/*******************************************************************************
 * Copyright (c) 2025, 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 <fstream>
#include <sstream>
#include <ctime>
#include <chrono>
#include <cctype>
#include <string>
#include <vector>
#include <iostream>
#include <stdexcept>
#include <memory>

#include "base64.h"
#include "x509.h"

using namespace netplus;

// ------------------------------------------------------------
// ctor
// ------------------------------------------------------------
netplus::x509cert::x509cert()
    : state(std::make_shared<X509State>()) {}

// ------------------------------------------------------------
// parseASN1Time (ptr/len)  -- Zero-copy
// ------------------------------------------------------------
time_t netplus::x509cert::parseASN1Time(const uint8_t* data, size_t len, uint8_t tag) {
    if (!data || len == 0) return 0;

    auto parse2 = [] (const char* s, size_t p, size_t n, int& v) {
        if (p + 2 > n) return false;
        if (!isdigit((unsigned char)s[p]) || !isdigit((unsigned char)s[p+1])) return false;
        v = (s[p]-'0')*10 + (s[p+1]-'0');
        return true;
    };

    auto parse4 = [] (const char* s, size_t p, size_t n, int& v) {
        if (p + 4 > n) return false;
        for (size_t i = 0; i < 4; ++i)
            if (!isdigit((unsigned char)s[p+i])) return false;
        v = (s[p]-'0')*1000 + (s[p+1]-'0')*100 + (s[p+2]-'0')*10 + (s[p+3]-'0');
        return true;
    };

    auto tm_to_utc_time_t = [](std::tm& t) {
#if defined(_WIN32)
        return _mkgmtime(&t);
#else
        return timegm(&t);
#endif
    };

    // Trim whitespace
    while (len > 0 && (data[len-1] == '\r' || data[len-1] == '\n' || data[len-1] == ' ' || data[len-1] == '\t'))
        --len;

    if (len == 0) return 0;

    // Must end with 'Z'
    if (data[len-1] != 'Z') return 0;
    --len;

    const char* s = reinterpret_cast<const char*>(data);

    int year = 0, mon = 0, day = 0, hour = 0, min = 0, sec = 0;

    if (tag == 0x17) {
        // UTCTime: YYMMDDHHMM[SS]
        if (len != 10 && len != 12) return 0;

        int yy = 0;
        if (!parse2(s, 0, len, yy)) return 0;
        if (!parse2(s, 2, len, mon)) return 0;
        if (!parse2(s, 4, len, day)) return 0;
        if (!parse2(s, 6, len, hour)) return 0;
        if (!parse2(s, 8, len, min)) return 0;

        if (len == 12) {
            if (!parse2(s, 10, len, sec)) return 0;
        } else {
            sec = 0;
        }

        // RFC 5280 mapping
        year = (yy <= 49) ? (2000 + yy) : (1900 + yy);

    } else if (tag == 0x18) {
        // GeneralizedTime: YYYYMMDDHHMM[SS]
        if (len != 12 && len != 14) return 0;

        if (!parse4(s, 0, len, year)) return 0;
        if (!parse2(s, 4, len, mon)) return 0;
        if (!parse2(s, 6, len, day)) return 0;
        if (!parse2(s, 8, len, hour)) return 0;
        if (!parse2(s, 10, len, min)) return 0;

        if (len == 14) {
            if (!parse2(s, 12, len, sec)) return 0;
        } else {
            sec = 0;
        }

    } else {
        return 0;
    }

    // Basic checks
    if (mon < 1 || mon > 12) return 0;
    if (day < 1 || day > 31) return 0;
    if (hour < 0 || hour > 23) return 0;
    if (min < 0 || min > 59) return 0;
    if (sec < 0 || sec > 60) return 0;

    std::tm t{};
    t.tm_year = year - 1900;
    t.tm_mon  = mon - 1;
    t.tm_mday = day;
    t.tm_hour = hour;
    t.tm_min  = min;
    t.tm_sec  = sec;
    t.tm_isdst = 0;

    return tm_to_utc_time_t(t);
}

// ------------------------------------------------------------
// decodeOID (ptr/len) -- Zero-copy
// ------------------------------------------------------------
std::string netplus::x509cert::decodeOID(const uint8_t* data, size_t len) {
    if (!data || len == 0) return "";

    std::string out;
    out.reserve(len * 3);

    auto appendNumber = [&](uint32_t v) {
        out += std::to_string(v);
    };

    // first byte = first two arcs
    uint8_t first = data[0];
    appendNumber(first / 40);
    out.push_back('.');
    appendNumber(first % 40);

    size_t i = 1;
    while (i < len) {
        uint32_t value = 0;
        uint8_t byte = 0;

        do {
            if (i >= len) break;
            byte = data[i++];
            value = (value << 7) | (byte & 0x7F);
        } while (byte & 0x80);

        out.push_back('.');
        appendNumber(value);
    }

    return out;
}

// ------------------------------------------------------------
// parseInternal -- Zero-copy tree builder
// ------------------------------------------------------------
size_t netplus::x509cert::parseInternal(const uint8_t* buf, size_t size, ASN1Node& outNode) {
    if (!buf || size < 2) return 0;

    size_t cursor = 0;
    outNode.tag = buf[cursor++];

    if (cursor >= size) return 0;
    size_t length = buf[cursor++];

    if (length & 0x80) {
        int octets = static_cast<int>(length & 0x7F);
        if (octets <= 0 || cursor + static_cast<size_t>(octets) > size) return 0;

        length = 0;
        for (int i = 0; i < octets; ++i) {
            length = (length << 8) | buf[cursor++];
        }
    }

    if (cursor + length > size) return 0;

    // ✅ Zero-copy view
    outNode.data = buf + cursor;
    outNode.len  = length;

    if (outNode.tag & 0x20) {
        outNode.children.clear();
        size_t sub = 0;

        while (sub < length) {
            ASN1Node child;
            size_t n = parseInternal(outNode.data + sub, length - sub, child);
            if (n == 0) break;
            outNode.children.push_back(std::move(child));
            sub += n;
        }
    } else {
        outNode.children.clear();
    }

    return cursor + length;
}

// ------------------------------------------------------------
// loadFromBuffer -- parse once, share state
// ------------------------------------------------------------
bool netplus::x509cert::loadFromBuffer(std::vector<uint8_t>&& buffer) {
    if (buffer.empty()) return false;

    state = std::make_shared<X509State>();
    state->der = std::move(buffer);
    state->root = ASN1Node();

    size_t bytesRead = parseInternal(state->der.data(), state->der.size(), state->root);
    return bytesRead > 0;
}

bool netplus::x509cert::loadFromBuffer(const std::vector<uint8_t>& buffer) {
    return loadFromBuffer(std::vector<uint8_t>(buffer));
}

// ------------------------------------------------------------
// loadFromFile -- move buffer, no copy
// ------------------------------------------------------------
bool netplus::x509cert::loadFromFile(const std::string& filePath) {
    std::ifstream file(filePath, std::ios::binary | std::ios::ate);
    if (!file.is_open()) return false;

    std::streamsize size = file.tellg();
    if (size <= 0 || size < 2) return false;

    file.seekg(0, std::ios::beg);
    std::vector<uint8_t> buffer(static_cast<size_t>(size));

    if (!file.read(reinterpret_cast<char*>(buffer.data()), size)) return false;

    // DER must start with SEQUENCE
    if (buffer[0] != 0x30) return false;

    return loadFromBuffer(std::move(buffer));
}

// ------------------------------------------------------------
// checkValidity
// ------------------------------------------------------------
bool netplus::x509cert::checkValidity() {
    auto& r = root();
    if (r.children.empty() || r.children[0].children.size() < 5)
        return false;

    ASN1Node* validity = nullptr;
    auto& tbs = r.children[0];

    // Validity = SEQUENCE( notBefore, notAfter )
    for (auto& node : tbs.children) {
        if (node.tag == 0x30 && node.children.size() >= 2) {
            uint8_t t1 = node.children[0].tag;
            uint8_t t2 = node.children[1].tag;

            bool ok1 = (t1 == 0x17 || t1 == 0x18);
            bool ok2 = (t2 == 0x17 || t2 == 0x18);

            if (ok1 && ok2) {
                validity = &node;
                break;
            }
        }
    }

    if (!validity) return false;

    const ASN1Node& nb = validity->children[0];
    const ASN1Node& na = validity->children[1];

    time_t notBefore = parseASN1Time(nb.data, nb.len, nb.tag);
    time_t notAfter  = parseASN1Time(na.data, na.len, na.tag);

    if (!notBefore || !notAfter) return false;

    time_t now = time(nullptr);

    if (now < notBefore) {
        std::cout << "Certificate is not yet valid (Future date)." << std::endl;
        return false;
    }
    if (now > notAfter) {
        std::cout << "Certificate has EXPIRED." << std::endl;
        return false;
    }

    std::cout << "Certificate validity: OK." << std::endl;
    return true;
}

// ------------------------------------------------------------
// bytesToBigInt -- vector-based (keeps header compat)
// ------------------------------------------------------------
void netplus::x509cert::bytesToBigInt(const std::vector<uint8_t>& bytes, netplus::rsa::bigInt& out) {
    out.reserve((bytes.size() / 4) + 1);
    out.used = 1;
    out.data[0] = 0;

    for (uint8_t b : bytes) {
        out.shiftLeft(8);
        out.data[0] |= b;
    }
}

// ------------------------------------------------------------
// extractPublicKey -- RSA SubjectPublicKeyInfo
// ------------------------------------------------------------
bool netplus::x509cert::extractPublicKey(netplus::rsa& outRsa) {
    auto& r = root();
    if (r.children.empty()) return false;

    auto& tbs = r.children[0];

    ASN1Node* pubKeyInfo = nullptr;

    // Find SubjectPublicKeyInfo with RSA OID
    for (auto& child : tbs.children) {
        if (child.tag == 0x30 && child.children.size() >= 2) {
            if (child.children[0].tag == 0x30 && !child.children[0].children.empty()) {
                const ASN1Node& oidNode = child.children[0].children[0];
                std::string oid = decodeOID(oidNode.data, oidNode.len);

                if (oid == "1.2.840.113549.1.1.1") { // rsaEncryption
                    pubKeyInfo = &child;
                    break;
                }
            }
        }
    }

    if (!pubKeyInfo) return false;

    // BIT STRING contains RSAPublicKey
    for (auto& node : pubKeyInfo->children) {
        if (node.tag != 0x03) continue;
        if (!node.data || node.len < 2) continue;

        ASN1Node rsaKey;

        // first byte = unused bits count
        if (parseInternal(node.data + 1, node.len - 1, rsaKey) == 0)
            continue;

        if (rsaKey.children.size() < 2) continue;

        // modulus (INTEGER)
        const uint8_t* mPtr = rsaKey.children[0].data;
        size_t mLen = rsaKey.children[0].len;
        if (!mPtr || mLen == 0) continue;

        // Remove leading 0x00 if present
        if (mLen > 1 && mPtr[0] == 0x00) { mPtr++; mLen--; }

        // exponent
        const uint8_t* ePtr = rsaKey.children[1].data;
        size_t eLen = rsaKey.children[1].len;
        if (!ePtr || eLen == 0) continue;

        // Build vectors only here (rare, not in parse hot path)
        std::vector<uint8_t> mBytes(mPtr, mPtr + mLen);
        std::vector<uint8_t> eBytes(ePtr, ePtr + eLen);

        bytesToBigInt(mBytes, outRsa.n);
        bytesToBigInt(eBytes, outRsa.e);
        return true;
    }

    return false;
}

// ------------------------------------------------------------
// getSubjectCN
// ------------------------------------------------------------
std::string netplus::x509cert::getSubjectCN() {
    auto& r = root();

    // TBSCertificate -> Subject (index 5 usually, but keep your current logic)
    if (r.children.empty() || r.children[0].children.size() < 6)
        return "";

    auto& subject = r.children[0].children[5];

    for (auto& rdnSet : subject.children) {
        for (auto& attrSeq : rdnSet.children) {
            if (attrSeq.children.size() < 2) continue;

            const ASN1Node& oidNode = attrSeq.children[0];
            std::string oid = decodeOID(oidNode.data, oidNode.len);

            if (oid == "2.5.4.3") { // CN
                const ASN1Node& valueNode = attrSeq.children[1];
                if (!valueNode.data || valueNode.len == 0) return "";
                return std::string(
                    reinterpret_cast<const char*>(valueNode.data),
                    reinterpret_cast<const char*>(valueNode.data) + valueNode.len
                );
            }
        }
    }

    return "";
}

// ------------------------------------------------------------
// getSubjectAltNames (SAN)
// ------------------------------------------------------------
std::vector<std::string> netplus::x509cert::getSubjectAltNames() {
    std::vector<std::string> dnsNames;

    auto& r = root();
    if (r.children.empty()) return dnsNames;

    auto& tbs = r.children[0];

    // Find Extensions wrapper ([3] EXPLICIT) tag A3
    ASN1Node* extWrapper = nullptr;
    for (auto& child : tbs.children) {
        if (child.tag == 0xA3) {
            extWrapper = &child;
            break;
        }
    }
    if (!extWrapper) return dnsNames;

    // Get Extensions sequence inside wrapper
    const ASN1Node* extensionsSeq = nullptr;
    ASN1Node parsedWrapper;

    if (!extWrapper->children.empty()) {
        extensionsSeq = &extWrapper->children[0];
    } else {
        if (!extWrapper->data || extWrapper->len == 0) return dnsNames;
        if (parseInternal(extWrapper->data, extWrapper->len, parsedWrapper) == 0) return dnsNames;
        if (parsedWrapper.children.empty()) return dnsNames;
        extensionsSeq = &parsedWrapper.children[0];
    }

    if (!extensionsSeq) return dnsNames;

    // Iterate extensions to find SAN OID 2.5.29.17
    for (const auto& ext : extensionsSeq->children) {
        if (ext.children.empty()) continue;

        const ASN1Node& oidNode = ext.children[0];
        std::string oid = decodeOID(oidNode.data, oidNode.len);
        if (oid != "2.5.29.17") continue;

        // extnValue is OCTET STRING (04) containing GeneralNames DER
        for (const auto& valNode : ext.children) {
            if (valNode.tag != 0x04) continue;
            if (!valNode.data || valNode.len == 0) continue;

            ASN1Node sanData;
            if (parseInternal(valNode.data, valNode.len, sanData) == 0) continue;

            // GeneralNames = SEQUENCE OF GeneralName
            for (const auto& name : sanData.children) {
                if (name.tag == 0x82 && name.data && name.len > 0) { // dNSName
                    dnsNames.emplace_back(
                        reinterpret_cast<const char*>(name.data),
                        reinterpret_cast<const char*>(name.data) + name.len
                    );
                }
            }
        }

        // SAN is normally unique -> stop here
        break;
    }

    return dnsNames;
}

// ------------------------------------------------------------
// PEM cert -> DER
// ------------------------------------------------------------
std::vector<uint8_t> netplus::x509cert::pemCertToDer(const std::string& pem) {
    const std::string begin = "-----BEGIN CERTIFICATE-----";
    const std::string end   = "-----END CERTIFICATE-----";

    auto b = pem.find(begin);
    if (b == std::string::npos)
        throw std::runtime_error("PEM: missing BEGIN");
    b += begin.size();

    auto e = pem.find(end, b);
    if (e == std::string::npos)
        throw std::runtime_error("PEM: missing END");

    std::string body = pem.substr(b, e - b);

    std::string b64;
    b64.reserve(body.size());
    for (unsigned char c : body) {
        if (std::isalnum(c) || c == '+' || c == '/' || c == '=')
            b64.push_back(static_cast<char>(c));
    }

    netplus::base64 dec;
    dec << b64;

    std::vector<char> decoded;
    dec >> decoded;

    std::vector<uint8_t> der;
    der.reserve(decoded.size());
    for (unsigned char c : decoded)
        der.push_back(static_cast<uint8_t>(c));

    return der;
}
