// Copyright 2025 XTX Markets Technologies Limited
//
// SPDX-License-Identifier: GPL-2.0-or-later

#include <rocksdb/db.h>
#include <vector>

#include "Bincode.hpp"
#include "BlockServicesCacheDB.hpp"
#include "Msgs.hpp"
#include "MsgsGen.hpp"
#include "RocksDBUtils.hpp"
#include "Time.hpp"

enum class BlockServicesCacheKey : uint8_t {
    CURRENT_BLOCK_SERVICES = 0,
    BLOCK_SERVICE = 1, // postfixed with the block service id
};

constexpr BlockServicesCacheKey CURRENT_BLOCK_SERVICES_KEY = BlockServicesCacheKey::CURRENT_BLOCK_SERVICES;
constexpr BlockServicesCacheKey BLOCK_SERVICE_KEY = BlockServicesCacheKey::BLOCK_SERVICE;

inline rocksdb::Slice blockServicesCacheKey(const BlockServicesCacheKey* k) {
    ALWAYS_ASSERT(*k != BLOCK_SERVICE_KEY);
    return rocksdb::Slice((const char*)k, sizeof(*k));
}

struct BlockServiceInfoShortBody {
    FIELDS(
        LE, uint64_t, blockServiceId, setBlockServiceId,
        LE, uint8_t, locationId, setLocationId,
        LE, uint8_t, storageClass, setStorageClass,
        EMIT_OFFSET, MIN_SIZE,
        END
    )

    BlockServiceInfoShortBody(char* data): _data(data) {}

    static size_t calcSize(const BlockServiceInfoShort& blockServiceInfo) {
        return SIZE;
    }

    void afterAlloc(const BlockServiceInfoShort& blockServiceInfo) {
        setBlockServiceId(blockServiceInfo.id.u64);
        setLocationId(blockServiceInfo.locationId);
        setStorageClass(blockServiceInfo.storageClass);
        BincodeBuf buf(_data + MIN_SIZE, FailureDomain::STATIC_SIZE);
        blockServiceInfo.failureDomain.pack(buf);
    }

    FailureDomain failureDomain() const {
        BincodeBuf buf(_data + MIN_SIZE, FailureDomain::STATIC_SIZE);
        FailureDomain fd;
        fd.unpack(buf);
        return fd;
    }

    size_t size() const {
        return SIZE;
    }
    static constexpr size_t SIZE = MIN_SIZE + FailureDomain::STATIC_SIZE;

};
struct CurrentBlockServicesBody {
    FIELDS(
        // if true (>0) then we have body in old version which is just []BlockServiceId
        // otherwise it's new version which is length followed by []BlockServiceInfoShortBody
        LE, uint8_t, oldVersion, _setOldVersion,
        EMIT_OFFSET, MIN_SIZE_V0,
        LE, uint8_t, _length, _setLength,
        EMIT_OFFSET, MIN_SIZE_V1,
        END
    )

    uint8_t length() const {
        return oldVersion() ? oldVersion() : _length();
    }

    void checkSize(size_t sz) {
        ALWAYS_ASSERT(sz >= MIN_SIZE_V0, "sz < MIN_SIZE (%s < %s)", sz, MIN_SIZE_V0);
        ALWAYS_ASSERT(sz == size(), "sz != size() (%s, %s)", sz, size());
    }

    static size_t calcSize(const std::vector<BlockServiceInfoShort>& blockServices) {
        ALWAYS_ASSERT(blockServices.size() < 256);
        return MIN_SIZE_V1 + blockServices.size()*BlockServiceInfoShortBody::SIZE;
    }

    void afterAlloc(const std::vector<BlockServiceInfoShort>& blockServices) {
        _setOldVersion(false);
        _setLength(blockServices.size());
        for(size_t i = 0; i < blockServices.size(); ++i) {
            blockServiceInfoAt(i).afterAlloc(blockServices[i]);
        }
    }

