smartbotic-database/service/src/persistence/snapshot.cpp

#include "snapshot.hpp"
#include "wal.hpp"
#include "../json_parse.hpp"

#include <algorithm>
#include <atomic>
#include <cerrno>
#include <chrono>
#include <cstring>
#include <fstream>
#include <iomanip>
#include <sstream>
#include <system_error>
#include <thread>

#include <fcntl.h>
#include <unistd.h>

#include <spdlog/spdlog.h>

#ifndef DISABLE_LZ4
#include <lz4.h>
#endif

#if defined(__GLIBC__) && !defined(__APPLE__)
#include <malloc.h>
#endif

namespace {
// v1.9.1 — `malloc_trim(0)` after a big-buffer phase finishes. The
// snapshot write and load paths both build multi-GB transient buffers
// (uncompressed body + LZ4 chunks; per-collection JSON parses). When
// those buffers free, glibc parks the pages on its arena freelist and
// only reclaims them lazily during subsequent small allocations. For a
// quiescent post-boot or post-snapshot process those allocations may
// not arrive for hours, so the operator-visible VmRSS stays inflated
// even though `estimatedMemoryBytes_` is correct. Trimming explicitly
// here turns a 2.8 GB no-op back into a 2.8 GB drop (measured on Zoe
// 2026-05-13 12:07 via the SIGUSR2 handler from v1.8.3 — same call,
// just automated). No-op on non-glibc platforms.
void releaseFreelistPages() {
#if defined(__GLIBC__) && !defined(__APPLE__)
    ::malloc_trim(0);
#endif
}
}  // namespace

namespace smartbotic::database {

namespace {

// Decompress a v4-chunked LZ4 body produced by createSnapshot.
// Body layout: sequence of [u32 uncomp_size, u32 comp_size, LZ4 bytes...] terminated
// by a (0, 0) end marker. Returns the full concatenated uncompressed payload.
std::vector<uint8_t> decompressChunked(const std::vector<uint8_t>& body,
                                        uint64_t uncompressedSize) {
#ifndef DISABLE_LZ4
    std::vector<uint8_t> out;
    out.reserve(uncompressedSize);

    size_t offset = 0;
    while (offset + 8 <= body.size()) {
        uint32_t uncompU32 = 0;
        uint32_t compU32 = 0;
        std::memcpy(&uncompU32, body.data() + offset, 4);
        std::memcpy(&compU32, body.data() + offset + 4, 4);
        offset += 8;

        if (uncompU32 == 0 && compU32 == 0) {
            break;  // end marker
        }

        if (offset + compU32 > body.size()) {
            throw std::runtime_error("Snapshot chunk extends past body");
        }

        size_t outStart = out.size();
        out.resize(outStart + uncompU32);
        int decompressed = LZ4_decompress_safe(
            reinterpret_cast<const char*>(body.data() + offset),
            reinterpret_cast<char*>(out.data() + outStart),
            static_cast<int>(compU32),
            static_cast<int>(uncompU32)
        );
        if (decompressed < 0 || static_cast<uint32_t>(decompressed) != uncompU32) {
            throw std::runtime_error("Snapshot chunk decompression failed");
        }
        offset += compU32;
    }

    if (out.size() != uncompressedSize) {
        throw std::runtime_error("Decompressed size mismatch: expected " +
                                 std::to_string(uncompressedSize) + " got " +
                                 std::to_string(out.size()));
    }
    return out;
#else
    (void)body;
    (void)uncompressedSize;
    throw std::runtime_error("LZ4 disabled — cannot decompress chunked snapshot body");
#endif
}

} // namespace

SnapshotManager::SnapshotManager(Config config)
    : config_(std::move(config)) {
    // Create snapshot directory if it doesn't exist
    std::error_code ec;
    std::filesystem::create_directories(config_.snapshotDir, ec);
}

std::filesystem::path SnapshotManager::createSnapshot(const MemoryStore& store, uint64_t walSequence) {
    // Serialize store data
    uint64_t docCount = 0;
    uint32_t collCount = 0;
    std::vector<uint8_t> uncompressedData = serializeStore(store, docCount, collCount);
    const uint64_t originalUncompressedSize = uncompressedData.size();

    // Build chunked body with streaming compression.
    // v4 body layout: sequence of [u32 uncompressed_size, u32 compressed_size, bytes...]
    // terminated by a (0, 0) end marker. Peak memory per chunk is ~16 MB output buffer
    // instead of LZ4_compressBound(full DB) which used to be ~3 GB for production sizes.
    constexpr size_t CHUNK_SIZE = 16 * 1024 * 1024;  // 16 MB uncompressed per chunk

    std::vector<uint8_t> bodyData;
    uint32_t compressionType = 0;

    if (config_.compressionEnabled) {
#ifndef DISABLE_LZ4
        compressionType = 1;
        // Rough estimate to reduce reallocation thrash; actual size depends on payload compressibility.
        bodyData.reserve(uncompressedData.size() / 2 + CHUNK_SIZE);

        size_t offset = 0;
        while (offset < uncompressedData.size()) {
            size_t thisChunk = std::min(CHUNK_SIZE, uncompressedData.size() - offset);
            int compressBound = LZ4_compressBound(static_cast<int>(thisChunk));
            std::vector<uint8_t> chunkOut(compressBound);

            int compressedSize = LZ4_compress_default(
                reinterpret_cast<const char*>(uncompressedData.data() + offset),
                reinterpret_cast<char*>(chunkOut.data()),
                static_cast<int>(thisChunk),
                compressBound
            );
            if (compressedSize <= 0) {
                throw std::runtime_error("LZ4 compression failed at offset " + std::to_string(offset));
            }

            // Append length prefixes and compressed data to bodyData
            uint32_t uncompU32 = static_cast<uint32_t>(thisChunk);
            uint32_t compU32   = static_cast<uint32_t>(compressedSize);
            const uint8_t* uncompBytes = reinterpret_cast<const uint8_t*>(&uncompU32);
            const uint8_t* compBytes   = reinterpret_cast<const uint8_t*>(&compU32);
            bodyData.insert(bodyData.end(), uncompBytes, uncompBytes + 4);
            bodyData.insert(bodyData.end(), compBytes,   compBytes   + 4);
            bodyData.insert(bodyData.end(), chunkOut.begin(), chunkOut.begin() + compressedSize);

            offset += thisChunk;
        }

        // End marker: two zero u32s
        uint32_t zero = 0;
        const uint8_t* zeroBytes = reinterpret_cast<const uint8_t*>(&zero);
        bodyData.insert(bodyData.end(), zeroBytes, zeroBytes + 4);
        bodyData.insert(bodyData.end(), zeroBytes, zeroBytes + 4);

        // Free the uncompressed buffer — we no longer need it
        std::vector<uint8_t>().swap(uncompressedData);
#else
        bodyData = std::move(uncompressedData);
        compressionType = 0;
#endif
    } else {
        bodyData = std::move(uncompressedData);
    }

    // Build header
    SnapshotHeader header;
    header.timestamp = static_cast<uint64_t>(
        std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::system_clock::now().time_since_epoch()
        ).count()
    );
    header.walSequence = walSequence;
    header.documentCount = docCount;
    header.collectionCount = collCount;
    header.uncompressedSize = originalUncompressedSize;
    header.compressedSize = bodyData.size();
    header.compressionType = compressionType;

