/*
 * tls.cpp
 * Shared TLS helper library for MontaukOS
 * Extracted from fetch, wiki, wikipedia, and weather apps
 * Copyright (c) 2025-2026 Daniel Hammer
 */

#include <tls/tls.hpp>
#include <montauk/syscall.h>

extern "C" {
#include <string.h>
#include <stdlib.h>
}

// ============================================================================
// File-local: PEM/DER certificate parsing helpers
// ============================================================================

namespace {

struct DerAccum { unsigned char* data; size_t len, cap; bool failed; };
struct DnAccum  { unsigned char* data; size_t len, cap; bool failed; };

void der_append(void* ctx, const void* buf, size_t len) {
    DerAccum* a = (DerAccum*)ctx;
    if (a->len + len > a->cap) {
        size_t nc = a->cap * 2;
        if (nc < a->len + len) nc = a->len + len + 4096;
        unsigned char* nb = (unsigned char*)malloc(nc);
        if (!nb) { a->failed = true; return; }
        if (a->data) { memcpy(nb, a->data, a->len); free(a->data); }
        a->data = nb; a->cap = nc;
    }
    memcpy(a->data + a->len, buf, len);
    a->len += len;
}

void dn_append(void* ctx, const void* buf, size_t len) {
    DnAccum* a = (DnAccum*)ctx;
    if (a->len + len > a->cap) {
        size_t nc = a->cap * 2;
        if (nc < a->len + len) nc = a->len + len + 256;
        unsigned char* nb = (unsigned char*)malloc(nc);
        if (!nb) { a->failed = true; return; }
        if (a->data) { memcpy(nb, a->data, a->len); free(a->data); }
        a->data = nb; a->cap = nc;
    }
    memcpy(a->data + a->len, buf, len);
    a->len += len;
}

bool ta_add(tls::TrustAnchors* tas, const br_x509_trust_anchor* ta) {
    if (tas->count >= tas->capacity) {
        size_t nc = tas->capacity == 0 ? 64 : tas->capacity * 2;
        br_x509_trust_anchor* na = (br_x509_trust_anchor*)malloc(nc * sizeof(*na));
        if (!na) return false;
        if (tas->anchors) { memcpy(na, tas->anchors, tas->count * sizeof(*na)); free(tas->anchors); }
        tas->anchors = na; tas->capacity = nc;
    }
    tas->anchors[tas->count++] = *ta;
    return true;
}

bool process_cert_der(tls::TrustAnchors* tas, const unsigned char* der, size_t der_len) {
    br_x509_decoder_context* dc =
        (br_x509_decoder_context*)malloc(sizeof(br_x509_decoder_context));
    if (!dc) return false;
    DnAccum dn = {nullptr, 0, 0, false};
    br_x509_decoder_init(dc, dn_append, &dn);
    br_x509_decoder_push(dc, der, der_len);
    br_x509_pkey* pk = br_x509_decoder_get_pkey(dc);
    if (!pk || dn.failed) { free(dc); if (dn.data) free(dn.data); return false; }

    br_x509_trust_anchor ta;
    memset(&ta, 0, sizeof(ta));
    ta.dn.data = dn.data; ta.dn.len = dn.len; ta.flags = 0;
    if (br_x509_decoder_isCA(dc)) ta.flags |= BR_X509_TA_CA;

    switch (pk->key_type) {
    case BR_KEYTYPE_RSA:
        ta.pkey.key_type = BR_KEYTYPE_RSA;
        ta.pkey.key.rsa.nlen = pk->key.rsa.nlen;
        ta.pkey.key.rsa.n = (unsigned char*)malloc(pk->key.rsa.nlen);
        if (!ta.pkey.key.rsa.n) { free(dc); free(dn.data); return false; }
        memcpy(ta.pkey.key.rsa.n, pk->key.rsa.n, pk->key.rsa.nlen);
        ta.pkey.key.rsa.elen = pk->key.rsa.elen;
        ta.pkey.key.rsa.e = (unsigned char*)malloc(pk->key.rsa.elen);
        if (!ta.pkey.key.rsa.e) {
            free(ta.pkey.key.rsa.n);
            free(dc);
            free(dn.data);
            return false;
        }
        memcpy(ta.pkey.key.rsa.e, pk->key.rsa.e, pk->key.rsa.elen);
        break;
    case BR_KEYTYPE_EC:
        ta.pkey.key_type = BR_KEYTYPE_EC;
        ta.pkey.key.ec.curve = pk->key.ec.curve;
        ta.pkey.key.ec.qlen = pk->key.ec.qlen;
        ta.pkey.key.ec.q = (unsigned char*)malloc(pk->key.ec.qlen);
        if (!ta.pkey.key.ec.q) { free(dc); free(dn.data); return false; }
        memcpy(ta.pkey.key.ec.q, pk->key.ec.q, pk->key.ec.qlen);
        break;
    default:
        free(dc);
        if (dn.data) free(dn.data);
        return false;
    }
    free(dc);
    if (ta_add(tas, &ta)) return true;
    free(ta.dn.data);
    if (ta.pkey.key_type == BR_KEYTYPE_RSA) {
        free(ta.pkey.key.rsa.n);
        free(ta.pkey.key.rsa.e);
    } else {
        free(ta.pkey.key.ec.q);
    }
    return false;
}

} // anonymous namespace

