// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0

#pragma once

#include <memory>
#include <string>
#include <vector>

#include "onnx/version_converter/adapters/adapter.h"

namespace ONNX_NAMESPACE::version_conversion {
class AxisInputToAttribute : public Adapter {
 public:
  explicit AxisInputToAttribute(
      const std::string& op_name,
      const OpSetID& initial,
      const OpSetID& target,
      size_t axis_index,
      int64_t default_axis)
      : Adapter(op_name, initial, target), axis_index(axis_index), default_axis(default_axis) {}

  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
    if (!HasAxisInput(node)) {
      node->i_(kaxis, this->default_axis);
      return EnsureAndReturnNode(node);
    }

    const ArrayRef<Value*>& inputs = node->inputs();
    Value* axis_val = inputs[this->axis_index];
    Node* axis_node = axis_val->node();

    if (axis_node->kind() == kConstant) {
      HandleConstantNode(node, axis_node, axis_val);
      return EnsureAndReturnNode(node);
    }

    if (graph->is_constant_initializer(axis_val)) {
      HandleInitializerNode(graph, node, axis_val);
      return EnsureAndReturnNode(node);
    }

    ONNX_ASSERTM(false, "Axis input must be a constant or initializer for promotion to attribute.")
  }

 private:
  size_t axis_index;
  int64_t default_axis;

  bool HasAxisInput(const Node* node) const {
    const ArrayRef<const Value*>& inputs = node->inputs();
    return inputs.size() > this->axis_index && inputs[this->axis_index]->node()->kind() != kUndefined;
  }

  void HandleConstantNode(Node* node, Node* axis_node, Value* axis_val) const {
    const std::vector<int64_t> values = ReadInt64Tensor(axis_node->t(kvalue));
    ONNX_ASSERTM(!values.empty(), "Axis tensor must contain at least one element.")
    node->i_(kaxis, values[0]);
    node->removeInput(this->axis_index);
    if (axis_val->uses().empty()) {
      axis_node->destroy();
    }
  }

  void HandleInitializerNode(const std::shared_ptr<Graph>& graph, Node* node, Value* axis_val) const {
    const std::string initializer_name = axis_val->uniqueName();
    for (const auto& initializer : graph->initializers()) {
      if (initializer.name() == initializer_name) {
        const std::vector<int64_t> values = ReadInt64Tensor(initializer);
        ONNX_ASSERTM(!values.empty(), "Axis tensor must contain at least one element.")
        node->i_(kaxis, values[0]);
        node->removeInput(this->axis_index);
        // Remove initializer
        if (axis_val->uses().empty())
          graph->eraseInitializer(initializer_name);
        break;
      }
    }
  }

  Node* EnsureAndReturnNode(Node* node) const {
    ONNX_ASSERTM(node->hasAttribute(kaxis), "Axis attribute not created. This may be a bug.")
    return node;
  }
};

} // namespace ONNX_NAMESPACE::version_conversion