    // Calculate header checksum (excluding the checksum field itself)
    header.headerChecksum = crc32::calculate(&header, sizeof(header) - sizeof(header.headerChecksum));

    // Generate filename — write to .tmp first, atomic rename at the end
    std::filesystem::path snapshotPath = config_.snapshotDir / generateFilename();
    std::filesystem::path tmpPath = snapshotPath;
    tmpPath += ".tmp";

    // Helper to remove tmp on error paths
    auto cleanupTmp = [&tmpPath]() {
        std::error_code ec;
        std::filesystem::remove(tmpPath, ec);
    };

    // Write to temp file
    {
        std::ofstream file(tmpPath, std::ios::binary | std::ios::trunc);
        if (!file) {
            throw std::runtime_error("Failed to create snapshot tmp file: " + tmpPath.string());
        }

        // Header + short-write check
        file.write(reinterpret_cast<const char*>(&header), sizeof(header));
        if (!file) {
            cleanupTmp();
            throw std::runtime_error("Failed to write snapshot header: " + tmpPath.string());
        }

        // Body + short-write check
        file.write(reinterpret_cast<const char*>(bodyData.data()),
                   static_cast<std::streamsize>(bodyData.size()));
        if (!file) {
            cleanupTmp();
            throw std::runtime_error("Failed to write snapshot body (short write): " + tmpPath.string());
        }

        // Trailer (body CRC) + short-write check
        uint32_t bodyChecksum = crc32::calculate(bodyData.data(), bodyData.size());
        file.write(reinterpret_cast<const char*>(&bodyChecksum), sizeof(bodyChecksum));
        if (!file) {
            cleanupTmp();
            throw std::runtime_error("Failed to write snapshot trailer: " + tmpPath.string());
        }

        // Flush C++ stream buffer
        file.flush();
        if (!file) {
            cleanupTmp();
            throw std::runtime_error("Failed to flush snapshot: " + tmpPath.string());
        }
    }  // ofstream destructor runs close()

    // fsync the file descriptor — C++ ofstream does not do this
    {
        int fd = ::open(tmpPath.c_str(), O_RDONLY);
        if (fd < 0) {
            cleanupTmp();
            throw std::runtime_error("Failed to reopen snapshot tmp for fsync: " + tmpPath.string());
        }
        if (::fsync(fd) != 0) {
            int saved_errno = errno;
            ::close(fd);
            cleanupTmp();
            throw std::runtime_error("fsync failed on snapshot tmp (errno=" +
                                     std::to_string(saved_errno) + "): " + tmpPath.string());
        }
        ::close(fd);
    }

