#pragma once

#include "common.cuh"


namespace gpucache {


    template<typename KeyType, typename ElemType>
    struct CacheStorage {
        CacheStorage(CacheConfig &cfg) : key_size(cfg.key_size), value_size(cfg.value_size), capacity(cfg.capacity),
                                         max_query_num(cfg.max_query_num), device_id(cfg.device_id), dim(cfg.dim) {
            nbucket = (capacity + warpSize - 1) / warpSize;
//            printf("Cache: keySize: %lu, valueSize: %u, dim: %u, capacity: %lu, "
//                   "maxQueryNum: %u, deviceId: %u\n",
//                   sizeof(KeyType), value_size, dim, capacity, max_query_num, device_id);
            CUDA_CHECK(cudaMalloc((void **) &keys, capacity * key_size));
            CUDA_CHECK(cudaMalloc((void **) &values, capacity * value_size));
            CUDA_CHECK(cudaMalloc((void **) &timestamps, capacity * sizeof(uint8_t)));
            CUDA_CHECK(cudaMemset(timestamps, 0, capacity * sizeof(uint8_t)));
            CUDA_CHECK(cudaMalloc((void **) &bucket_mutexes, nbucket * sizeof(WarpMutex)));
            dim3 block(defaultBlockX);
            dim3 grid((nbucket + defaultBlockX - 1) / defaultBlockX);
            initLocks<<<grid, block>>>(nbucket, bucket_mutexes);
            CUDA_CHECK(cudaMalloc((void **) &query_key_buffer, max_query_num * sizeof(KeyType)));
            CUDA_CHECK(cudaMalloc((void **) &query_indice_buffer,
                                  max_query_num * sizeof(uint32_t)));
        }

        ~CacheStorage() {
            CUDA_CHECK(cudaFree(keys));
            CUDA_CHECK(cudaFree(values));
            CUDA_CHECK(cudaFree(timestamps));
            CUDA_CHECK(cudaFree(bucket_mutexes));
            CUDA_CHECK(cudaFree(query_key_buffer));
            CUDA_CHECK(cudaFree(query_indice_buffer));
        }

        KeyType *keys;
        ElemType *values;
        uint8_t *timestamps;
        uint32_t nbucket; // 32 values for one bucket
        void *bucket_mutexes{};

        // CacheConfig::CacheEvictStrategy strategy;
        uint64_t capacity;
        uint32_t key_size;
        uint32_t value_size;
        int8_t device_id;
        uint32_t dim;

        // store missing keys and indices for Evict
        KeyType *query_key_buffer{};
        uint32_t *query_indice_buffer{};
        uint32_t max_query_num;
    };


    template<typename KeyType, typename ElemType, CacheConfig::CacheEvictStrategy Strategy>
    struct BucketView;


    template<typename KeyType, typename ElemType, CacheConfig::CacheEvictStrategy Strategy>
    __device__ BucketView<KeyType, ElemType, Strategy> SetBucketView(KeyType *cache_keys, ElemType *cache_values,
                                                                     uint8_t *cache_timestamps, void *cache_mutexes,
                                                                     uint32_t num_elem_per_value, uint32_t bucket_id) {
        return BucketView<KeyType, ElemType, Strategy>(cache_keys + bucket_id * warpSize,
                                                       cache_values + bucket_id * warpSize * num_elem_per_value,
                                                       cache_timestamps + bucket_id * warpSize,
                                                       reinterpret_cast<WarpMutex *>(cache_mutexes) + bucket_id,
                                                       num_elem_per_value);
    }



    // Cache Interface
    template<typename KeyType, typename ElemType>
    class Cache {
    public:

        explicit Cache(CacheConfig cfg) : storage(cfg), cfg(cfg) {}

        virtual ~Cache() = default;

        virtual uint32_t KeySize() { return cfg.key_size; };

        virtual uint32_t ValueSize() { return cfg.value_size; };

        virtual CacheConfig::CacheEvictStrategy Strategy() { return cfg.strategy; };

        virtual uint64_t Capacity() { return cfg.capacity; };

        virtual int8_t DeviceId() { return cfg.device_id; };

        virtual uint32_t Dim() { return cfg.dim; };

        virtual uint32_t MaxQueryNum() { return cfg.max_query_num; };

        virtual void Get(cudaStream_t stream, uint32_t num_query, KeyType *queries, ElemType *results, bool *find_mask);

