#pragma once

#include "common.cuh"
#include "cache.h"


namespace gpucache {


//    template<typename KeyType, typename ElemType>
//    class LRUCache;
//
//    template<typename KeyType, typename ElemType>
//    struct BucketView;
//
//    struct ThreadCtx;
//
//    template<typename KeyType, typename ElemType>
//    __device__ __host__ BucketView<KeyType, ElemType>
//    setBucketView(ThreadCtx ctx, KeyType *cache_keys, ElemType *cache_values,
//                  uint8_t *cache_timestamps, void *cache_mutexes,
//                  uint32_t num_elem_per_value, uint32_t bucket_id);
//
//    template<typename KeyType, typename ElemType>
//    class LRUCache : public Cache<KeyType, ElemType> {
//
//        friend BucketView<KeyType, ElemType> __device__ __host__ setBucketView<KeyType, ElemType>(
//                ThreadCtx ctx, KeyType *cache_keys, ElemType *cache_values,
//                uint8_t *cache_timestamps, void *cache_mutexes,
//                uint32_t num_elem_per_value, uint32_t bucket_id);
//
//    public:
//        explicit LRUCache(const CacheConfig &cfg);
//
//
//        ~LRUCache();
//
//        uint32_t KeySize() override;
//
//        uint32_t ValueSize() override;
//
//        uint64_t Capacity() override;
//
//        uint32_t NumElemsPerValue() override;
//
//        uint32_t MaxQueryNum();
//
//        uint32_t NBucket();
//
//        int8_t DeviceId() override;
//
//        uint32_t Dim() override;
//
//        // for test
//        // void *Mutex() { return bucketMutexes; }
//
//        CacheConfig::CacheEvictStrategy Strategy() override;
//
//        void Clear() override;
//
//        void Get(cudaStream_t stream, uint32_t num_query, KeyType *queries,
//                 ElemType *results, bool *find_mask) override;
//
//        void Put(cudaStream_t stream, uint32_t num_query, KeyType *putkeys,
//                 ElemType *putvalues, uint32_t *n_evict = nullptr,
//                 KeyType *evict_keys = nullptr) override;
//
//    private:
//        KeyType *keys;
//        ElemType *values;
//        uint8_t *timestamps{};
//        uint32_t nbucket; // 32 values for one bucket
//        void *bucketMutexes{};
//
//        // CacheConfig::CacheEvictStrategy strategy;
//        uint64_t capacity;
//        uint32_t keySize;
//        uint32_t valueSize;
//        uint32_t numElemPerValue; // embedding dim
//        int8_t device_id;
//        uint32_t dim;
//
//        // store missing keys and indices for Evict
//        KeyType *queryKeyBuffer{};
//        uint32_t *queryIndiceBuffer{};
//        uint32_t maxQueryNum;
//    };

    class LRUCacheWrapper {
    public:
        LRUCacheWrapper(at::Tensor t, CacheConfig cfg);

        ~LRUCacheWrapper();

        std::pair<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 *lru_cache;
        c10::ScalarType dtype;
        c10::ScalarType kdtype;
        bool key_is_int32;
        CacheConfig cache_cfg;
    };

    std::unique_ptr<LRUCacheWrapper> NewLRUCache(at::Tensor t, CacheConfig cfg);
}