    // Atomic rename — from this point on, either the final file exists and is complete,
    // or the .tmp is gone and no snapshot was produced.
    {
        std::error_code ec;
        std::filesystem::rename(tmpPath, snapshotPath, ec);
        if (ec) {
            cleanupTmp();
            throw std::runtime_error("Snapshot rename failed: " + ec.message());
        }
    }

    // fsync the containing directory so the rename is durable after crash
    {
        int dfd = ::open(config_.snapshotDir.c_str(), O_RDONLY | O_DIRECTORY);
        if (dfd >= 0) {
            ::fsync(dfd);
            ::close(dfd);
        }
        // Non-fatal on failure — some filesystems reject O_DIRECTORY
    }

    // Post-write verification (opt-out via config)
    bool verified = true;
    if (config_.validateAfterWrite) {
        if (!verifySnapshot(snapshotPath)) {
            // The file exists on disk but is corrupt — delete it so it can't be used
            std::error_code ec;
            std::filesystem::remove(snapshotPath, ec);
            throw std::runtime_error("Snapshot verification failed after write: " + snapshotPath.string());
        }
        spdlog::debug("Snapshot verified after write: {}", snapshotPath.string());
    }

    // Cleanup — only evict old snapshots if the new one passed verification
    // (or if cleanup_only_if_verified is disabled)
    if (verified || !config_.cleanupOnlyIfVerified) {
        cleanupOldSnapshots();
    }

    // v1.9.1 — release freelist pages after the snapshot write. We just
    // built and threw away the uncompressed serialized buffer
    // (potentially several GB) plus LZ4 chunk staging buffers; same
    // post-burst freelist-retention story as the boot-time load above.
    // Trimming here keeps the periodic snapshot path from drifting RSS
    // up over the day. Cheap relative to the snapshot write itself.
    releaseFreelistPages();

    return snapshotPath;
}

uint64_t SnapshotManager::loadLatestSnapshot(MemoryStore& store) {
    auto snapshots = listSnapshots();
    if (snapshots.empty()) {
        return 0;
    }

    return loadSnapshot(snapshots.front(), store);
}

uint64_t SnapshotManager::loadWithFallback(MemoryStore& store,
                                            std::filesystem::path* outUsed,
                                            size_t* outAttempted) {
    auto snapshots = listSnapshots();   // already sorted newest-first

    if (outAttempted) *outAttempted = 0;

    if (snapshots.empty()) {
        throw std::runtime_error("loadWithFallback: no snapshots available");
    }

    std::string lastError;
    for (size_t i = 0; i < snapshots.size(); ++i) {
        const auto& path = snapshots[i];
        if (outAttempted) (*outAttempted)++;

        try {
            uint64_t seq = loadSnapshot(path, store);
            if (outUsed) *outUsed = path;
            if (i > 0) {
                spdlog::error(
                    "Loaded fallback snapshot (index {} of {}): {} "
                    "— newer snapshots were corrupt",
                    i, snapshots.size(), path.string());
            }
            return seq;
        } catch (const std::exception& e) {
            lastError = e.what();
            spdlog::warn("Snapshot load failed for {}: {} — trying next older",
                         path.string(), e.what());
            // Note: store.clear() is called by loadSnapshot() before writing,
            // so a partial load followed by retry on the next snapshot is safe.
        }
    }

    throw std::runtime_error("loadWithFallback: all " +
                             std::to_string(snapshots.size()) +
                             " snapshots failed to load (last error: " +
                             lastError + ")");
}

uint64_t SnapshotManager::loadSnapshot(const std::filesystem::path& path, MemoryStore& store) {
    std::ifstream file(path, std::ios::binary);
    if (!file) {
        throw std::runtime_error("Failed to open snapshot file: " + path.string());
    }

    // Read header
    SnapshotHeader header;
    file.read(reinterpret_cast<char*>(&header), sizeof(header));
    if (!file) {
        throw std::runtime_error("Failed to read snapshot header");
    }

    // Verify magic
    if (std::memcmp(header.magic, "CALSNAP1", 8) != 0) {
        throw std::runtime_error("Invalid snapshot file magic");
    }

    // Verify header checksum
    uint32_t expectedHeaderCrc = crc32::calculate(&header, sizeof(header) - sizeof(header.headerChecksum));
    if (header.headerChecksum != expectedHeaderCrc) {
        throw std::runtime_error("Snapshot header checksum mismatch");
    }

    // Read body
    std::vector<uint8_t> bodyData(header.compressedSize);
    file.read(reinterpret_cast<char*>(bodyData.data()),
              static_cast<std::streamsize>(bodyData.size()));
    if (!file) {
        throw std::runtime_error("Failed to read snapshot body");
    }

    // Read and verify body checksum
    uint32_t storedBodyCrc;
    file.read(reinterpret_cast<char*>(&storedBodyCrc), sizeof(storedBodyCrc));
    uint32_t calculatedBodyCrc = crc32::calculate(bodyData.data(), bodyData.size());
    if (storedBodyCrc != calculatedBodyCrc) {
        throw std::runtime_error("Snapshot body checksum mismatch");
    }

    // Decompress if needed
    std::vector<uint8_t> uncompressedData;
    if (header.compressionType == 1) {
        if (header.version >= 4) {
            // v4: chunked LZ4 body
            uncompressedData = decompressChunked(bodyData, header.uncompressedSize);
        } else {
            // v3 and earlier: single LZ4 block
            uncompressedData = decompress(bodyData, header.uncompressedSize);
            if (uncompressedData.empty()) {
                throw std::runtime_error("Failed to decompress snapshot");
            }
        }
    } else {
        uncompressedData = std::move(bodyData);
    }

    // Clear store and deserialize
    store.clear();
    deserializeStore(uncompressedData, store, header.version);

    // v1.9.1 — release allocator freelist pages now that the multi-GB
    // transient buffers (decompressed body + per-doc JSON parses) are
    // gone. Without this, glibc keeps those pages parked on the per-
    // thread freelist; VmRSS stays high until subsequent small allocs
    // happen to land in the right bin. On Zoe (5 GB snapshot) this
    // drops post-boot RSS by ~2.8 GB.
    releaseFreelistPages();

    return header.walSequence;
}