    size_t size() const {
        if (oldVersion()) {
            return MIN_SIZE_V0 + length()*sizeof(uint64_t);
        }
        return MIN_SIZE_V1 + length() * BlockServiceInfoShortBody::SIZE;
    }

    uint64_t blockIdAt(uint64_t ix) const {
        ALWAYS_ASSERT(oldVersion());
        ALWAYS_ASSERT(ix < length());
        uint64_t v;
        memcpy(&v, _data + MIN_SIZE_V0 + (ix*sizeof(uint64_t)), sizeof(uint64_t));
        return v;
    }

    BlockServiceInfoShortBody blockServiceInfoAt(uint64_t ix) const {
        ALWAYS_ASSERT(!oldVersion());
        ALWAYS_ASSERT(ix < length());
        return BlockServiceInfoShortBody(_data + MIN_SIZE_V1 + ix *BlockServiceInfoShortBody::SIZE);
    }
};

struct BlockServiceKey {
    FIELDS(
        BE, BlockServicesCacheKey, key, setKey, // always BLOCK_SERVICE_KEY
        BE, uint64_t,              blockServiceId, setBlockServiceId,
        END_STATIC
    )
};

struct BlockServiceBody {
    FIELDS(
        LE, uint8_t,  version,          setVersion,
        LE, uint64_t, id,               setId,
        FBYTES, 4,    ip1,              setIp1,
        LE, uint16_t, port1,            setPort1,
        FBYTES, 4,    ip2,              setIp2,
        LE, uint16_t, port2,            setPort2,
        LE, uint8_t,  storageClass,     setStorageClass,
        FBYTES, 16,   failureDomain,    setFailureDomain,
        FBYTES, 16,   secretKey,        setSecretKey,
        EMIT_OFFSET, V0_OFFSET,
        LE, BlockServiceFlags, flagsV1, setFlagsV1,
        EMIT_OFFSET, V1_OFFSET,
        LE, uint8_t,  locationIdV2,     setLocationIdV2,
        LE, uint64_t, capacityBytesV2,  setCapacityBytesV2,
        LE, uint64_t, availableBytesV2, setAvailableBytesV2,
        LE, uint64_t, blocksV2,         setBlocksV2,
        LE, uint64_t, firstSeenV2,      setFirstSeenV2,
        LE, uint64_t, lastSeenV2,       setLastSeenV2,
        LE, uint64_t, lastInfoChangeV2, setLastInfoChangeV2,
        LE, uint8_t,  hasFilesV2,       setHasFilesV2,
        BYTES, pathV2, setPathV2,
        END
    )

    static constexpr size_t V2_OFFSET = V1_OFFSET + 1 + 8*6 + 1;  // locationId + 6 uint64_t + hasFiles

    static constexpr size_t MIN_SIZE = V2_OFFSET + sizeof(uint8_t);  // +1 for path length byte
    static constexpr size_t MAX_SIZE = MIN_SIZE + 255;  // max path size

    size_t size() const {
        switch (version()) {
        case 0: return V0_OFFSET;
        case 1: return V1_OFFSET;
        case 2: return MIN_SIZE + pathV2().size();
        default: throw TERN_EXCEPTION("bad version %s", version());
        }
    }

    void checkSize(size_t sz) {
        ALWAYS_ASSERT(sz >= MIN_SIZE || version() < 2, "expected %s >= %s", sz, MIN_SIZE);
        ALWAYS_ASSERT(sz == size());
    }

    BlockServiceFlags flags() const {
        if (unlikely(version() == 0)) { return BlockServiceFlags::EMPTY; }
        return flagsV1();
    }

    void setFlags(BlockServiceFlags f) {
        ALWAYS_ASSERT(version() > 0);
        setFlagsV1(f);
    }

    uint8_t locationId() const {
        ALWAYS_ASSERT(version() >= 2);
        return locationIdV2();
    }

    void setLocationId(uint8_t v) {
        ALWAYS_ASSERT(version() >= 2);
        setLocationIdV2(v);
    }

    uint64_t capacityBytes() const {
        ALWAYS_ASSERT(version() >= 2);
        return capacityBytesV2();
    }

