#include "cuda_utils.hpp"

#include <cstdint>
#include <optional>
#include <torch/extension.h>

#include "utils/log.hpp"

// Neuron (native PrivateUse1 device). Device memory is not addressable through raw
// pointers from the host, so allocation goes through the torch allocator and slot
// copies are done on tensor views (see clone_and_zero_states_and_blocks).
namespace deepinfer::cuda::utils {

std::optional<torch::Tensor> alloc_tensor(ssize_t size) {
    try {
        auto t = torch::empty({(int64_t)size}, at::TensorOptions().device(get_default_device()).dtype(torch::kByte));
        LOG_INFO("alloc_tensor: allocated neuron memory of size {}", size);
        return {t};
    } catch (const c10::Error &e) {
        LOG_ERROR("alloc_tensor: failed to allocate neuron memory of size {}: {}", size, e.what());
        return std::nullopt;
    }
}

std::optional<torch::Tensor> import_ipc_tensor(py::bytes data) { TORCH_CHECK(false, "IPC tensors are not supported on Neuron"); }

py::bytes export_ipc_tensor(size_t ptr, ssize_t nbytes) { TORCH_CHECK(false, "IPC tensors are not supported on Neuron"); }

void cuda_copy_async(void *dst, void *src, size_t size, uint64_t stream) {
    TORCH_CHECK(false, "raw device-pointer copies are not supported on Neuron; copy tensor views instead");
}

uint64_t get_current_stream() { return 0; }

void pymodule_init(pybind11::module &m) {
    m.def("alloc_tensor", &deepinfer::cuda::utils::alloc_tensor, "Allocate a 1D byte tensor on the Neuron device.", py::arg("size"))
        .def("export_ipc_tensor", &deepinfer::cuda::utils::export_ipc_tensor, "Not supported on Neuron.", py::arg("ptr"),
            py::arg("nbytes"))
        .def("import_ipc_tensor", &deepinfer::cuda::utils::import_ipc_tensor, "Not supported on Neuron.", py::arg("data"));
}

} // namespace deepinfer::cuda::utils