std::vector<std::filesystem::path> SnapshotManager::listSnapshots() const {
    std::vector<std::pair<uint64_t, std::filesystem::path>> snapshots;

    // Clean up .tmp files left from crashed writes (they start with "snapshot-" and end with ".tmp")
    {
        std::error_code scan_ec;
        for (const auto& entry : std::filesystem::directory_iterator(config_.snapshotDir, scan_ec)) {
            if (!entry.is_regular_file()) continue;
            const auto& p = entry.path();
            if (p.extension() == ".tmp" &&
                p.filename().string().starts_with("snapshot-")) {
                std::error_code rm_ec;
                std::filesystem::remove(p, rm_ec);
                spdlog::info("Removed orphaned snapshot tmp file: {}", p.string());
            }
        }
    }

    std::error_code ec;
    for (const auto& entry : std::filesystem::directory_iterator(config_.snapshotDir, ec)) {
        if (entry.is_regular_file()) {
            const auto& path = entry.path();
            if (path.filename().string().starts_with("snapshot-") &&
                path.extension() == ".dat") {
                // Read header to get timestamp
                std::ifstream file(path, std::ios::binary);
                if (file) {
                    SnapshotHeader header;
                    file.read(reinterpret_cast<char*>(&header), sizeof(header));
                    if (file && std::memcmp(header.magic, "CALSNAP1", 8) == 0) {
                        snapshots.emplace_back(header.timestamp, path);
                    }
                }
            }
        }
    }

    // Sort by timestamp descending (newest first)
    std::sort(snapshots.begin(), snapshots.end(),
              [](const auto& a, const auto& b) { return a.first > b.first; });

    std::vector<std::filesystem::path> result;
    result.reserve(snapshots.size());
    for (const auto& [_, path] : snapshots) {
        result.push_back(path);
    }

    return result;
}

std::optional<SnapshotHeader> SnapshotManager::getSnapshotInfo(const std::filesystem::path& path) const {
    std::ifstream file(path, std::ios::binary);
    if (!file) {
        return std::nullopt;
    }

    SnapshotHeader header;
    file.read(reinterpret_cast<char*>(&header), sizeof(header));
    if (!file) {
        return std::nullopt;
    }

    if (std::memcmp(header.magic, "CALSNAP1", 8) != 0) {
        return std::nullopt;
    }

    return header;
}

void SnapshotManager::cleanupOldSnapshots() {
    auto snapshots = listSnapshots();

    // Delete snapshots beyond the max count
    while (snapshots.size() > config_.maxSnapshots) {
        const auto& oldestPath = snapshots.back();
        std::error_code ec;
        std::filesystem::remove(oldestPath, ec);
        snapshots.pop_back();
    }
}

bool SnapshotManager::verifySnapshot(const std::filesystem::path& path) const {
    try {
        std::ifstream file(path, std::ios::binary);
        if (!file) {
            return false;
        }

        // Read and verify header
        SnapshotHeader header;
        file.read(reinterpret_cast<char*>(&header), sizeof(header));
        if (!file) {
            return false;
        }

        if (std::memcmp(header.magic, "CALSNAP1", 8) != 0) {
            return false;
        }

        uint32_t expectedHeaderCrc = crc32::calculate(&header, sizeof(header) - sizeof(header.headerChecksum));
        if (header.headerChecksum != expectedHeaderCrc) {
            return false;
        }

        // Read and verify body
        std::vector<uint8_t> bodyData(header.compressedSize);
        file.read(reinterpret_cast<char*>(bodyData.data()),
                  static_cast<std::streamsize>(bodyData.size()));
        if (!file) {
            return false;
        }

        uint32_t storedBodyCrc;
        file.read(reinterpret_cast<char*>(&storedBodyCrc), sizeof(storedBodyCrc));
        uint32_t calculatedBodyCrc = crc32::calculate(bodyData.data(), bodyData.size());

        return storedBodyCrc == calculatedBodyCrc;
    } catch (...) {
        return false;
    }
}

