// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include <ostream>
#include <vector>

#include "paddle/common/enforce.h"
#include "paddle/common/macros.h"
#include "paddle/pir/include/core/block.h"
#include "paddle/pir/include/core/ir_mapping.h"
#include "paddle/pir/include/core/iterator.h"
#include "paddle/pir/include/core/op_info.h"
#include "paddle/pir/include/core/operation_utils.h"
#include "paddle/pir/include/core/type.h"
#include "paddle/pir/include/core/visitors.h"
namespace pir {
class OpBase;
class Program;
class OpOperand;
class OpResult;

namespace detail {
class OpResultImpl;
class OpOperandImpl;
}  // namespace detail

class CloneOptions {
 public:
  CloneOptions()
      : clone_regions_{false},
        clone_operands_{false},
        clone_successors_{false} {}
  CloneOptions(bool clone_regions, bool clone_operands, bool clone_successors)
      : clone_regions_(clone_regions),
        clone_operands_(clone_operands),
        clone_successors_(clone_successors) {}

  bool IsCloneRegions() const { return clone_regions_; }
  bool IsCloneOperands() const { return clone_operands_; }
  bool IsCloneSuccessors() const { return clone_successors_; }

  static CloneOptions &All() {
    static CloneOptions all{true, true, true};
    return all;
  }

 private:
  bool clone_regions_{true};
  bool clone_operands_{true};
  bool clone_successors_{true};
};

class alignas(8) Operation final : public DoubleLevelContainer<Operation> {
 public:
  ///
  /// \brief Malloc memory and construct objects in the following order:
  /// OpResultImpls|Operation|OpOperandImpls.
  /// NOTE: Similar to new and delete, the destroy() and the create() need to be
  /// used in conjunction.
  ///
  PADDLE_API static Operation *Create(
      const std::vector<Value> &inputs,
      const AttributeMap &attributes,
      const std::vector<Type> &output_types,
      const OpInfo &op_info,
      size_t num_regions = 0,
      const std::vector<Block *> &successors = {},
      bool verify = true);
  PADDLE_API static Operation *Create(OperationArgument &&op_argument);

  ///
  /// \brief Deep copy all information and create a new operation.
  ///
  PADDLE_API Operation *Clone(IrMapping &ir_mapping,
                              CloneOptions options = CloneOptions()) const;
  ///
  /// \brief Destroy the operation objects and free memory by create().
  ///
  PADDLE_API void Destroy();

  PADDLE_API IrContext *ir_context() const;

  Dialect *dialect() const;

  bool operator==(const Operation &other) const { return this == &other; }

  ///
  /// \brief op attribute related public interfaces
  ///
  const AttributeMap &attributes() const { return attributes_; }
  // return nullptr if attribute not found.
  Attribute attribute(const std::string &key) const {
    auto iter = attributes_.find(key);
    return iter == attributes_.end() ? nullptr : iter->second;
  }

  template <typename T>
  T attribute(const std::string &key) const {
    return attribute(key).dyn_cast<T>();
  }
  void set_attribute(const std::string &key, Attribute value) {
    attributes_[key] = value;
  }
  void erase_attribute(const std::string &key) { attributes_.erase(key); }
  bool HasAttribute(const std::string &key) const {
    return attributes_.find(key) != attributes_.end();
  }

  void set_value_property(const std::string &key,
                          const Property &value,
                          size_t index);

  void *value_property(const std::string &key, size_t index) const;

  ///
  /// \brief op output related public interfaces
  ///
  uint32_t num_results() const { return num_results_; }
  Value result(uint32_t index) const { return OpResult(op_result_impl(index)); }
  template <typename T = Type>
  T result_type(uint32_t index) const {
    return result(index).type().dyn_cast<T>();
  }
  PADDLE_API std::vector<Value> results() const;

  ///
  /// \brief op input related public interfaces
  ///
  uint32_t num_operands() const { return num_operands_; }
  OpOperand operand(uint32_t index) const { return op_operand_impl(index); }
  PADDLE_API std::vector<OpOperand> operands() const;
  PADDLE_API Value operand_source(uint32_t index) const;
  std::vector<Value> operands_source() const;
  Type operand_type(uint32_t index) const { return operand(index).type(); }

  ///
  /// \brief op successor related public interfaces
  ///
  uint32_t num_successors() const { return num_successors_; }
  PADDLE_API BlockOperand block_operand(uint32_t index) const;
  PADDLE_API Block *successor(uint32_t index) const;
  void set_successor(Block *block, unsigned index);
  bool HasSuccessors() { return num_successors_ != 0; }