    void setCapacityBytes(uint64_t v) {
        ALWAYS_ASSERT(version() >= 2);
        setCapacityBytesV2(v);
    }

    uint64_t availableBytes() const {
        ALWAYS_ASSERT(version() >= 2);
        return availableBytesV2();
    }

    void setAvailableBytes(uint64_t v) {
        ALWAYS_ASSERT(version() >= 2);
        setAvailableBytesV2(v);
    }

    uint64_t blocks() const {
        ALWAYS_ASSERT(version() >= 2);
        return blocksV2();
    }

    void setBlocks(uint64_t v) {
        ALWAYS_ASSERT(version() >= 2);
        setBlocksV2(v);
    }

    uint64_t firstSeen() const {
        ALWAYS_ASSERT(version() >= 2);
        return firstSeenV2();
    }

    void setFirstSeen(uint64_t v) {
        ALWAYS_ASSERT(version() >= 2);
        setFirstSeenV2(v);
    }

    uint64_t lastSeen() const {
        ALWAYS_ASSERT(version() >= 2);
        return lastSeenV2();
    }

    void setLastSeen(uint64_t v) {
        ALWAYS_ASSERT(version() >= 2);
        setLastSeenV2(v);
    }

    uint64_t lastInfoChange() const {
        ALWAYS_ASSERT(version() >= 2);
        return lastInfoChangeV2();
    }

    void setLastInfoChange(uint64_t v) {
        ALWAYS_ASSERT(version() >= 2);
        setLastInfoChangeV2(v);
    }

    bool hasFiles() const {
        ALWAYS_ASSERT(version() >= 2);
        return hasFilesV2() != 0;
    }

    void setHasFiles(bool v) {
        ALWAYS_ASSERT(version() >= 2);
        setHasFilesV2(v ? 1 : 0);
    }

    std::string_view path() const {
        if (version() < 2) { return ""; }
        auto ref = pathV2();
        return std::string_view(ref.data(), ref.size());
    }

    void setPath(const std::string& p) {
        ALWAYS_ASSERT(version() >= 2);
        setPathV2(BincodeBytesRef(p.data(), p.size()));
    }
};

std::vector<rocksdb::ColumnFamilyDescriptor> BlockServicesCacheDB::getColumnFamilyDescriptors() {
    return std::vector<rocksdb::ColumnFamilyDescriptor> {{"blockServicesCache", {}}};
}

