/**
 * Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
 * 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 INC_EXTERNAL_GE_GE_API_H_
#define INC_EXTERNAL_GE_GE_API_H_

#include "stub/defines.h"

#include <map>
#include <string>
#include <vector>
#include <cstring>
#include "graph/ascend_string.h"
#include "tensorflow/core/graph/graph.h"

namespace ge {
Status GEInitialize(const std::map<ge::AscendString, ge::AscendString> &options);

Status GEFinalize();

ge::AscendString GEGetErrorMsgV2();
ge::AscendString GEGetWarningMsgV2();

class Shape {
 public:
  Shape() = default;
  ~Shape() = default;
  explicit Shape(const std::vector<int64_t> &dims) : dims_(dims) {}
  std::vector<int64_t> GetDims() const { return dims_; }

 private:
  std::vector<int64_t> dims_;
};

class TensorDesc {
 public:
  TensorDesc() = default;
  ~TensorDesc() = default;
  explicit TensorDesc(Shape shape, Format format = FORMAT_ND, DataType type = DT_FLOAT)
      : shape_(shape), format_(format), type_(type) {}
  Shape GetShape() const { return shape_; }
  void SetShape(const Shape &shape) { shape_ = shape; }
  DataType GetDataType() const { return type_; }

 private:
  Shape shape_;
  Format format_ = FORMAT_ND;
  DataType type_ = DT_FLOAT;
};

class Tensor {
  using DeleteFunc = std::function<void(uint8_t *)>;

 public:
  Tensor() = default;
  ~Tensor() = default;
  Tensor(const Tensor &other) {
    const static DeleteFunc deleter = [](uint8_t *p) { delete[] p; };
    desc_ = other.desc_;
    size_ = other.size_;
    data_ = std::unique_ptr<uint8_t[], DeleteFunc>(new uint8_t[size_], deleter);
    std::memcpy(data_.get(), other.GetData(), size_);
  }
  Tensor &operator=(const Tensor &) = default;

  TensorDesc GetTensorDesc() const { return desc_; }
  graphStatus SetTensorDesc(const TensorDesc &desc) {
    desc_ = desc;
    return GRAPH_SUCCESS;
  }

  const uint8_t *GetData() const { return data_.get(); }
  size_t GetSize() const { return size_; }
  std::unique_ptr<uint8_t[], DeleteFunc> ResetData() { return std::move(data_); }
  graphStatus SetData(const uint8_t *data, size_t size) {
    const static DeleteFunc deleter = [](uint8_t *p) { delete[] p; };
    data_ = std::unique_ptr<uint8_t[], DeleteFunc>(new uint8_t[size], deleter);
    std::memcpy(data_.get(), data, size);
    size_ = size;
    return GRAPH_SUCCESS;
  }

  graphStatus SetData(uint8_t *data, size_t size, const Tensor::DeleteFunc &deleter_func) {
    data_ = std::unique_ptr<uint8_t[], DeleteFunc>(new uint8_t[2], deleter_func);
    data_.reset(data);
    size_ = size;
    return GRAPH_SUCCESS;
  }

 private:
  TensorDesc desc_;
  std::unique_ptr<uint8_t[], DeleteFunc> data_;
  size_t size_;
};

class ComputeGraph {
 public:
  explicit ComputeGraph(const std::string &name) {}
  ~ComputeGraph() = default;
  std::shared_ptr<tensorflow::Graph> graph;
  size_t GetAllNodesSize() const;
  size_t GetInputSize() const;
  size_t GetOutputSize() const;
};

class GNode {
 public:
  graphStatus GetType(AscendString &type) const;
  graphStatus SetAttr(const AscendString &name, bool &attr_value) const;
};

void ConfigureGNodeStub(graphStatus get_type_status, const std::string &node_type);
bool IsHostTensorSet();

class Graph {
 public:
  Graph() = default;
  ~Graph() = default;
  std::vector<GNode> GetDirectNode() const;
  void SetNeedIteration(bool v) { need_iteration = v; }
  bool need_iteration = false;
  std::shared_ptr<tensorflow::Graph> graph;
};

class Session {
  using RunAsyncCallback = std::function<void(Status, std::vector<ge::Tensor> &)>;

 public:
  explicit Session(const std::map<std::string, std::string> &options) {}
  explicit Session(const std::map<ge::AscendString, ge::AscendString> &options) {}

  ~Session() = default;

  Status AddGraph(uint32_t graphId, const Graph &graph);

  Status AddGraph(uint32_t graphId, const Graph &graph, const std::map<ge::AscendString, ge::AscendString> &options);

  Status RemoveGraph(uint32_t graphId);

  Status RunGraphAsync(uint32_t graphId, const std::vector<ge::Tensor> &inputs, RunAsyncCallback callback);

  bool IsGraphNeedRebuild(uint32_t graphId);

 private:
  std::map<uint32_t, std::shared_ptr<tensorflow::Graph>> graphs_;
  std::map<uint32_t, bool> graph_need_rebuild_;
};

struct GraphUtilsEx {
  static Graph CreateGraphFromComputeGraph(const ge::ComputeGraphPtr compute_graph);
  static ge::ComputeGraphPtr GetComputeGraph(const ge::Graph &compute_graph);
};  // namespace GraphUtilsEx

}  // namespace ge

#endif  // INC_EXTERNAL_GE_GE_API_H_