#include "common.cuh"
#include "lru_cache.h"
#include <pybind11/pybind11.h>



namespace py = pybind11;

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){


    py::class_<gpucache::CacheConfig> cfg(m,"CacheConfig");

    cfg.def(py::init<>())
       .def(py::init<gpucache::CacheConfig::CacheEvictStrategy, uint64_t, uint32_t, uint32_t, uint32_t, int8_t,uint32_t>())
       .def_readwrite("strategy", &gpucache::CacheConfig::strategy)
       .def_readwrite("capacity",&gpucache::CacheConfig::capacity)
       .def_readwrite("key_size", &gpucache::CacheConfig::keySize)
       .def_readwrite("value_size",&gpucache::CacheConfig::valueSize)
       .def_readwrite("max_query_num", &gpucache::CacheConfig::maxQueryNum)
       .def_readwrite("device_id",&gpucache::CacheConfig::deviceId)
       .def_readwrite("dim",&gpucache::CacheConfig::dim);

    py::enum_<gpucache::CacheConfig::CacheEvictStrategy>(cfg,"CacheEvictStrategy")
        .value("LRU",gpucache::CacheConfig::CacheEvictStrategy::LRU)
        .value("LFU",gpucache::CacheConfig::CacheEvictStrategy::LFU)
        .value("FIFO",gpucache::CacheConfig::CacheEvictStrategy::FIFO)
        .export_values();

    py::class_<gpucache::LRUCacheWrapper> lru_cache(m, "LRUCache");

    lru_cache
        .def(py::init<at::Tensor, gpucache::CacheConfig>())
        .def("Get",&gpucache::LRUCacheWrapper::Get,"get values for keys,  find_mask return whether each key exists in cache")
        .def("Put",&gpucache::LRUCacheWrapper::Put,"put key-value pairs")
        .def("Strategy",&gpucache::LRUCacheWrapper::Strategy,"get evict strategy")
        .def("Capacity",&gpucache::LRUCacheWrapper::Capacity,"return cache capacity")
        .def("KeySize",&gpucache::LRUCacheWrapper::KeySize,"return key size")
        .def("ValueSize",&gpucache::LRUCacheWrapper::ValueSize,"return value size")
        .def("MaxQueryNum",&gpucache::LRUCacheWrapper::MaxQueryNum,"return max number of keys to get or key-values to put once")
        .def("Clear",&gpucache::LRUCacheWrapper::Clear,"clear cache")
        .def("Device",&gpucache::LRUCacheWrapper::DeviceId,"return device id")
        .def("Dim",&gpucache::LRUCacheWrapper::Dim,"return value dim");

    m.def("NewLRUCache", &gpucache::NewLRUCache, "create a lru cache",py::return_value_policy::reference);
    
}