BlockServicesCacheDB::BlockServicesCacheDB(Logger& logger, std::shared_ptr<XmonAgent>& xmon, const SharedRocksDB& sharedDB, Duration blockServiceWritableDelay) :
    _env(logger, xmon, "bs_cache_db"),
    _db(sharedDB.db()),
    _blockServicesCF(sharedDB.getCF("blockServicesCache")),
    _picker(15, blockServiceWritableDelay)
{
    LOG_INFO(_env, "Initializing block services cache DB");

    std::unique_lock _(_mutex);

    const auto keyExists = [this](rocksdb::ColumnFamilyHandle* cf, const rocksdb::Slice& key) -> bool {
        std::string value;
        auto status = _db->Get({}, cf, key, &value);
        if (status.IsNotFound()) {
            return false;
        } else {
            ROCKS_DB_CHECKED(status);
            return true;
        }
    };

    if (!keyExists(_blockServicesCF, blockServicesCacheKey(&CURRENT_BLOCK_SERVICES_KEY))) {
        LOG_INFO(_env, "initializing current block services (as empty)");
        OwnedValue<CurrentBlockServicesBody> v(std::vector<BlockServiceInfoShort>{});
        ROCKS_DB_CHECKED(_db->Put({}, _blockServicesCF, blockServicesCacheKey(&CURRENT_BLOCK_SERVICES_KEY), v.toSlice()));
    } else {
        LOG_INFO(_env, "initializing block services cache (from db)");
        std::string buf;
        {
            rocksdb::ReadOptions options;
            static_assert(sizeof(BlockServicesCacheKey) == sizeof(uint8_t));
            auto upperBound = (BlockServicesCacheKey)((uint8_t)BLOCK_SERVICE_KEY + 1);
            auto upperBoundSlice = blockServicesCacheKey(&upperBound);
            options.iterate_upper_bound = &upperBoundSlice;
            std::unique_ptr<rocksdb::Iterator> it(_db->NewIterator(options, _blockServicesCF));
            StaticValue<BlockServiceKey> beginKey;
            beginKey().setKey(BLOCK_SERVICE_KEY);
            beginKey().setBlockServiceId(0);
            for (it->Seek(beginKey.toSlice()); it->Valid(); it->Next()) {
                auto k = ExternalValue<BlockServiceKey>::FromSlice(it->key());
                ALWAYS_ASSERT(k().key() == BLOCK_SERVICE_KEY);
                auto v = ExternalValue<BlockServiceBody>::FromSlice(it->value());
                auto& cache = _blockServices[k().blockServiceId()];
                cache.addrs[0].ip = v().ip1();
                cache.addrs[0].port = v().port1();
                cache.addrs[1].ip = v().ip2();
                cache.addrs[1].port = v().port2();
                expandKey(v().secretKey(), cache.secretKey);
                cache.storageClass = v().storageClass();
                cache.flags = v().flags();
                cache.failureDomain = v().failureDomain();

                if (v().version() >= 2) {
                    cache.locationId = v().locationId();
                    cache.capacityBytes = v().capacityBytes();
                    cache.availableBytes = v().availableBytes();
                    cache.blocks = v().blocks();
                    cache.firstSeen = TernTime(v().firstSeen());
                    cache.lastSeen = TernTime(v().lastSeen());
                    cache.lastInfoChange = TernTime(v().lastInfoChange());
                    cache.hasFiles = v().hasFiles();
                    cache.path = (v().version() >= 2) ? std::string(v().path()) : "";
                } else {
                    // Set defaults for v1 format
                    cache.locationId = 0;
                    cache.capacityBytes = 0;
                    cache.availableBytes = 0;
                    cache.blocks = 0;
                    cache.firstSeen = TernTime();
                    cache.lastSeen = TernTime();
                    cache.lastInfoChange = TernTime();
                    cache.hasFiles = false;
                    cache.path = "";
                }
            }
            ROCKS_DB_CHECKED(it->status());
        }
        ROCKS_DB_CHECKED(_db->Get({}, _blockServicesCF, blockServicesCacheKey(&CURRENT_BLOCK_SERVICES_KEY), &buf));
        ExternalValue<CurrentBlockServicesBody> v(buf);
        _currentBlockServices.resize(v().length());
        for (int i = 0; i < v().length(); i++) {
            auto& current = _currentBlockServices[i];
            if (v().oldVersion()) {
                current.id = v().blockIdAt(i);
                auto& blockServiceInfo = _blockServices[current.locationId];
                current.locationId = DEFAULT_LOCATION;
                current.failureDomain.name = blockServiceInfo.failureDomain;
                current.storageClass = blockServiceInfo.storageClass;
            } else {
                auto blockServiceInfo = v().blockServiceInfoAt(i);
                current.id = blockServiceInfo.blockServiceId();
                current.locationId = blockServiceInfo.locationId();
                current.storageClass = blockServiceInfo.storageClass();
                current.failureDomain = blockServiceInfo.failureDomain();
            }
        }
        _haveBlockServices = true;
    }
}

BlockServicesCache BlockServicesCacheDB::getCache() const {
    return BlockServicesCache(_mutex, _blockServices, _currentBlockServices);
}

