/**
 * Copyright (c) 2026 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.
 */

#include "aclnn_add_n.h"
#include "add_n.h"
#include "aclnn_kernels/cast.h"
#include "aclnn_kernels/reshape.h"
#include "aclnn_kernels/contiguous.h"
#include "conversion/broadcast_to/op_api/broadcast_to.h"
#include "aclnn_kernels/common/op_error_check.h"
#include "opdev/common_types.h"
#include "opdev/data_type_utils.h"
#include "opdev/format_utils.h"
#include "opdev/op_dfx.h"
#include "opdev/op_executor.h"
#include "opdev/op_log.h"
#include "opdev/shape_utils.h"
#include "opdev/tensor_view_utils.h"
#include "opdev/platform.h"
#include "op_api/aclnn_check.h"
#include "acl/acl_base.h"

using namespace op;
#ifdef __cplusplus
extern "C" {
#endif

#define MAX_SUPPORT_DIMS_NUMS 8

// 根据API定义，列出所有能支持的dtype
static const std::initializer_list<op::DataType> ADD_N_DTYPE_SUPPORT_LIST = {
    op::DataType::DT_INT8, op::DataType::DT_INT16, op::DataType::DT_INT32, op::DataType::DT_INT64, 
    op::DataType::DT_UINT8, op::DataType::DT_UINT16, op::DataType::DT_UINT32, op::DataType::DT_UINT64,
    op::DataType::DT_FLOAT16, op::DataType::DT_BF16, op::DataType::DT_FLOAT, op::DataType::DT_DOUBLE, 
    op::DataType::DT_COMPLEX64, op::DataType::DT_COMPLEX128
};

// 检查入参是否为nullptr
static bool CheckNotNull(const aclTensorList *tensors, const aclTensor* out)
{
    OP_CHECK_NULL(tensors, return false);
    for (uint64_t i = 0; i < tensors->Size(); i++) {
        if ((*tensors)[i] == nullptr) {
            OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Expected a proper Tensor but got null for tensor %lu.", i);
            return false;
        }
    }
    OP_CHECK_NULL(out, return false);
    return true;
}

// 检查参数dtype是否支持
static bool CheckDtypeValid(const aclTensorList* tensors, const aclTensor* out)
{
    for (uint64_t i = 0; i < tensors->Size(); i++) {
        if (!CheckType((*tensors)[i]->GetDataType(), ADD_N_DTYPE_SUPPORT_LIST)) {
            OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Tensor %lu not implemented for %s, should be in dtype support list [%s].", i,
                    op::ToString((*tensors)[i]->GetDataType()).GetString(), op::ToString(ADD_N_DTYPE_SUPPORT_LIST).GetString());
            return false;
        }
    } 
    OP_CHECK_DTYPE_NOT_SUPPORT(out, ADD_N_DTYPE_SUPPORT_LIST, return false);
    return true;
}

// 检查当前NPU架构是否支持
static bool CheckArch() 
{
    auto npuArch = op::GetCurrentPlatformInfo().GetCurNpuArch();
    if (npuArch != NpuArch::DAV_2201) {
        return false;
    }
    return true;
}

// 获取BroadcastShape
static bool GetTensorBroadcastShape(const aclTensorList* tensors, op::Shape &broadcastShape)
{
    broadcastShape = (*tensors)[0]->GetViewShape();
    for (uint64_t i = 1; i < tensors->Size(); i++) {
        if (!BroadcastInferShape((*tensors)[i]->GetViewShape(), broadcastShape, broadcastShape)) {
            return false;
        }
    }
    return true;
}

// 进行Shape检查
static bool CheckShape(const aclTensorList* tensors, const aclTensor* out)
{
    for (uint64_t i = 0; i < tensors->Size(); i++) {
        auto dimNum = (*tensors)[i]->GetViewShape().GetDimNum();
        if (dimNum > MAX_SUPPORT_DIMS_NUMS) {
            OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Dim of tensor %lu is %zu, can't be greater than %zu.", i, dimNum,
                    MAX_SUPPORT_DIMS_NUMS);
            return false;
        }
        if (dimNum == 0) {
            OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Input tensor %lu has 0 dimension (scalar), which is not supported in this operator.", i);
            return false;
        }
    }
    OP_CHECK_MAX_DIM(out, MAX_SUPPORT_DIMS_NUMS, return false);
    
    op::Shape broadcastShape;
    if (!GetTensorBroadcastShape(tensors, broadcastShape)) {
        // 检查输入tensors中的tensor是否都满足broadcast规则
        OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Input tensors can't broadcast.");
        return false;
    }
    // 输出shape应该等于输入tensors经过broadcast后的shape
    OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(out, broadcastShape, return false);
    return true;
}

