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

#include "snapshot.hpp"
#include "wal.hpp"

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

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

#include <spdlog/spdlog.h>

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

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();
    }

    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);

    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;

        // Write version history (v2)
        auto allHistory = store.getAllVersionHistory(name);
        uint64_t historyEntryCount = allHistory.size();
        for (int i = 0; i < 8; ++i) {
            result.push_back(static_cast<uint8_t>((historyEntryCount >> (i * 8)) & 0xFF));
        }

        for (const auto& [docId, versions] : allHistory) {
            // 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 version count
            uint32_t versionCount = static_cast<uint32_t>(versions.size());
            for (int i = 0; i < 4; ++i) {
                result.push_back(static_cast<uint8_t>((versionCount >> (i * 8)) & 0xFF));
            }

            // Write each version
            for (const auto& ver : versions) {
                std::string verJson = ver.toJson().dump();
                uint32_t verLen = static_cast<uint32_t>(verJson.size());
                for (int i = 0; i < 4; ++i) {
                    result.push_back(static_cast<uint8_t>((verLen >> (i * 8)) & 0xFF));
                }
                result.insert(result.end(), verJson.begin(), verJson.end());
            }
        }

        // 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;
}

void SnapshotManager::deserializeStore(const std::vector<uint8_t>& data, MemoryStore& store,
                                       uint32_t snapshotVersion) const {
    size_t offset = 0;

    while (offset < data.size()) {
        // Read 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;
        std::string collName(reinterpret_cast<const char*>(&data[offset]), nameLen);
        offset += nameLen;

        // Read collection options
        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;

        CollectionOptions options;
        try {
            options = CollectionOptions::fromJson(nlohmann::json::parse(optionsJson));
        } catch (const nlohmann::json::exception&) {
            // Use default options
        }

        // Create collection
        store.createCollection(collName, options);

        // Read document count
        if (offset + 8 > data.size()) break;
        uint64_t docCount = 0;
        for (int i = 0; i < 8; ++i) {
            docCount |= static_cast<uint64_t>(data[offset++]) << (i * 8);
        }

        // Read each document
        for (uint64_t i = 0; i < docCount; ++i) {
            if (offset + 4 > 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()) break;
            std::string docJson(reinterpret_cast<const char*>(&data[offset]), docLen);
            offset += docLen;

            try {
                Document doc = Document::fromJson(nlohmann::json::parse(docJson));
                store.loadDocument(collName, doc);
            } catch (const nlohmann::json::exception&) {
                // Skip invalid document
            }
        }

        // Read version history (v2+)
        if (snapshotVersion >= 2) {
            if (offset + 8 > data.size()) break;
            uint64_t historyEntryCount = 0;
            for (int i = 0; i < 8; ++i) {
                historyEntryCount |= static_cast<uint64_t>(data[offset++]) << (i * 8);
            }

            for (uint64_t h = 0; h < historyEntryCount; ++h) {
                // Read doc ID
                if (offset + 2 > 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()) break;
                std::string docId(reinterpret_cast<const char*>(&data[offset]), docIdLen);
                offset += docIdLen;

                // Read version count
                if (offset + 4 > data.size()) break;
                uint32_t versionCount = 0;
                for (int i = 0; i < 4; ++i) {
                    versionCount |= static_cast<uint32_t>(data[offset++]) << (i * 8);
                }

                // Read each version into a deque
                std::deque<DocumentVersion> history;
                for (uint32_t v = 0; v < versionCount; ++v) {
                    if (offset + 4 > 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()) break;
                    std::string verJson(reinterpret_cast<const char*>(&data[offset]), verLen);
                    offset += verLen;

                    try {
                        auto ver = DocumentVersion::fromJson(nlohmann::json::parse(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));
                }
            }
        }

        // Read vector data (v3+)
        if (snapshotVersion >= 3) {
            if (offset + 8 > data.size()) break;
            uint64_t vecCount = 0;
            for (int i = 0; i < 8; ++i) {
                vecCount |= static_cast<uint64_t>(data[offset++]) << (i * 8);
            }

            for (uint64_t v = 0; v < vecCount; ++v) {
                // Read doc ID
                if (offset + 2 > 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()) break;
                std::string docId(reinterpret_cast<const char*>(&data[offset]), docIdLen);
                offset += docIdLen;

                // Read vector dimension
                if (offset + 4 > data.size()) break;
                uint32_t dim = 0;
                for (int i = 0; i < 4; ++i) {
                    dim |= static_cast<uint32_t>(data[offset++]) << (i * 8);
                }

                // Read raw float data
                size_t byteCount = static_cast<size_t>(dim) * sizeof(float);
                if (offset + byteCount > data.size()) break;
                std::vector<float> vec(dim);
                std::memcpy(vec.data(), &data[offset], byteCount);
                offset += byteCount;

                store.loadVector(collName, docId, std::move(vec));
            }
        }
    }
}

} // namespace smartbotic::database