#ifndef HUMANUS_MEMORY_MEM0_VECTOR_STORE_HNSWLIB_H
#define HUMANUS_MEMORY_MEM0_VECTOR_STORE_HNSWLIB_H

#include "base.h"
#include "hnswlib/hnswlib.h"

namespace humanus::mem0 {

class HNSWLibVectorStore : public VectorStore {
private:
    std::shared_ptr<hnswlib::HierarchicalNSW<float>> hnsw;
    std::shared_ptr<hnswlib::SpaceInterface<float>> space; // 保持space对象的引用以确保其生命周期
    std::unordered_map<size_t, MemoryItem> metadata_store; // 存储向量的元数据

public:
    HNSWLibVectorStore(const std::shared_ptr<VectorStoreConfig>& config) : VectorStore(config) {
        reset();
    }

    void reset() override;

    void insert(const std::vector<float>& vector, const size_t vector_id, const MemoryItem& metadata) override;

    std::vector<MemoryItem> search(const std::vector<float>& query, size_t limit, const FilterFunc& filter = nullptr) override;

    void delete_vector(size_t vector_id) override;

    void update(size_t vector_id, const std::vector<float>& vector = std::vector<float>(), const MemoryItem& metadata = MemoryItem()) override;

    MemoryItem get(size_t vector_id) override;

    std::vector<MemoryItem> list(size_t limit, const FilterFunc& filter = nullptr) override;
};

class HNSWLibFilterFunctorWrapper : public hnswlib::BaseFilterFunctor {
private:
    HNSWLibVectorStore& vector_store;
    FilterFunc filter_func;

public:
    HNSWLibFilterFunctorWrapper(HNSWLibVectorStore& store, const FilterFunc& filter_func)
    : vector_store(store), filter_func(filter_func) {}

    bool operator()(hnswlib::labeltype id) override {
        if (filter_func == nullptr) {
            return true;
        }
        
        try {
            return filter_func(vector_store.get(id));
        } catch (...) {
            return false;
        }
    }
};

}

#endif // HUMANUS_MEMORY_MEM0_VECTOR_STORE_HNSWLIB_H