// ============================================================================
// Exported functions
// ============================================================================

namespace tls {

TrustAnchors load_trust_anchors() {
    TrustAnchors tas = {nullptr, 0, 0};
    int fh = montauk::open("0:/os/certs/ca-certificates.crt");
    if (fh < 0) return tas;
    uint64_t fsize = montauk::getsize(fh);
    if (fsize == 0 || fsize > 512 * 1024) { montauk::close(fh); return tas; }

    unsigned char* pem = (unsigned char*)malloc(fsize + 1);
    if (!pem) { montauk::close(fh); return tas; }
    uint64_t readOffset = 0;
    while (readOffset < fsize) {
        int n = montauk::read(fh, pem + readOffset, readOffset, fsize - readOffset);
        if (n <= 0) break;
        readOffset += (uint64_t)n;
    }
    montauk::close(fh);
    if (readOffset != fsize) { free(pem); return tas; }
    pem[readOffset] = 0;

    br_pem_decoder_context* pc =
        (br_pem_decoder_context*)malloc(sizeof(br_pem_decoder_context));
    if (!pc) { free(pem); return tas; }
    br_pem_decoder_init(pc);
    DerAccum der = {nullptr, 0, 0, false};
    bool inCert = false;
    size_t offset = 0;

    while (offset < fsize) {
        size_t pushed = br_pem_decoder_push(pc, pem + offset, fsize - offset);
        offset += pushed;
        int ev = br_pem_decoder_event(pc);
        if (ev == BR_PEM_BEGIN_OBJ) {
            inCert = (strcmp(br_pem_decoder_name(pc), "CERTIFICATE") == 0);
            br_pem_decoder_setdest(pc, inCert ? der_append : nullptr, inCert ? &der : nullptr);
            if (inCert) { der.len = 0; der.failed = false; }
        } else if (ev == BR_PEM_END_OBJ) {
            if (inCert && der.len > 0) process_cert_der(&tas, der.data, der.len);
            inCert = false;
        } else if (ev == BR_PEM_ERROR) {
            break;
        }
        if (der.failed) break;
        if (pushed == 0 && ev == 0) break;
    }
    if (der.data) free(der.data);
    free(pc);
    free(pem);
    return tas;
}

void free_trust_anchors(TrustAnchors* tas) {
    if (!tas) return;
    for (size_t i = 0; i < tas->count; ++i) {
        br_x509_trust_anchor& ta = tas->anchors[i];
        free(ta.dn.data);
        if (ta.pkey.key_type == BR_KEYTYPE_RSA) {
            free(ta.pkey.key.rsa.n);
            free(ta.pkey.key.rsa.e);
        } else if (ta.pkey.key_type == BR_KEYTYPE_EC) {
            free(ta.pkey.key.ec.q);
        }
    }
    free(tas->anchors);
    tas->anchors = nullptr;
    tas->count = 0;
    tas->capacity = 0;
}

void get_bearssl_time(uint32_t* days, uint32_t* seconds) {
    montauk::abi::DateTime dt;
    montauk::gettime(&dt);
    int y = dt.Year, m = dt.Month, d = dt.Day;
    uint32_t total = 365u * (uint32_t)y
        + (uint32_t)((y+3)/4) - (uint32_t)((y+99)/100) + (uint32_t)((y+399)/400);
    const int md[] = {0,31,28,31,30,31,30,31,31,30,31,30,31};
    for (int mo = 1; mo < m && mo <= 12; mo++) total += md[mo];
    if (y%4==0 && (y%100!=0 || y%400==0) && m > 2) total++;
    total += d - 1;
    *days = total;
    *seconds = (uint32_t)(dt.Hour*3600 + dt.Minute*60 + dt.Second);
}

int tls_send_all(int fd, const unsigned char* data, size_t len) {
    size_t sent = 0;
    uint64_t deadline = montauk::get_milliseconds() + 15000;
    while (sent < len) {
        int r = montauk::send(fd, data + sent, (uint32_t)(len - sent));
        if (r > 0) { sent += r; deadline = montauk::get_milliseconds() + 15000; }
        else if (r < 0) return -1;
        else {
            uint64_t now = montauk::get_milliseconds();
            if (now >= deadline) return -1;
            uint32_t signals = montauk::wait_handle(
                fd, montauk::abi::IPC_SIGNAL_WRITABLE |
                    montauk::abi::IPC_SIGNAL_PEER_CLOSED, deadline - now);
            if (signals == (uint32_t)-1 ||
                (signals & montauk::abi::IPC_SIGNAL_PEER_CLOSED)) return -1;
        }
    }
    return (int)sent;
}

int tls_recv_some(int fd, unsigned char* buf, size_t maxlen) {
    uint64_t deadline = montauk::get_milliseconds() + 15000;
    while (true) {
        int r = montauk::recv(fd, buf, (uint32_t)maxlen);
        if (r > 0) return r;
        if (r < 0) return -1;
        uint64_t now = montauk::get_milliseconds();
        if (now >= deadline) return -1;
        uint32_t signals = montauk::wait_handle(
            fd, montauk::abi::IPC_SIGNAL_READABLE |
                montauk::abi::IPC_SIGNAL_PEER_CLOSED, deadline - now);
        if (signals == (uint32_t)-1) return -1;
        if ((signals & montauk::abi::IPC_SIGNAL_PEER_CLOSED) &&
            !(signals & montauk::abi::IPC_SIGNAL_READABLE)) return -1;
    }
}

int tls_exchange(int fd, br_ssl_engine_context* eng,
                 const char* request, int reqLen,
                 char* respBuf, int respMax,
                 AbortCheckFn abort_check) {
    int requestOffset = 0;
    int respLen = 0;
    uint64_t deadline = montauk::get_milliseconds() + 30000;

    while (true) {
        unsigned state = br_ssl_engine_current_state(eng);
        if (state & BR_SSL_CLOSED) {
            int err = br_ssl_engine_last_error(eng);
            if (err != BR_ERR_OK && err != BR_ERR_IO && respLen == 0) return -1;
            return respLen;
        }
        if (abort_check && abort_check()) {
            br_ssl_engine_close(eng);
            return respLen > 0 ? respLen : -1;
        }
        if (state & BR_SSL_SENDREC) {
            size_t len; unsigned char* buf = br_ssl_engine_sendrec_buf(eng, &len);
            int sent = tls_send_all(fd, buf, len);
            if (sent < 0) { br_ssl_engine_close(eng); return respLen > 0 ? respLen : -1; }
            br_ssl_engine_sendrec_ack(eng, len);
            deadline = montauk::get_milliseconds() + 30000; continue;
        }
        if (state & BR_SSL_RECVAPP) {
            size_t len; unsigned char* buf = br_ssl_engine_recvapp_buf(eng, &len);
            size_t toCopy = len;
            if (respLen + (int)toCopy > respMax - 1) toCopy = respMax - 1 - respLen;
            if (toCopy > 0) { memcpy(respBuf + respLen, buf, toCopy); respLen += toCopy; }
            br_ssl_engine_recvapp_ack(eng, len);
            deadline = montauk::get_milliseconds() + 30000; continue;
        }
        if ((state & BR_SSL_SENDAPP) && requestOffset < reqLen) {
            size_t len; unsigned char* buf = br_ssl_engine_sendapp_buf(eng, &len);
            size_t toWrite = (size_t)(reqLen - requestOffset);
            if (toWrite > len) toWrite = len;
            memcpy(buf, request + requestOffset, toWrite);
            br_ssl_engine_sendapp_ack(eng, toWrite);
            requestOffset += (int)toWrite;
            if (requestOffset == reqLen) br_ssl_engine_flush(eng, 0);
            deadline = montauk::get_milliseconds() + 30000; continue;
        }
        if (state & BR_SSL_RECVREC) {
            size_t len; unsigned char* buf = br_ssl_engine_recvrec_buf(eng, &len);
            int got = tls_recv_some(fd, buf, len);
            if (got < 0) { br_ssl_engine_close(eng); return respLen > 0 ? respLen : -1; }
            br_ssl_engine_recvrec_ack(eng, got);
            deadline = montauk::get_milliseconds() + 30000; continue;
        }
        if (montauk::get_milliseconds() >= deadline) return respLen > 0 ? respLen : -1;
        montauk::sleep_ms(1);
    }
}

int https_fetch(const char* host, uint32_t ip, uint16_t port,
                const char* request, int reqLen,
                const TrustAnchors& tas,
                char* respBuf, int respMax,
                AbortCheckFn abort_check) {
    if (!host || !*host || ip == 0 || port == 0 || !request || reqLen <= 0 ||
        !respBuf || respMax < 2 || !tas.anchors || tas.count == 0) return -1;
    int fd = montauk::socket(montauk::abi::SOCK_TCP);
    if (fd < 0) return -1;
    if (montauk::connect(fd, ip, port) < 0) { montauk::closesocket(fd); return -1; }

    br_ssl_client_context* cc  = (br_ssl_client_context*)malloc(sizeof(*cc));
    br_x509_minimal_context* xc = (br_x509_minimal_context*)malloc(sizeof(*xc));
    void* iobuf = malloc(BR_SSL_BUFSIZE_BIDI);
    if (!cc || !xc || !iobuf) {
        free(cc); free(xc); free(iobuf);
        montauk::closesocket(fd); return -1;
    }

    br_ssl_client_init_full(cc, xc, tas.anchors, tas.count);
    uint32_t days, secs;
    get_bearssl_time(&days, &secs);
    br_x509_minimal_set_time(xc, days, secs);

    unsigned char seed[32];
    if (montauk::getrandom(seed, sizeof(seed)) != (int64_t)sizeof(seed)) {
        montauk::closesocket(fd); free(cc); free(xc); free(iobuf); return -1;
    }
    br_ssl_engine_set_buffer(&cc->eng, iobuf, BR_SSL_BUFSIZE_BIDI, 1);
    br_ssl_engine_inject_entropy(&cc->eng, seed, sizeof(seed));

    if (!br_ssl_client_reset(cc, host, 0)) {
        montauk::closesocket(fd); free(cc); free(xc); free(iobuf); return -1;
    }

    int respLen = tls_exchange(fd, &cc->eng, request, reqLen, respBuf, respMax, abort_check);
    montauk::closesocket(fd);
    free(cc); free(xc); free(iobuf);
    return respLen;
}

} // namespace tls