std::string SnapshotManager::generateFilename() const {
    auto now = std::chrono::system_clock::now();
    auto time = std::chrono::system_clock::to_time_t(now);
    auto tm = *std::localtime(&time);

    std::ostringstream oss;
    oss << "snapshot-"
        << std::put_time(&tm, "%Y%m%d-%H%M%S")
        << ".dat";

    return oss.str();
}

std::vector<uint8_t> SnapshotManager::compress(const std::vector<uint8_t>& data) const {
#ifndef DISABLE_LZ4
    if (data.empty()) {
        return {};
    }

    int maxCompressedSize = LZ4_compressBound(static_cast<int>(data.size()));
    std::vector<uint8_t> compressed(maxCompressedSize);

    int compressedSize = LZ4_compress_default(
        reinterpret_cast<const char*>(data.data()),
        reinterpret_cast<char*>(compressed.data()),
        static_cast<int>(data.size()),
        maxCompressedSize
    );

    if (compressedSize <= 0) {
        return {};  // Compression failed
    }

    compressed.resize(compressedSize);
    return compressed;
#else
    return {};  // Compression disabled
#endif
}

std::vector<uint8_t> SnapshotManager::decompress(const std::vector<uint8_t>& data,
                                                  uint64_t uncompressedSize) const {
#ifndef DISABLE_LZ4
    if (data.empty() || uncompressedSize == 0) {
        return {};
    }

    std::vector<uint8_t> decompressed(uncompressedSize);

    int decompressedSize = LZ4_decompress_safe(
        reinterpret_cast<const char*>(data.data()),
        reinterpret_cast<char*>(decompressed.data()),
        static_cast<int>(data.size()),
        static_cast<int>(uncompressedSize)
    );

    if (decompressedSize < 0 || static_cast<uint64_t>(decompressedSize) != uncompressedSize) {
        return {};  // Decompression failed
    }

    return decompressed;
#else
    return {};  // Compression disabled
#endif
}