        virtual void Put(cudaStream_t stream, uint32_t num_keys, KeyType *keys, ElemType *values, uint32_t *n_evict = nullptr,
            KeyType *evict_keys = nullptr);

        virtual void Clear();


    protected:
        CacheConfig cfg;
        CacheStorage<KeyType, ElemType> storage;
        // CacheManager<KeyType, ElemType> *manager;
    };

    template<typename KeyType, typename ElemType>
    void Cache<KeyType, ElemType>::Clear() {
        CUDA_CHECK(cudaMemset(storage.timestamps, 0, storage.capacity * sizeof(uint8_t)));
        dim3 block(defaultBlockX);
        dim3 grid((storage.nbucket + defaultBlockX - 1) / defaultBlockX);
        initLocks<<<grid, block>>>(storage.nbucket, storage.bucket_mutexes);
    }


    template<typename KeyType, typename ElemType, CacheConfig::CacheEvictStrategy Strategy>
    __global__ void GetInternal(KeyType *cache_keys, ElemType *cache_values,
                                uint8_t *cache_timestamps, void *cache_mutexes,
                                uint32_t nbucket, uint32_t num_elem_per_value,
                                uint32_t num_query, KeyType *queries,
                                ElemType *results, bool *find_mask) {
        ThreadCtx ctx{};
        __shared__ KeyType blockQueryKeys[defaultNumWarpsPerBlock][warpSize];
        __shared__ uint32_t blockBucketIds[defaultNumWarpsPerBlock][warpSize];
        for (uint32_t offset = ctx.global_warp_idx * warpSize; offset < num_query;
             offset += ctx.num_warps * warpSize) {
            uint32_t n_query = min(warpSize, num_query - offset);
            if (ctx.lane_id < n_query) {
                uint32_t idx = offset + ctx.lane_id;
                KeyType key = queries[idx];
                size_t hash = getHash<KeyType>(key);
                uint32_t bucket_id = hash % nbucket;
                blockQueryKeys[ctx.block_warp_idx][ctx.lane_id] = key;
                blockBucketIds[ctx.block_warp_idx][ctx.lane_id] = bucket_id;
            }
            __syncwarp();

            //  32 threads compare it own slot with key
            //  if find write to result
            for (uint32_t i = 0; i < n_query; i++) {
                uint32_t idx = offset + i;
                KeyType key = blockQueryKeys[ctx.block_warp_idx][i];
                uint32_t bucket_id = blockBucketIds[ctx.block_warp_idx][i];


                auto bucket = SetBucketView<KeyType, ElemType, Strategy>(cache_keys, cache_values, cache_timestamps,
                                                                         cache_mutexes,
                                                                         num_elem_per_value, bucket_id);

                bucket.mutex->Lock(ctx);

                int slot_num = bucket.Get(ctx, key);
                if (slot_num != -1) {
                    // printf("warp_id:%d,lane_id:%d find key %d\n",ctx.global_warp_idx,ctx.lane_id,key);
                    find_mask[idx] = true;
                    bucket.ReadOneValue(ctx, slot_num, &results[idx * num_elem_per_value]);
                }
                bucket.mutex->UnLock(ctx);
            }
            __syncwarp();
        }
    }

    template<typename KeyType, typename ElemType, CacheConfig::CacheEvictStrategy Strategy>
    __global__ void