void BlockServicesCacheDB::updateCache(const std::vector<FullBlockServiceInfo>& blockServices, const std::vector<BlockServiceInfoShort>& currentBlockServices) {
    LOG_INFO(_env, "Updating block service cache");

    std::unique_lock _(_mutex);

    rocksdb::WriteBatch batch;
    StaticValue<BlockServiceKey> blockKey;
    blockKey().setKey(BLOCK_SERVICE_KEY);
    StaticValue<BlockServiceBody> blockBody;
    for (int i = 0; i < blockServices.size(); i++) {
        const auto& entryBlock = blockServices[i];
        blockKey().setBlockServiceId(entryBlock.id.u64);
        blockBody().setVersion(2);
        blockBody().setId(entryBlock.id.u64);
        blockBody().setIp1(entryBlock.addrs[0].ip.data);
        blockBody().setPort1(entryBlock.addrs[0].port);
        blockBody().setIp2(entryBlock.addrs[1].ip.data);
        blockBody().setPort2(entryBlock.addrs[1].port);
        blockBody().setStorageClass(entryBlock.storageClass);
        blockBody().setFailureDomain(entryBlock.failureDomain.name.data);
        blockBody().setSecretKey(entryBlock.secretKey.data);
        blockBody().setFlags(entryBlock.flags);
        blockBody().setLocationId(entryBlock.locationId);
        blockBody().setCapacityBytes(entryBlock.capacityBytes);
        blockBody().setAvailableBytes(entryBlock.availableBytes);
        blockBody().setBlocks(entryBlock.blocks);
        blockBody().setLastSeen(entryBlock.lastSeen.ns);
        blockBody().setLastInfoChange(entryBlock.lastInfoChange.ns);
        blockBody().setHasFiles(entryBlock.hasFiles);

        uint64_t firstSeen;
        auto it = _blockServices.find(entryBlock.id.u64);
        if (it == _blockServices.end()) {
            firstSeen = ternNow().ns;
        } else {
            firstSeen = it->second.firstSeen.ns;
        }
        blockBody().setFirstSeen(firstSeen);
        blockBody().setPath(std::string(entryBlock.path.data(), entryBlock.path.size()));

        ROCKS_DB_CHECKED(batch.Put(_blockServicesCF, blockKey.toSlice(), blockBody.toSlice()));

        // Update in-memory cache
        auto& cache = _blockServices[entryBlock.id.u64];
        expandKey(entryBlock.secretKey.data, cache.secretKey);
        cache.addrs = entryBlock.addrs;
        cache.storageClass = entryBlock.storageClass;
        cache.failureDomain = entryBlock.failureDomain.name.data;
        cache.flags = entryBlock.flags;
        cache.locationId = entryBlock.locationId;
        cache.capacityBytes = entryBlock.capacityBytes;
        cache.availableBytes = entryBlock.availableBytes;
        cache.blocks = entryBlock.blocks;
        cache.lastSeen = entryBlock.lastSeen;
        cache.lastInfoChange = entryBlock.lastInfoChange;
        cache.hasFiles = entryBlock.hasFiles;
        cache.firstSeen = TernTime(firstSeen);
        cache.path = std::string(entryBlock.path.data(), entryBlock.path.size());
    }

    ALWAYS_ASSERT(currentBlockServices.size() < 256); // TODO handle this properly
    _currentBlockServices.clear();
    for (auto& bs: currentBlockServices) {
        _currentBlockServices.emplace_back(bs);
    }
    OwnedValue<CurrentBlockServicesBody> currentBody(_currentBlockServices);

    ROCKS_DB_CHECKED(batch.Put(_blockServicesCF, blockServicesCacheKey(&CURRENT_BLOCK_SERVICES_KEY), currentBody.toSlice()));

    // We intentionally do not flush here, it's not critical
    // and shard writes will flush anwyay.
    ROCKS_DB_CHECKED(_db->Write({}, &batch));

    _haveBlockServices = true;
    // Rebuild precomputed picker state
    _picker.update(_blockServices, _currentBlockServices);
}

TernError BlockServicesCacheDB::pickBlockServices(
    uint8_t locationId,
    uint8_t storageClass,
    int needed,
    const std::vector<BlacklistEntry>& blacklist,
    std::vector<BlockServiceId>& out
) const {
    return _picker.pick(locationId, storageClass, needed, blacklist, out);
}

bool BlockServicesCacheDB::haveBlockServices() const {
    return _haveBlockServices;
}