// 参数检查
static aclnnStatus CheckParams(const aclTensorList* tensors, const aclTensor* out)
{
    // 固定写法，检查参数是否为空指针
    CHECK_RET(CheckNotNull(tensors, out), ACLNN_ERR_PARAM_NULLPTR);

    // 检查数据类型是否一致
    OP_CHECK_DTYPE_NOT_SAME((*tensors)[0], out, return false);
    if (!CheckArch()) {
        OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ACLNN only support ASCEND910B(A2) and ASCEND910_93(A3) series");
    }

    // 检查输入的数据类型是否为支持的数据类型
    CHECK_RET(CheckDtypeValid(tensors, out), ACLNN_ERR_PARAM_INVALID);

    // 检查输出的shape
    CHECK_RET(CheckShape(tensors, out), ACLNN_ERR_PARAM_INVALID);

    // 检查数据格式是否为ND
    if (IsPrivateFormat((*tensors)[0]->GetStorageFormat())) {
        OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format only support ND、NCHW、NHWC、HWCN、NDHWC、NCDHW.");
    }

    return ACLNN_SUCCESS;
}

aclnnStatus aclnnAddNGetWorkspaceSize(const aclTensorList* tensors, aclTensor *out, 
                                      uint64_t *workspaceSize, aclOpExecutor** executor)
{
    L2_DFX_PHASE_1(aclnnAddN, DFX_IN(tensors), DFX_OUT(out));

    // 固定写法，创建OpExecutor
    auto uniqueExecutor = CREATE_EXECUTOR();
    CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);

    // 固定写法，检查参数
    auto ret = CheckParams(tensors, out);
    CHECK_RET(ret == ACLNN_SUCCESS, ret);

    // 空tensorList处理
    if (tensors->Size() == 0) {
        *workspaceSize = 0;
        uniqueExecutor.ReleaseTo(executor);
        return ACLNN_SUCCESS;
    }

    // 空tensor处理
    for (uint64_t i = 0; i < tensors->Size(); i++) {
        if ((*tensors)[i]->IsEmpty()) {
            *workspaceSize = 0;
            uniqueExecutor.ReleaseTo(executor);
            return ACLNN_SUCCESS;
        }
    }
    
    // 对输入tensors进行broadcast
    op::Shape broadcastShape = (*tensors)[0]->GetViewShape();
    for (uint64_t i = 1; i < tensors->Size(); i++) {
        BroadcastInferShape((*tensors)[i]->GetViewShape(), broadcastShape, broadcastShape);
    }

    op::FVector<int64_t, op::MAX_DIM_NUM> broadcastDims = op::ToShapeVector(broadcastShape);
    auto broadcastShapeArray = uniqueExecutor.get()->AllocIntArray(broadcastDims.data(), broadcastDims.size());
    CHECK_RET(broadcastShapeArray != nullptr, ACLNN_ERR_INNER_NULLPTR);
    
    op::FVector<const aclTensor *> tensorList;
    for (uint64_t i = 0; i < tensors->Size(); i++) {
        auto tensorsContiguous = l0op::Contiguous((*tensors)[i], uniqueExecutor.get());
        CHECK_RET(tensorsContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
        auto tensorsBroadcast = l0op::BroadcastTo(tensorsContiguous, broadcastShapeArray, uniqueExecutor.get());
        CHECK_RET(tensorsBroadcast != nullptr, ACLNN_ERR_INNER_NULLPTR);
        tensorList.push_back(tensorsBroadcast);
    }

    // 调用AddN算子计算
    const aclTensorList *inputList = uniqueExecutor.get()->AllocTensorList(tensorList.data(), tensorList.size());
    const aclTensor *addnOut;
    if (tensors->Size() == 1 || broadcastDims.size() == 0) {
        addnOut = (*tensors)[0];
    } else {
        addnOut = l0op::AddN(inputList, uniqueExecutor.get());
    }

    CHECK_RET(addnOut != nullptr, ACLNN_ERR_INNER_NULLPTR);

    // 固定写法，将计算结果拷贝到输出out上
    auto viewCopyResult = l0op::ViewCopy(addnOut, out, uniqueExecutor.get());
    CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);

    // 固定写法，获取计算过程中需要使用的workspace大小
    *workspaceSize = uniqueExecutor->GetWorkspaceSize();
    // 需要把 uniqueExecutor持有executor转移给executor
    uniqueExecutor.ReleaseTo(executor);
    return ACLNN_SUCCESS;
}

aclnnStatus aclnnAddN(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)
{
    L2_DFX_PHASE_2(aclnnAddN);
    // 固定写法，调用框架能力，完成计算
    return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
}

#ifdef __cplusplus
}
#endif