/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
#ifndef ATB_TYPES_H
#define ATB_TYPES_H
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include <acl/acl.h>
#include <aclnn/acl_meta.h>
#include "atb/svector.h"

//!
//! \file types.h
//!
//! \brief 定义加速库各种数据类型及日志错误类型
//!

namespace atb {
//! \brief 用一系列状态值表示加速库中的返回值
using Status = int32_t;
//! \brief 数据最大维度定义
constexpr uint32_t MAX_DIM = 8;

//!
//! \enum ErrorType
//!
//! \brief 加速库日志中定义的错误类型
//!
enum ErrorType : int {
    NO_ERROR = 0,                         //!< 正确
    ERROR_INVALID_PARAM,                  //!< 无效参数
    ERROR_INVALID_GRAPH,                  //!< 计算图错误
    ERROR_INTERNAL_ERROR,                 //!< 内部错误
    ERROR_RT_FAIL,                        //!< 调用Runtime接口失败
    ERROR_INVALID_IN_TENSOR_NUM,          //!< 算子输入Tensor数量与定义不一致
    ERROR_INVALID_TENSOR_DTYPE,           //!< Tensor数据类型错误
    ERROR_INVALID_TENSOR_FORMAT,          //!< Tensor数据格式错误
    ERROR_INVALID_TENSOR_DIM,             //!< Tensor数据维度错误
    ERROR_INVALID_TENSOR_SIZE,            //!< Tensor数据size错误
    ERROR_OPERATION_NULL_RUNNER,          //!< Operation内部错误
    ERROR_GRAPH_INFERSHAPE_FUNC_FAIL,     //!< 图infershapeFunc错误
    ERROR_CANN_ERROR,                     //!< 调用CANN接口错误
    ERROR_INVALID_TENSOR_INI_MATCH,       //!< ini配置文件校验失败
    ERROR_INVALID_TENSOR_ADDR,            //!< Tensor地址错误
    ERROR_INVALID_TENSOR_NUM,             //!< Tensor数量错误
    ERROR_INVALID_TENSOR_DIM_NUM,         //!< Tensor维度数错误
    ERROR_INVALID_SINGLE_OPERATION_PARAM, //!< 单Operation参数错误
    ERROR_GRAPH_NODE_RESHAPE_FUNC_FAIL,   //!< 图节点reshapeFuncs设置有误
    ERROR_INVALID_GRAPH_NODE_CHUNK,       //!< 图节点Chunk参数错误
    ERROR_INVALID_CONTEXT_ADDR,           //!< Context地址错误
    ERROR_INVALID_STREAM,                 //!< Stream错误
    ERROR_INVALID_WORKSPACE_SIZE,         //!< Workspace大小有误
    ERROR_INVALID_WORKSPACE_ADDR,         //!< Workspace地址有误
    ERROR_INVALID_OPERATION_ADDR,         //!< Operation地址有误
    ERROR_HCCL_FAIL,                      //!< HCCL接口调用失败
    ERROR_OUT_OF_DEVICE_MEMORY,           //!< Device内存不足
    ERROR_OUT_OF_HOST_MEMORY,             //!< Host内存不足
    ERROR_COMM_EMPTY,                     //!< 通信域指针为空
    ERROR_COPY_HOST_MEMORY_FAIL           //!< Host内存拷贝失败
};

//!
//! \enum LogLevel
//!
//! \brief 加速库日志级别
//!
enum class LogLevel : int {
    DEBUG = 0,          //!< DEBUG级别
    INFO,               //!< INFO级别
    WARN,               //!< WARNING级别
    ERROR,              //!< ERROR级别
    NONE                //!< 关闭日志
};

//!
//! \struct Dims
//!
//! \brief Shape维度信息。
//!
struct Dims {
    //! \brief 每一维的大小，要求大于0。
    int64_t dims[MAX_DIM];
    //! \brief Tensor的维数，取值范围为(0, 8]。
    uint64_t dimNum = 0;
};

//!
//! \struct TensorDesc
//!
//! \brief 包含对Tensor的相关描述信息：每个Tensor的数据类型，数据排布格式和形状维度信息。
//!
//! \warning Atlas 推理系列产品 中不支持ACL_BF16（bf16）类型数据。
//!
struct TensorDesc {
    //! \brief Tensor数据类型
    aclDataType dtype = ACL_DT_UNDEFINED;
    //! \brief Tensor数据排布格式
    aclFormat format = ACL_FORMAT_UNDEFINED;
    //! \brief Tensor数据的形状
    Dims shape;
};

//!
//! \struct Tensor
//!
//! \brief 加速库的Tensor定义，包含每个Tensor的描述信息、NPU内存地址、CPU内存地址和内存大小等。
//!
//!
struct Tensor {
    //! \brief Tensor描述信息
    TensorDesc desc;
    //! \brief TensorNPU内存地址。
    void *deviceData = nullptr;
    //! \brief TensorCPU内存地址。
    void *hostData = nullptr;
    //! \brief “deviceData”或“hostData”指向内容的内存大小。
    uint64_t dataSize = 0;
};

//!
//! \struct VariantPack
//!
//! \brief 加速库算子执行时需要构造VariantPack存放输入及最终输出
//!
//! \see Tensor,SVector
//!
struct VariantPack {
    //! \brief 存放所有输入Tensor的SVector
    SVector<Tensor> inTensors;
    //! \brief 存放所有输出Tensor的SVector
    SVector<Tensor> outTensors;
};

//!
//! \struct Chunk
//!
//! \brief 在host侧对Tensor执行split切分操作
//!
struct Chunk {
    //! \brief 切分数量
    uint32_t chunkNum = 1;
    //! \brief 标记均分后使用是的第几份数据
    uint32_t chunkIndex = 0;
};

//! \brief Reshape功能，改变Tensor的shape。
using ReshapeFunc = std::function<void(const Dims &oldShape, Dims &newShape)>;
//! \brief 根据输入inTensor信息推导输出outTensor信息。
using InferShapeFunc =
    std::function<Status(const SVector<TensorDesc> &inTensorDescs, SVector<TensorDesc> &outTensorDescs)>;

class Operation;

//!
//! \struct Node
//!
//! \brief 图算子中的Operation节点，每个Node表示一个Operation或者GraphOperation，所有的Node组成一个完整的图算子。
//!
//! \note inTensorIds、outTensorIds和inTensorReshapeFuncs均为SVector，每个元素的顺序需要和对应Tensor的顺序保持一致。
//!
struct Node {
    //! \brief Node对应的operation或者graphOperation。
    Operation *operation = nullptr;
    //! \brief Node对应的operation或者graphOperation的输入tensorId SVector。
    SVector<uint32_t> inTensorIds;
    //! \brief Node对应的operation或者graphOperation的输出tensorId SVector。
    SVector<uint32_t> outTensorIds;
    //! \brief Node对应的operation或者graphOperation的每个输入Tensor的reshape函数SVector。
    SVector<ReshapeFunc> inTensorReshapeFuncs;
    //! \brief 存放chunk
    SVector<Chunk> inTensorChunks;
};

//!
//! \struct GraphParam
//!
//! \brief 图算子参数。
//!
struct GraphParam {
    //! \brief 图名称。仅允许字母、数字、下划线，名称长度不超过128。
    std::string name;
    //! \brief 图算子输入Tensor的数量，需小于或等于256
    uint32_t inTensorNum = 0;
    //! \brief 图算子输出Tensor的数量，需小于或等于256
    uint32_t outTensorNum = 0;
    //! \brief 图算子中间Tensor的数量，需小于或等于256
    uint32_t internalTensorNum = 0;
    //!
    //! \brief 图算子Node Vector。
    //!
    //! nodes的长度满足小于1024。
    //! nodes的顺序，需要满足各个node对应operation的执行顺序依赖关系，先执行的在前面，后执行的在后面。
    //! 如果inTensorNum，outTensorNum与internalTensorNum之和为S，则nodes中各个元素的inTensorIds，outTensorIds中各元素值均要求落在[0, S-1]范围内。
    //!
    std::vector<Node> nodes;
    //! \brief inferShape函数指针。
    InferShapeFunc inferShapeFunc = nullptr;
};

//!
//! \struct AclNNIntArray
//!
//! \brief aclnn host tensor
//!
struct AclNNIntArray {
    /// This struct is created by calling `aclCreateIntArray` and should be destroyed by calling `aclDestroyIntArray`.
    /// \brief It is used to create the `aclOpExecutor`.
    aclIntArray *intArray = nullptr;
    /// \brief Data used to create the `aclIntArray*`. It is copied from atb::Tensor's hostData.
    std::vector<int64_t> data = {};
    /// \brief the origin of the data
    std::vector<int32_t> dataOri = {};
    /// \brief The size of `data` in bytes.
    uint64_t dataSize = 0;
};

//!
//! \struct AclNNTensor
//!
//! \brief 用于aclnn算子，包装Tensor和intArray等信息
//!
class AclNNTensor {
public:
    /// An const value to indicate that the `tensorListidx` is invalid.
    static const int64_t notInTensorList = -1;