std::vector<uint8_t> SnapshotManager::serializeStore(const MemoryStore& store,
                                                      uint64_t& docCount,
                                                      uint32_t& collCount) const {
    std::vector<uint8_t> result;
    docCount = 0;
    collCount = 0;

    auto collections = store.getAllCollectionsWithOptions();
    collCount = static_cast<uint32_t>(collections.size());

    for (const auto& [name, options] : collections) {
        // Write collection name
        uint16_t nameLen = static_cast<uint16_t>(name.size());
        result.push_back(static_cast<uint8_t>(nameLen & 0xFF));
        result.push_back(static_cast<uint8_t>((nameLen >> 8) & 0xFF));
        result.insert(result.end(), name.begin(), name.end());

        // Write collection options
        std::string optionsJson = options.toJson().dump();
        uint32_t optionsLen = static_cast<uint32_t>(optionsJson.size());
        for (int i = 0; i < 4; ++i) {
            result.push_back(static_cast<uint8_t>((optionsLen >> (i * 8)) & 0xFF));
        }
        result.insert(result.end(), optionsJson.begin(), optionsJson.end());

        // Get all documents in collection. Snapshot must include both
        // in-memory docs AND docs that have been evicted to stubs — otherwise
        // truncateBefore(walSeq) after this snapshot deletes the WAL entries
        // that hold the only remaining copy of the evicted docs, and they're
        // gone forever. Pre-v1.7.3 only wrote in-memory docs; v1.7.0
        // introduced eviction without a corresponding snapshot path, so any
        // snapshot taken after eviction silently lost the evicted set.
        auto docs = store.getAllDocuments(name);
        auto evictedIds = store.getEvictedDocumentIds(name);

        // Reserve 8 bytes for the doc count and back-patch after writing,
        // so we can skip evicted docs that fail to load from WAL without
        // desynchronising the deserializer's counter.
        size_t countOffset = result.size();
        for (int i = 0; i < 8; ++i) result.push_back(0);

        uint64_t collDocCount = 0;

        // Write each in-memory document
        for (const auto& doc : docs) {
            std::string docJson = doc.toJson().dump();
            uint32_t docLen = static_cast<uint32_t>(docJson.size());
            for (int i = 0; i < 4; ++i) {
                result.push_back(static_cast<uint8_t>((docLen >> (i * 8)) & 0xFF));
            }
            result.insert(result.end(), docJson.begin(), docJson.end());
            collDocCount++;
        }

        // Write each evicted document, loading from WAL via the document
        // load callback. peekEvictedDocument is const and uses the same
        // path that read queries use to fault evicted docs back in.
        uint64_t evictedFailed = 0;
        for (const auto& id : evictedIds) {
            auto docOpt = store.peekEvictedDocument(name, id);
            if (!docOpt) {
                evictedFailed++;
                continue;
            }
            std::string docJson = docOpt->toJson().dump();
            uint32_t docLen = static_cast<uint32_t>(docJson.size());
            for (int i = 0; i < 4; ++i) {
                result.push_back(static_cast<uint8_t>((docLen >> (i * 8)) & 0xFF));
            }
            result.insert(result.end(), docJson.begin(), docJson.end());
            collDocCount++;
        }

        if (evictedFailed > 0) {
            spdlog::warn(
                "Snapshot: {} evicted documents in collection '{}' could not "
                "be loaded from WAL — they will be missing from the snapshot. "
                "Likely a pre-v1.7.3 snapshot already lost them.",
                evictedFailed, name);
        }

        // Back-patch the doc count
        for (int i = 0; i < 8; ++i) {
            result[countOffset + i] =
                static_cast<uint8_t>((collDocCount >> (i * 8)) & 0xFF);
        }

        docCount += collDocCount;

        // v1.9.0 — version history block is now an empty section. History
        // moved to disk-resident HistoryStore (separate `.hlog` files per
        // collection under <data_dir>/history/). We still emit the
        // 8-byte zero count so v1.8.x readers can parse the snapshot
        // (they'll find no entries and move on). v2.x can drop this
        // block entirely with a snapshot format bump.
        uint64_t historyEntryCount = 0;
        for (int i = 0; i < 8; ++i) {
            result.push_back(static_cast<uint8_t>((historyEntryCount >> (i * 8)) & 0xFF));
        }

        // Write vector data (v3)
        const auto* vectors = store.getCollectionVectors(name);
        uint64_t vecCount = (vectors != nullptr) ? static_cast<uint64_t>(vectors->size()) : 0;
        for (int i = 0; i < 8; ++i) {
            result.push_back(static_cast<uint8_t>((vecCount >> (i * 8)) & 0xFF));
        }

        if (vectors != nullptr) {
            for (const auto& [docId, vec] : *vectors) {
                // Write doc ID
                uint16_t docIdLen = static_cast<uint16_t>(docId.size());
                result.push_back(static_cast<uint8_t>(docIdLen & 0xFF));
                result.push_back(static_cast<uint8_t>((docIdLen >> 8) & 0xFF));
                result.insert(result.end(), docId.begin(), docId.end());

                // Write vector dimension
                uint32_t dim = static_cast<uint32_t>(vec.size());
                for (int i = 0; i < 4; ++i) {
                    result.push_back(static_cast<uint8_t>((dim >> (i * 8)) & 0xFF));
                }

                // Write raw float data
                const uint8_t* floatBytes = reinterpret_cast<const uint8_t*>(vec.data());
                result.insert(result.end(), floatBytes, floatBytes + dim * sizeof(float));
            }
        }
    }

    return result;
}

namespace {

// v1.8.0 — slice describing where one collection's bytes live in the
// decompressed snapshot body, used by the two-phase parallel deserializer.
struct CollectionSlice {
    std::string name;
    CollectionOptions options;
    uint64_t docCount = 0;
    size_t docSectionStart = 0;
    size_t docSectionEnd = 0;
    uint64_t historyEntryCount = 0;
    size_t historySectionStart = 0;
    size_t historySectionEnd = 0;
    uint64_t vecCount = 0;
    size_t vecSectionStart = 0;
    size_t vecSectionEnd = 0;
};

}  // namespace

