/**
 * 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.
 */

/*!
 * \file tile_infershape.cpp
 * \brief
 */
#include <limits>

#include "register/op_impl_registry.h"
#include "op_api/op_util.h"
#include "log/log.h"
#include "util/shape_util.h"

using namespace ge;
using namespace Ops::Base;

namespace ops {
constexpr size_t ATTR_INDEX_MULTIPLES = 0;
constexpr size_t TILE_IN_IDX = 0;
constexpr size_t TILE_OUT_IDX = 0;
static constexpr size_t MAXDIMNUM = 8;
constexpr size_t INPUT_INDEX_MULTIPLES = 1;

template <typename T>
ge::graphStatus TileInferShapeCommon(gert::InferShapeContext* context, const T* multiples_data, size_t multiples_len)
{
    auto in_shape = context->GetInputShape(TILE_IN_IDX);
    OP_CHECK_NULL_WITH_CONTEXT(context, in_shape);
    auto out_shape = context->GetOutputShape(TILE_OUT_IDX);
    OP_CHECK_NULL_WITH_CONTEXT(context, out_shape);
    OP_CHECK_IF(multiples_data == nullptr, OP_LOGE_WITH_INVALID_INPUT(context->GetNodeName(), "multiples"),
                return ge::GRAPH_FAILED);
    auto in_shape_len = in_shape->GetDimNum();
    if (Ops::Base::IsUnknownRank(*in_shape)) {
        OP_LOGI(context->GetNodeName(), "input shape is unknown rank [-2], output rank = multiples len.");
        Ops::Base::SetUnknownShape(static_cast<int64_t>(multiples_len), *out_shape);
        return GRAPH_SUCCESS;
    }
    OP_CHECK_IF(multiples_len > MAXDIMNUM,
                OP_LOGE(context->GetNodeName(), "the tile multiples len is more than MaxDimNum 8"),
                return ge::GRAPH_FAILED);
    // align shape for input
    gert::Shape in_shape_new;
    if (in_shape_len < multiples_len) {
        OP_LOGI(context->GetNodeName(), "The tile multiples len is more than the input len.");
        int32_t len_diff = multiples_len - in_shape_len;
        for (int32_t i = 0; i < len_diff; i++) {
            in_shape_new.AppendDim(1);
        }
        for (size_t i = 0; i < in_shape_len; i++) {
            in_shape_new.AppendDim(in_shape->GetDim(i));
        }
        in_shape_len = multiples_len;
    } else {
        OP_LOGI(context->GetNodeName(), "The tile multiples len is less than or equal to the input len.");
        in_shape_new = *in_shape;
    }
    // in shape == [], out shape = []
    if (in_shape_len == 0) {
        OP_LOGI(context->GetNodeName(), "input shape is [], output shape is [].");
        *out_shape = *in_shape;
        return GRAPH_SUCCESS;
    }
    // calculate output shape dim value
    out_shape->SetDimNum(in_shape_len);
    for (uint64_t i = 0; i < in_shape_len; i++) {
        if (in_shape_new[i] >= 0) {
            int32_t multiples_index = static_cast<int32_t>(multiples_len - in_shape_len + i);
            int64_t multiple = 1;
            if (multiples_index >= 0) {
                multiple = static_cast<int64_t>(multiples_data[multiples_index]);
            }
            if (multiple < 0) {
                OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(
                    context->GetNodeName(), "multiples",
                    ("multiples[" + std::to_string(multiples_index) + "]=" + std::to_string(multiple)).c_str(),
                    "the value of multiples must be greater than zero");
                return ge::GRAPH_FAILED;
            }
            OP_CHECK_IF((multiple != 0) && (in_shape_new[i] > std::numeric_limits<int64_t>::max() / multiple),
                        OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(
                            context->GetNodeName(), "output.shape",
                            ("dim[" + std::to_string(i) + "]=" + std::to_string(in_shape_new[i]) + "x" +
                             std::to_string(multiple))
                                .c_str(),
                            "the output dim exceeds the int64 max value"),
                        return ge::GRAPH_FAILED);
            out_shape->SetDim(i, in_shape_new[i] * multiple);
        } else {
            // 动态shape场景：输入维度未知（值为-1等），输出对应维度保持未知，不报错
            out_shape->SetDim(i, in_shape_new[i]);
            OP_LOGI(context->GetNodeName(),
                    "Runtime infershape input dim:%lu is unknown (value is %ld), "
                    "propagate unknown dim to output.",
                    i, in_shape_new[i]);
        }
    }
    return GRAPH_SUCCESS;
}

// 非 const multiples：输出 rank 取 multiples shape 长度，未知时取 x rank
static ge::graphStatus TileInferShapeWithUnknownMultiples(gert::InferShapeContext* context, const gert::Shape* x_shape,
                                                          gert::Shape* out_shape)
{
    const auto* multiples_shape = context->GetInputShape(INPUT_INDEX_MULTIPLES);
    OP_CHECK_NULL_WITH_CONTEXT(context, multiples_shape);
    OP_CHECK_IF(multiples_shape->GetDimNum() > 1,
                OP_LOGE(context->GetNodeName(), "the rank of multiples should not be more than 1"),
                return ge::GRAPH_FAILED);
    int64_t rank = -1;
    if (multiples_shape->GetDimNum() == 0) {
        rank = 1;
    } else if (multiples_shape->GetDim(0) > 0) {
        rank = multiples_shape->GetDim(0);
    }
    if (rank < 0) {
        if (Ops::Base::IsUnknownRank(*x_shape)) {
            Ops::Base::SetUnknownRank(*out_shape);
            return ge::GRAPH_SUCCESS;
        }
        rank = static_cast<int64_t>(x_shape->GetDimNum());
    }
    Ops::Base::SetUnknownShape(rank, *out_shape);
    return ge::GRAPH_SUCCESS;
}

static ge::graphStatus InferShape4Tile(gert::InferShapeContext* context)
{
    auto x_shape = context->GetInputShape(0);
    OP_CHECK_NULL_WITH_CONTEXT(context, x_shape);
    auto multiples_tensor = context->GetInputTensor(INPUT_INDEX_MULTIPLES);
    auto out_shape = context->GetOutputShape(0);
    OP_CHECK_NULL_WITH_CONTEXT(context, out_shape);
    if (multiples_tensor == nullptr || !IsConstTensor(multiples_tensor)) {
        return TileInferShapeWithUnknownMultiples(context, x_shape, out_shape);
    }
    auto multiples_size = static_cast<size_t>(multiples_tensor->GetShapeSize());
    OP_LOGD(context->GetNodeName(), "multiples_size is %zu", multiples_size);
    DataType data_type = multiples_tensor->GetDataType();
    OP_CHECK_IF(
        (data_type != DT_INT32) && (data_type != DT_INT64),
        OP_LOGE(context->GetNodeName(), "multiples' dtype %s must be in (int32,int64)!", ToString(data_type).c_str()),
        return ge::GRAPH_FAILED);

    if (data_type == DT_INT32) {
        const int32_t* multiples_data = multiples_tensor->GetData<int32_t>();
        if (multiples_data == nullptr) {
            return TileInferShapeWithUnknownMultiples(context, x_shape, out_shape);
        }
        return TileInferShapeCommon<int32_t>(context, multiples_data, multiples_size);
    }
    const int64_t* multiples_data = multiples_tensor->GetData<int64_t>();
    if (multiples_data == nullptr) {
        return TileInferShapeWithUnknownMultiples(context, x_shape, out_shape);
    }
    return TileInferShapeCommon<int64_t>(context, multiples_data, multiples_size);
}

IMPL_OP_INFERSHAPE(Tile).InferShape(InferShape4Tile).InputsDataDependency({INPUT_INDEX_MULTIPLES});
} // namespace ops