#include <pybind11/pybind11.h>
#include "cache_wrapper.cuh"


namespace py = pybind11;

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {

/* CacheConfig */
/*-----------------------------------------------------------------------------------------------------------------------------------*/
    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("keySize", &gpucache::CacheConfig::key_size)
            .def_readwrite("valueSize", &gpucache::CacheConfig::value_size)
            .def_readwrite("maxQueryNum", &gpucache::CacheConfig::max_query_num)
            .def_readwrite("deviceId", &gpucache::CacheConfig::device_id)
            .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();
/*-----------------------------------------------------------------------------------------------------------------------------------*/


/* CacheWrapper */
/*-----------------------------------------------------------------------------------------------------------------------------------*/
    py::class_<gpucache::CacheWrapper> cache(m, "Cache");

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

    m.def("NewCache", &gpucache::NewCache, "create a cache",py::return_value_policy::reference);
/*-----------------------------------------------------------------------------------------------------------------------------------*/
}