    PutWithoutEvictInternal(KeyType *cache_keys, ElemType *cache_values,
                            uint8_t *cache_timestamps, void *cache_mutexes,
                            uint32_t nbucket, uint32_t num_elem_per_value,
                            uint32_t num_query, KeyType *put_keys,
                            ElemType *put_values, uint32_t *n_missing,
                            KeyType *missing_keys, uint32_t *missing_indices) {
        ThreadCtx ctx{};
        __shared__ KeyType blockPutKeys[defaultNumWarpsPerBlock][warpSize];
        __shared__ uint32_t blockBucketIds[defaultNumWarpsPerBlock][warpSize];
        for (uint32_t offset = ctx.global_warp_idx * warpSize; offset < num_query;
             offset += ctx.num_warps * warpSize) {
            uint32_t n_query = min(warpSize, num_query - offset);
            if (ctx.lane_id < n_query) {
                uint32_t idx = offset + ctx.lane_id;
                KeyType key = put_keys[idx];
                size_t hash = getHash(key);
                uint32_t bucket_id = hash % nbucket;
                blockPutKeys[ctx.block_warp_idx][ctx.lane_id] = key;
                blockBucketIds[ctx.block_warp_idx][ctx.lane_id] = bucket_id;
            }
            __syncwarp();

            uint32_t n_warp_missing = 0;
            KeyType warp_missing_key;
            uint32_t warp_missing_index = 0;

            // 32 threads handle n_query keys together instead of 1 thread for 1 key
            for (uint32_t i = 0; i < n_query; i++) {
                uint32_t idx = offset + i;
                KeyType key = blockPutKeys[ctx.block_warp_idx][i];
                uint32_t bucket_id = blockBucketIds[ctx.block_warp_idx][i];


                auto bucket = SetBucketView<KeyType, ElemType, Strategy>(
                        cache_keys, cache_values, cache_timestamps, cache_mutexes,
                        num_elem_per_value, bucket_id);
                bucket.mutex->Lock(ctx);
                int slot_num = bucket.TryPut(ctx, key);
                if (slot_num != -1) { // insert value
                    // if(ctx.lane_id == 0){
                    //     printf("warp %u put key %u value %u\n",ctx.block_warp_idx, key,
                    //     put_values[idx*num_elem_per_value]);
                    // }
                    bucket.WriteOneValue(ctx, slot_num,
                                         &put_values[idx * num_elem_per_value]);
                } else { // bucket full record missing_key + idx for Evict
                    if (ctx.lane_id == n_warp_missing) { // i-th thread keep i-th missing key
                        // printf("warp %u lane_id %u try put key %u fail\n",
                        // ctx.global_warp_idx, ctx.lane_id, key);
                        warp_missing_key = key;
                        warp_missing_index = idx;
                    }
                    n_warp_missing += 1;
                }
                bucket.mutex->UnLock(ctx);
            }
            // reduce missing_key & idx
            if (n_warp_missing > 0) {
                uint32_t base_missing_idx = 0;
                if (ctx.lane_id == 0) {
                    base_missing_idx = atomicAdd(n_missing, n_warp_missing);
                }
                base_missing_idx = __shfl_sync(warpFullMask, base_missing_idx, 0);
                if (ctx.lane_id < n_warp_missing) {
                    // printf("warp %u thread %u add missing_key %u index %u in idx
                    // %u\n",ctx.global_warp_idx, ctx.lane_id,warp_missing_key,
                    // warp_missing_index,base_missing_idx + ctx.lane_id);
                    missing_keys[base_missing_idx + ctx.lane_id] = warp_missing_key;
                    missing_indices[base_missing_idx + ctx.lane_id] = warp_missing_index;
                }
            }
            __syncwarp();
        }
    }

    template<typename KeyType, typename ElemType, CacheConfig::CacheEvictStrategy Strategy>
    __global__ void
    EvictInternal(KeyType *cache_keys, ElemType *cache_values,
                  uint8_t *cache_timestamps, void *cache_mutexes, uint32_t nbucket,
                  uint32_t num_elem_per_value, ElemType *put_values,
                  uint32_t n_missing, KeyType *missing_keys,
                  uint32_t *missing_indices, uint32_t *num_evict,
                  KeyType *evict_keys) {
        ThreadCtx ctx{};
        __shared__ KeyType blockPutKeys[defaultNumWarpsPerBlock][warpSize];
        __shared__ uint32_t blockBucketIds[defaultNumWarpsPerBlock][warpSize];
        for (uint32_t offset = ctx.global_warp_idx * warpSize; offset < n_missing;
             offset += ctx.num_warps * warpSize) {
            uint32_t n_evict = min(warpSize, n_missing - offset);
            if (ctx.lane_id < n_evict) {
                uint32_t idx = offset + ctx.lane_id;
                KeyType key = missing_keys[idx];
                size_t hash = getHash(missing_keys[idx]);
                uint32_t bucket_id = hash % nbucket;
                blockPutKeys[ctx.block_warp_idx][ctx.lane_id] = key;
                blockBucketIds[ctx.block_warp_idx][ctx.lane_id] = bucket_id;
            }
            __syncwarp();
            for (uint32_t i = 0; i < n_evict; i++) {
                uint32_t idx = offset + i;
                KeyType key = blockPutKeys[ctx.block_warp_idx][i];
                uint32_t bucket_id = blockBucketIds[ctx.block_warp_idx][i];

                auto bucket = SetBucketView<KeyType, ElemType, Strategy>(
                        cache_keys, cache_values, cache_timestamps, cache_mutexes,
                        num_elem_per_value, bucket_id);
                bucket.mutex->Lock(ctx);
                KeyType evict_key = 0;
                int slot_num = bucket.Evict(ctx, key, &evict_key);

                uint32_t evict_keys_idx = 0;
                if (num_evict && ctx.lane_id == slot_num) {
                    evict_keys_idx = atomicAdd(num_evict, 1);
                    evict_keys[evict_keys_idx] = evict_key;
                }
                bucket.WriteOneValue(ctx, slot_num,
                                     put_values +
                                     missing_indices[idx] * num_elem_per_value);
                bucket.mutex->UnLock(ctx);
            }
        }
    }