void SnapshotManager::deserializeStore(const std::vector<uint8_t>& data, MemoryStore& store,
                                       uint32_t snapshotVersion) const {
    // v1.8.0 — two-phase parallel deserializer. Before this, a 5 GB snapshot
    // with many JSON-encoded documents pinned a single core at 100% for tens
    // of seconds during boot. The serialized format is a concatenation of
    // independent per-collection blocks, so we:
    //   1. Walk the body once with pure byte arithmetic (no JSON parsing) to
    //      record each collection's name, options, and the byte ranges that
    //      hold its docs / history / vectors. This is cheap and stays in L2.
    //   2. Create the collections sequentially so global state (collection
    //      map, options) is consistent before workers start.
    //   3. Spawn N worker threads (N = hardware_concurrency capped at the
    //      number of collections) that pull work off a shared index counter
    //      and parse + load each collection's payload in isolation. JSON
    //      parsing — the actual hot path — runs in parallel; per-collection
    //      MemoryStore mutexes mean no two workers contend on the same
    //      collection's documents map.
    //
    // For workloads with many small/medium collections (e.g. ShadowMan, with
    // ~30 collections) this drops boot time roughly linearly with cores. For
    // a single-giant-collection workload there is no within-collection
    // parallelism yet; that would be a follow-up.

    // ===== Phase 1: scan body byte-by-byte to find collection slices =====
    std::vector<CollectionSlice> slices;
    size_t offset = 0;
    while (offset < data.size()) {
        CollectionSlice slice;

        // Collection name
        if (offset + 2 > data.size()) break;
        uint16_t nameLen = static_cast<uint16_t>(data[offset]) |
                          (static_cast<uint16_t>(data[offset + 1]) << 8);
        offset += 2;
        if (offset + nameLen > data.size()) break;
        slice.name.assign(reinterpret_cast<const char*>(&data[offset]), nameLen);
        offset += nameLen;

        // Collection options (small JSON — keep parsing here, sequential)
        if (offset + 4 > data.size()) break;
        uint32_t optionsLen = 0;
        for (int i = 0; i < 4; ++i) {
            optionsLen |= static_cast<uint32_t>(data[offset++]) << (i * 8);
        }
        if (offset + optionsLen > data.size()) break;
        std::string optionsJson(reinterpret_cast<const char*>(&data[offset]), optionsLen);
        offset += optionsLen;
        try {
            slice.options = CollectionOptions::fromJson(smartbotic::db::parse_to_nlohmann(optionsJson));
        } catch (const nlohmann::json::exception&) {
            // Defaults — same fallback as the pre-v1.8 path.
        }

        // Document section — record its byte range, skip past without parsing
        if (offset + 8 > data.size()) break;
        for (int i = 0; i < 8; ++i) {
            slice.docCount |= static_cast<uint64_t>(data[offset++]) << (i * 8);
        }
        slice.docSectionStart = offset;
        for (uint64_t i = 0; i < slice.docCount; ++i) {
            if (offset + 4 > data.size()) { offset = data.size(); break; }
            uint32_t docLen = 0;
            for (int j = 0; j < 4; ++j) {
                docLen |= static_cast<uint32_t>(data[offset++]) << (j * 8);
            }
            if (offset + docLen > data.size()) { offset = data.size(); break; }
            offset += docLen;
        }
        slice.docSectionEnd = offset;

        // History section (v2+)
        if (snapshotVersion >= 2) {
            if (offset + 8 > data.size()) break;
            for (int i = 0; i < 8; ++i) {
                slice.historyEntryCount |= static_cast<uint64_t>(data[offset++]) << (i * 8);
            }
            slice.historySectionStart = offset;
            for (uint64_t h = 0; h < slice.historyEntryCount; ++h) {
                if (offset + 2 > data.size()) { offset = data.size(); break; }
                uint16_t docIdLen = static_cast<uint16_t>(data[offset]) |
                                   (static_cast<uint16_t>(data[offset + 1]) << 8);
                offset += 2;
                if (offset + docIdLen > data.size()) { offset = data.size(); break; }
                offset += docIdLen;
                if (offset + 4 > data.size()) { offset = data.size(); break; }
                uint32_t versionCount = 0;
                for (int i = 0; i < 4; ++i) {
                    versionCount |= static_cast<uint32_t>(data[offset++]) << (i * 8);
                }
                for (uint32_t v = 0; v < versionCount; ++v) {
                    if (offset + 4 > data.size()) { offset = data.size(); break; }
                    uint32_t verLen = 0;
                    for (int i = 0; i < 4; ++i) {
                        verLen |= static_cast<uint32_t>(data[offset++]) << (i * 8);
                    }
                    if (offset + verLen > data.size()) { offset = data.size(); break; }
                    offset += verLen;
                }
            }
            slice.historySectionEnd = offset;
        }

        // Vector section (v3+)
        if (snapshotVersion >= 3) {
            if (offset + 8 > data.size()) break;
            for (int i = 0; i < 8; ++i) {
                slice.vecCount |= static_cast<uint64_t>(data[offset++]) << (i * 8);
            }
            slice.vecSectionStart = offset;
            for (uint64_t v = 0; v < slice.vecCount; ++v) {
                if (offset + 2 > data.size()) { offset = data.size(); break; }
                uint16_t docIdLen = static_cast<uint16_t>(data[offset]) |
                                   (static_cast<uint16_t>(data[offset + 1]) << 8);
                offset += 2;
                if (offset + docIdLen > data.size()) { offset = data.size(); break; }
                offset += docIdLen;
                if (offset + 4 > data.size()) { offset = data.size(); break; }
                uint32_t dim = 0;
                for (int i = 0; i < 4; ++i) {
                    dim |= static_cast<uint32_t>(data[offset++]) << (i * 8);
                }
                size_t byteCount = static_cast<size_t>(dim) * sizeof(float);
                if (offset + byteCount > data.size()) { offset = data.size(); break; }
                offset += byteCount;
            }
            slice.vecSectionEnd = offset;
        }

        slices.push_back(std::move(slice));
    }

    // ===== Phase 2: create all collections sequentially =====
    // Doing this before workers start means each worker only mutates its own
    // collection's per-coll mutex, never the global collections map.
    for (const auto& slice : slices) {
        store.createCollection(slice.name, slice.options);
    }

    // ===== Phase 3: parallel parse + load per collection =====
    // Worker: pop a slice index, parse + load all of its docs / history /
    // vectors. JSON parsing is the dominant cost; per-collection isolation
    // means no shared-mutex contention except on the lock-free atomics.
    const auto loadOneSlice = [&data, &store, snapshotVersion](const CollectionSlice& slice) {
        const std::string& collName = slice.name;

        // --- Documents ---
        size_t off = slice.docSectionStart;
        for (uint64_t i = 0; i < slice.docCount && off < slice.docSectionEnd; ++i) {
            if (off + 4 > slice.docSectionEnd) break;
            uint32_t docLen = 0;
            for (int j = 0; j < 4; ++j) {
                docLen |= static_cast<uint32_t>(data[off++]) << (j * 8);
            }
            if (off + docLen > slice.docSectionEnd) break;
            std::string docJson(reinterpret_cast<const char*>(&data[off]), docLen);
            off += docLen;
            try {
                Document doc = Document::fromJson(smartbotic::db::parse_to_nlohmann(docJson));
                store.loadDocument(collName, doc);
            } catch (const nlohmann::json::exception&) {
                // Skip invalid document — same tolerance as pre-v1.8.
            }
        }

        // --- History ---
        if (snapshotVersion >= 2) {
            off = slice.historySectionStart;
            for (uint64_t h = 0; h < slice.historyEntryCount && off < slice.historySectionEnd; ++h) {
                if (off + 2 > slice.historySectionEnd) break;
                uint16_t docIdLen = static_cast<uint16_t>(data[off]) |
                                   (static_cast<uint16_t>(data[off + 1]) << 8);
                off += 2;
                if (off + docIdLen > slice.historySectionEnd) break;
                std::string docId(reinterpret_cast<const char*>(&data[off]), docIdLen);
                off += docIdLen;
                if (off + 4 > slice.historySectionEnd) break;
                uint32_t versionCount = 0;
                for (int i = 0; i < 4; ++i) {
                    versionCount |= static_cast<uint32_t>(data[off++]) << (i * 8);
                }
                std::deque<DocumentVersion> history;
                for (uint32_t v = 0; v < versionCount; ++v) {
                    if (off + 4 > slice.historySectionEnd) break;
                    uint32_t verLen = 0;
                    for (int i = 0; i < 4; ++i) {
                        verLen |= static_cast<uint32_t>(data[off++]) << (i * 8);
                    }
                    if (off + verLen > slice.historySectionEnd) break;
                    std::string verJson(reinterpret_cast<const char*>(&data[off]), verLen);
                    off += verLen;
                    try {
                        auto ver = DocumentVersion::fromJson(smartbotic::db::parse_to_nlohmann(verJson));
                        history.push_back(std::move(ver));
                    } catch (const nlohmann::json::exception&) {
                        // Skip invalid version entry.
                    }
                }
                if (!history.empty()) {
                    store.loadVersionHistory(collName, docId, std::move(history));
                }
            }
        }

        // --- Vectors ---
        if (snapshotVersion >= 3) {
            off = slice.vecSectionStart;
            for (uint64_t v = 0; v < slice.vecCount && off < slice.vecSectionEnd; ++v) {
                if (off + 2 > slice.vecSectionEnd) break;
                uint16_t docIdLen = static_cast<uint16_t>(data[off]) |
                                   (static_cast<uint16_t>(data[off + 1]) << 8);
                off += 2;
                if (off + docIdLen > slice.vecSectionEnd) break;
                std::string docId(reinterpret_cast<const char*>(&data[off]), docIdLen);
                off += docIdLen;
                if (off + 4 > slice.vecSectionEnd) break;
                uint32_t dim = 0;
                for (int i = 0; i < 4; ++i) {
                    dim |= static_cast<uint32_t>(data[off++]) << (i * 8);
                }
                size_t byteCount = static_cast<size_t>(dim) * sizeof(float);
                if (off + byteCount > slice.vecSectionEnd) break;
                std::vector<float> vec(dim);
                std::memcpy(vec.data(), &data[off], byteCount);
                off += byteCount;
                store.loadVector(collName, docId, std::move(vec));
            }
        }
    };

    if (slices.size() <= 1) {
        // Trivial fast path — single collection, no parallelism overhead.
        for (const auto& slice : slices) loadOneSlice(slice);
        return;
    }

    const size_t hwc = std::thread::hardware_concurrency();
    const size_t workerCount = std::min(slices.size(), hwc > 0 ? hwc : 4);
    std::atomic<size_t> nextSlice{0};
    std::vector<std::thread> workers;
    workers.reserve(workerCount);
    for (size_t w = 0; w < workerCount; ++w) {
        workers.emplace_back([&]() {
            for (;;) {
                size_t i = nextSlice.fetch_add(1, std::memory_order_relaxed);
                if (i >= slices.size()) break;
                loadOneSlice(slices[i]);
            }
        });
    }
    for (auto& t : workers) t.join();

    spdlog::info("Snapshot deserialized: {} collections via {} workers",
                 slices.size(), workerCount);
}

} // namespace smartbotic::database