    /// \brief Tensor passed through the ATB framework.
    atb::Tensor atbTensor;
    /// \brief The stride of each dimension in the tensor's view shape. Used when creating `aclTensor`.
    atb::SVector<int64_t> strides = {};
    /// \brief Tensor passed into the AclNN operation.
    aclTensor *tensor = nullptr;
    /// \brief An AclNNIntArray object contain tensor's host data in the int array format.
    AclNNIntArray intArrayHostData;
    /// \brief The index of the tensor in the tensor list. Used when `aclTensor` is passed into `aclTensorList`.
    int tensorListidx = notInTensorList;
    /// \brief The index of the tensor in `aclOpExecutor`'s parameter list.
    int tensorIdx = -1;
    /// An indicator that shows whether the tensor's device data needs to be updated in the execution.
    bool needUpdateTensorDataPtr = false;
};

//!
//! \struct AclNNVariantPack
//!
//! \brief aclnn算子
//!
struct AclNNVariantPack {
    /// \brief A container stores an AclNN operation's in tensor in order.
    /// Each `AclNNTensor` object contains one `aclTensor`.
    atb::SVector<std::shared_ptr<AclNNTensor>> aclInTensors;
    /// \brief A container stores an AclNN operation's out tensor in order.
    /// Each `AclNNTensor` object contains one `aclTensor`.
    atb::SVector<std::shared_ptr<AclNNTensor>> aclOutTensors;
    /// \brief A container stores an AclNN operation's input `aclTensorList` in order.
    /// Each `aclTensorList` object may contain multiple `aclTensor`.
    atb::SVector<aclTensorList *> aclInTensorList;
    /// \brief A container stores an AclNN operation's output `aclTensorList` in order.
    /// Each `aclTensorList` object may contain multiple `aclTensor`.
    atb::SVector<aclTensorList *> aclOutTensorList;
};
} // namespace atb
#endif