#pragma once
#include "common.cuh"

namespace gpucache{
    class CacheWrapper {
    public:
        CacheWrapper(at::Tensor t, CacheConfig cfg);

        ~CacheWrapper();

        std::tuple<torch::Tensor,torch::Tensor,torch::Tensor, torch::Tensor> Get(uint32_t num_query, const torch::Tensor queries);

        void Put(uint32_t num_query, const torch::Tensor keys, const torch::Tensor values);

        CacheConfig::CacheEvictStrategy Strategy();

        uint64_t Capacity();

        uint32_t KeySize();

        uint32_t ValueSize();

        uint32_t MaxQueryNum();

        uint64_t DeviceId();

        uint32_t Dim();

        void Clear();

    private:
        void *cache;
        c10::ScalarType vtype; // value dtype
        c10::ScalarType ktype; // key dtype
        bool key_is_int32; // only support int32 and int64
        CacheConfig cache_cfg;
    };

    std::unique_ptr<CacheWrapper> NewCache(at::Tensor t, CacheConfig cfg);
}
