// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project
// Modified from
// https://github.com/vllm-project/vllm-ascend/blob/main/csrc/utils.h

#pragma once

// #include "kernels/types.h"
#include <Python.h>
#include <c10/core/ScalarType.h>

#define _CONCAT(A, B) A##B
#define CONCAT(A, B) _CONCAT(A, B)

#define _STRINGIFY(A) #A
#define STRINGIFY(A) _STRINGIFY(A)

// A version of the TORCH_LIBRARY macro that expands the NAME, i.e. so NAME
// could be a macro instead of a literal token.
#define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE)

// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME
// could be a macro instead of a literal token.
#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE)

// REGISTER_EXTENSION allows the shared library to be loaded and initialized
// via python's import statement.
#define REGISTER_EXTENSION(NAME)                                                                \
    PyMODINIT_FUNC CONCAT(PyInit_, NAME)()                                                      \
    {                                                                                           \
        static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, STRINGIFY(NAME), nullptr, 0, \
                                            nullptr};                                           \
        return PyModule_Create(&module);                                                        \
    }

class TrochBindException : public std::exception {
private:
    std::string message = {};

public:
    explicit TrochBindException(const char* name, const char* file, const int line,
                                const std::string& error)
    {
        message = std::string("Failed: ") + name + " error " + file + ":" + std::to_string(line) +
                  " error message or error code is '" + error + "'";
    }

    const char* what() const noexcept override { return message.c_str(); }
};

#define TORCH_BIND_ASSERT(cond)                                                              \
    ;                                                                                        \
    do {                                                                                     \
        if (not(cond)) { throw TrochBindException("Assertion", __FILE__, __LINE__, #cond); } \
    } while (0)