  ///
  /// \brief region related public interfaces
  ///
  using Element = Region;
  using Iterator = Region *;
  using ConstIterator = const Region *;
  uint32_t num_regions() const { return num_regions_; }
  PADDLE_API Region &region(unsigned index);
  const Region &region(unsigned index) const;
  ConstIterator begin() const { return regions_; }
  ConstIterator end() const { return regions_ + num_regions_; }
  Iterator begin() { return regions_; }
  Iterator end() { return regions_ + num_regions_; }

  /// \brief block related public interfaces
  using BlockContainer = DoubleLevelContainer<Operation>;
  BlockContainer &blocks() { return *this; }

  ///
  /// \brief parent related public interfaces
  ///
  Block *GetParent() const { return parent_; }
  Region *GetParentRegion() const;
  Operation *GetParentOp() const;
  PADDLE_API Program *GetParentProgram();
  operator Block::Iterator() { return position_; }
  operator Block::ConstIterator() const { return position_; }
  void MoveTo(Block *block, Block::Iterator position);

  PADDLE_API void Print(std::ostream &os) const;
  OpInfo info() const { return info_; }
  PADDLE_API std::string name() const;

  ///
  /// \brief Operation Walkers
  ///
  template <WalkOrder Order = WalkOrder::PostOrder, typename FuncT>
  void Walk(FuncT &&callback) {
    return pir::Walk<Order>(this, std::forward<FuncT>(callback));
  }

  ///
  /// \brief Remove this operation from its parent block and delete it.
  ///
  void Erase();

  ///
  /// \brief Returns true if this operation has no uses.
  ///
  bool use_empty();

  template <typename T>
  T dyn_cast() const {
    return CastUtil<T>::call(this);
  }

  template <typename T>
  bool isa() const {
    return T::classof(this);
  }

  template <typename Trait>
  bool HasTrait() const {
    return info_.HasTrait<Trait>();
  }

  template <typename Interface>
  bool HasInterface() const {
    return info_.HasInterface<Interface>();
  }

  /// Replace all uses of results of this operation with the provided 'values'.
  PADDLE_API void ReplaceAllUsesWith(const std::vector<Value> &values);

  void ReplaceAllUsesWith(const std::vector<OpResult> &op_results);

  inline void ReplaceAllUsesWith(Value value) {
    ReplaceAllUsesWith(std::vector<Value>{value});
  }

  void Verify();

  uint64_t id() const { return id_; }

 private:
  DISABLE_COPY_AND_ASSIGN(Operation);
  Operation(const AttributeMap &attribute,
            OpInfo op_info,
            uint32_t num_results,
            uint32_t num_operands,
            uint32_t num_regions,
            uint32_t num_successors);

  PADDLE_API int32_t ComputeOpResultOffset(uint32_t index) const;
  PADDLE_API detail::OpResultImpl *op_result_impl(uint32_t index) const;

  int32_t ComputeOpOperandOffset(uint32_t index) const;
  PADDLE_API detail::OpOperandImpl *op_operand_impl(uint32_t index) const;

  template <typename To, typename Enabler = void>
  struct CastUtil {
    static To call(const Operation *op) {
      throw("Can't dyn_cast to To, To should be a Op or Trait or Interface");
    }
  };

  // Allow access to 'SetParent'.
  friend class Block;
  void SetParent(Block *parent, const Block::Iterator &position);

  template <typename To>
  struct CastUtil<
      To,
      typename std::enable_if<std::is_base_of<OpBase, To>::value>::type> {
    static To call(const Operation *op) { return To::dyn_cast(op); }
  };

  AttributeMap attributes_;

  // store data that user create by Python
  std::vector<PropertyMap> value_properties_;

  OpInfo info_;

  static uint64_t GenerateId() {
    static std::atomic<std::uint64_t> uid{0};
    return ++uid;
  }

  const uint32_t num_results_ = 0;
  const uint32_t num_operands_ = 0;
  const uint32_t num_regions_ = 0;
  const uint32_t num_successors_ = 0;
  const uint64_t id_ = -1;

  detail::BlockOperandImpl *block_operands_{nullptr};
  Region *regions_{nullptr};
  Block *parent_{nullptr};
  Block::Iterator position_;
};

IR_API std::ostream &operator<<(std::ostream &os, const Operation &op);

}  // namespace pir