#pragma once
#include "cache.cuh"

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

    template<typename KeyType, typename ElemType>
    struct BucketView<KeyType,ElemType,CacheConfig::CacheEvictStrategy::FIFO> {
        __device__ BucketView(KeyType *k, ElemType *v, uint8_t *ts, WarpMutex *m, uint32_t num_elems_per_value)
                : mutex(m), bkeys(k), bvalues(v), btimestamps(ts), num_elems_per_value(num_elems_per_value) {}

        __device__ int Get(const ThreadCtx &ctx, const KeyType key) {
            KeyType lane_key = bkeys[ctx.lane_id];
            KeyType ts = btimestamps[ctx.lane_id];
            bool exist = (ts > 0 && lane_key == key);
            uint32_t exist_mask = __ballot_sync(warpFullMask, exist);
            int slot_num = __ffs(exist_mask) - 1; // if not exist slot_num is -1
            return slot_num;
        }

        __device__ int TryPut(const ThreadCtx &ctx, const KeyType key) {
            KeyType lane_key = bkeys[ctx.lane_id];
            KeyType ts = btimestamps[ctx.lane_id];
            bool exist = (ts > 0 && lane_key == key);
            uint32_t exist_mask = __ballot_sync(warpFullMask, exist);
            int slot_num = __ffs(exist_mask) - 1;
            if (slot_num == -1) { // key doesn't exist
//                    if(ctx.lane_id == 0){
//                        printf("key %u doesn't exist\n",key);
//                    }
                uint32_t empty_mask = __ballot_sync(warpFullMask, ts != 0);
//                    if(ctx.lane_id == 0){
//                        printf("try put key %u, empty_mask %u\n",key,empty_mask);
//                    }
                if (empty_mask != warpFullMask) { // bucket not full
                    slot_num = __popc(empty_mask);
                    if (ts > 0) {
                        ts++;
                    } else if (ctx.lane_id == slot_num) {
                        bkeys[slot_num] = key;
                        ts = 1;
                    }
                }
            }
            btimestamps[ctx.lane_id] = ts;
            //printf("thread %u set btimestamps[%u] = %u\n",ctx.lane_id,ctx.lane_id,ts);
            return slot_num;
        }

        __device__ int Evict(const ThreadCtx &ctx, const KeyType key, KeyType *evict_key) {
            KeyType lane_key = bkeys[ctx.lane_id];
            uint8_t ts = btimestamps[ctx.lane_id];
            uint32_t evict_mask = __ballot_sync(warpFullMask, ts == 32);
            int slot_num = __ffs(evict_mask) - 1;
            if (slot_num == ctx.lane_id) {
                ts = 1;
                *evict_key = lane_key;
                bkeys[ctx.lane_id] = key;
            } else {
                ts++;
            }
            btimestamps[ctx.lane_id] = ts;
            return slot_num;
        }

        __device__ void ReadOneValue(const ThreadCtx &ctx, uint8_t slot_num,
                                     ElemType *out) {
            for (size_t i = ctx.lane_id; i < num_elems_per_value; i += warpSize) {
                out[i] = bvalues[slot_num * num_elems_per_value + i];
            }
        }

        __device__ void WriteOneValue(const ThreadCtx &ctx, uint8_t slot_num,
                                      ElemType *v) {
            for (size_t i = ctx.lane_id; i < num_elems_per_value; i += warpSize) {
                bvalues[slot_num * num_elems_per_value + i] = v[i];
            }
        }


        WarpMutex *mutex;
        KeyType *bkeys;
        ElemType *bvalues;
        uint8_t *btimestamps;
        uint32_t num_elems_per_value;
    };

    template<typename KeyType, typename ElemType>
    class FIFOCache: public Cache<KeyType, ElemType>{
    public:
        explicit FIFOCache(CacheConfig cfg):Cache<KeyType,ElemType>(cfg){
            // this->manager = new LRUCacheManager<KeyType,ElemType>();
        }
        ~FIFOCache(){
            // delete this->manager;
        }
    };

    template<typename KeyType, typename ElemType>
    FIFOCache<KeyType,ElemType>* NewFIFOCache(CacheConfig cfg){
        return new FIFOCache<KeyType,ElemType>(cfg);
    }

}