    template<typename KeyType, typename ElemType>
    void Cache<KeyType, ElemType>::Get(cudaStream_t stream, uint32_t num_query,
                                       KeyType *queries, ElemType *results,
                                       bool *find_mask) {
        assert(num_query <= storage.max_query_num && "num_query should less than max_query_num");
        if (num_query == 0) {
            return;
        }
        dim3 block(defaultBlockX);
        dim3 grid((num_query + defaultBlockX - 1) / defaultBlockX);
        DISPATCH_GET(cfg.strategy,
                     grid, block, SharedMemorySize, stream, storage.keys, storage.values, storage.timestamps,
                     storage.bucket_mutexes, storage.nbucket, storage.dim, num_query, queries, results, find_mask);
        // GetInternal<KeyType, ElemType,><<<grid, block, SharedMemorySize, stream>>>(
//                storage.keys, storage.values, storage.timestamps, storage.bucket_mutexes, storage.nbucket,
//                storage.dim,
//                num_query, queries, results, find_mask);

    }


    template<typename KeyType, typename ElemType>
    void Cache<KeyType, ElemType>::Put(cudaStream_t stream, uint32_t num_query,
                                       KeyType *put_keys,

                                       ElemType *put_values, uint32_t *n_evict,
                                       KeyType *evict_keys) {

        assert(num_query <= storage.max_query_num && "num_query should less than max_query_num");
        if (num_query == 0) {
            return;
        }
        dim3 block(defaultBlockX);
        dim3 grid((num_query + defaultBlockX - 1) / defaultBlockX);
        uint32_t *n_missing;
        CUDA_CHECK(cudaMallocManaged(&n_missing, sizeof(uint32_t)));
        CUDA_CHECK(cudaMemset(n_missing, 0, sizeof(uint32_t)));


        DISPATCH_PUT(cfg.strategy, grid, block, SharedMemorySize, stream, storage.keys, storage.values,
                     storage.timestamps,
                     storage.bucket_mutexes, storage.nbucket,
                     storage.dim, num_query, put_keys, put_values,
                     n_missing, storage.query_key_buffer, storage.query_indice_buffer);
//        PutWithoutEvictInternal<KeyType, ElemType, cfg.strategy><<<grid, block, SharedMemorySize, stream>>>(
//                storage.keys, storage.values, storage.timestamps,
//                storage.bucket_mutexes, storage.nbucket,
//                storage.dim, num_query, put_keys, put_values,
//                n_missing, storage.query_key_buffer, storage.query_indice_buffer);
        CUDA_CHECK(cudaStreamSynchronize(stream));
        DISPATCH_EVICT(cfg.strategy, grid, block, SharedMemorySize, stream, storage.keys,
                       storage.values, storage.timestamps, storage.bucket_mutexes,
                       storage.nbucket, storage.dim, put_values, *n_missing,
                       storage.query_key_buffer, storage.query_indice_buffer, n_evict, evict_keys);
//        EvictInternal<KeyType, ElemType, cfg.strategy><<<grid, block, SharedMemorySize, stream>>>(storage.keys,
//                                                                                                  storage.values,
//                                                                                                  storage.timestamps,
//                                                                                                  storage.bucket_mutexes,
//                                                                                                  storage.nbucket,
//                                                                                                  storage.dim,
//                                                                                                  put_values,
//                                                                                                  *n_missing,
//                                                                                                  storage.query_key_buffer,
//                                                                                                  storage.query_indice_buffer,
//                                                                                                  n_evict, evict_keys);


        // printf("n_missing = %u\n",*n_missing);
        CUDA_CHECK(cudaFree(n_missing));
    }

}
