/**
 * Copyright (c) 2025 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.
 */
 
#ifndef MAIN_TUNING_UTILS_H
#define MAIN_TUNING_UTILS_H

#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <algorithm>
#include <cstring>
#include <fstream>
#include <iomanip>
#include <queue>
#include <mutex>

#include <graph/anchor.h>
#include <graph/ge_tensor.h>
#include <graph/graph.h>
#include <graph/model.h>
#include <graph/node.h>
#include <graph/utils/graph_utils.h>
#include <graph/utils/type_utils.h>

#include "framework/common/debug/ge_log.h"
#include "external/ge_common/ge_api_types.h"
#include "utils/node_utils.h"
#include "graph/debug/ge_attr_define.h"
#include "graph/utils/op_desc_utils.h"
#include "graph/utils/tensor_utils.h"
namespace ge {
// Configure build mode, default value is "normal"
constexpr char_t BUILD_MODE[] = "ge.buildMode";
constexpr char_t BUILD_STEP[] = "ge.buildStep";
// Configure tuning path
constexpr char_t TUNING_PATH[] = "ge.tuningPath";
// for interface: aclgrphBuildModel
extern const std::set<std::string> ir_builder_supported_options_for_lx_fusion;

// Build model
constexpr char_t BUILD_MODE_NORMAL[] = "normal";
constexpr char_t BUILD_MODE_TUNING[] = "tuning";
constexpr char_t BUILD_MODE_BASELINE[] = "baseline";
constexpr char_t BUILD_MODE_OPAT_RESULT[] = "opat_result";
extern const std::set<std::string> build_mode_options;

// Build step
constexpr char_t BUILD_STEP_BEFORE_UB_MATCH[] = "before_ub_match";
constexpr char_t BUILD_STEP_AFTER_UB_MATCH[] = "after_ub_match";
constexpr char_t BUILD_STEP_AFTER_BUILDER[] = "after_builder";
constexpr char_t BUILD_STEP_AFTER_BUILDER_SUB[] = "after_builder_sub";
constexpr char_t BUILD_STEP_AFTER_MERGE[] = "after_merge";
constexpr char_t BUILD_STEP_BEFORE_BUILD[] = "before_build";
constexpr char_t BUILD_STEP_AFTER_BUILD[] = "after_build";
extern const std::set<std::string> build_step_options;

using SubgraphCreateOutNode = std::unordered_map<ComputeGraphPtr, NodePtr>;
using NodetoNodeMap = std::unordered_map<NodePtr, NodePtr>;
using NodeVec = std::vector<NodePtr>;
using NodeNametoNodeNameMap = std::map<std::string, std::string>;
using NodetoNodeNameMap = std::unordered_map<NodePtr, std::string>;
class TuningUtils {
 public:
  TuningUtils() = default;
  ~TuningUtils() = default;
  // Dump all the subgraphs and modify
  // the subgraphs in them to be executable subgraphs if exe_flag is true
  // `tuning_path` means path to save the graphs
  static graphStatus ConvertGraphToFile(std::vector<ComputeGraphPtr> tuning_subgraphs,
                                        std::vector<ComputeGraphPtr> non_tuning_subgraphs = {},
                                        const bool exe_flag = false,
                                        const std::string &path = "",
                                        const std::string &user_path = "");
  // Recovery `graph` from graph dump files configured in options
  static graphStatus ConvertFileToGraph(const std::map<int64_t, std::string> &options, ge::Graph &graph);

  static graphStatus LinkSubgraph(ComputeGraphPtr &root_graph, const ComputeGraphPtr &graph,
                                  const std::map<std::string, ComputeGraphPtr> &name_to_merged_subgraph);

private:
  // part 1
  class HelpInfo {
    HelpInfo(const int64_t index, const bool exe_flag, const bool is_tuning_graph, const std::string &path,
             const std::string &user_path) : index_(index),
                                             exe_flag_(exe_flag),
                                             is_tuning_graph_(is_tuning_graph),
                                             path_(path),
                                             user_path_(user_path) {}
    ~HelpInfo() = default;
   private:
    int64_t index_;
    bool exe_flag_;
    bool is_tuning_graph_;
    const std::string &path_;
    const std::string &user_path_;
    bool need_preprocess_ = false;
    friend class TuningUtils;
  };
  static graphStatus MakeExeGraph(ComputeGraphPtr &exe_graph,
                                  const HelpInfo& help_info);
  static graphStatus ConvertConstToWeightAttr(const ComputeGraphPtr &exe_graph);
  static graphStatus SetFileConstInfo(const NodePtr &node, const GeTensorPtr &tensor, const std::string &aoe_path,
                                      const OpDescPtr &op_desc);
  static graphStatus HandlePld(NodePtr &node, const std::string &aoe_path);
  static graphStatus HandleConst(NodePtr &node, const std::string &aoe_path);
  static graphStatus PreProcessNode(const NodePtr &node);
  static graphStatus HandleEnd(NodePtr &node);
  static graphStatus ChangePld2Data(const NodePtr &node, const NodePtr &data_node);
  static graphStatus ChangeEnd2NetOutput(NodePtr &end_node, NodePtr &out_node);
  static graphStatus LinkEnd2NetOutput(NodePtr &end_node, NodePtr &out_node);
  static graphStatus CreateDataNode(NodePtr &node, const std::string &aoe_path, NodePtr &data_node);
  static graphStatus CreateNetOutput(const NodePtr &node, NodePtr &out_node);
  static graphStatus AddAttrToDataNodeForMergeGraph(const NodePtr &pld, const NodePtr &data_node);
  static graphStatus AddAttrToNetOutputForMergeGraph(const NodePtr &end, const NodePtr &out_node, const int64_t index);
  static void DumpGraphToPath(const ComputeGraphPtr &exe_graph, const int64_t index,
                              const bool is_tuning_graph, std::string path);
  static void TryGetWeight(const NodePtr &node, std::vector<ge::GeTensorPtr> &weight);

  static SubgraphCreateOutNode create_output_;
  // part 2
  static graphStatus MergeGraph(const std::vector<ComputeGraphPtr> &subgraphs,
                                ComputeGraphPtr &output_merged_compute_graph);
  static graphStatus MergeAllSubGraph(const std::vector<ComputeGraphPtr> &subgraphs,
                                      ComputeGraphPtr &output_merged_compute_graph);
  static graphStatus MergeSubGraph(const ComputeGraphPtr &subgraph);
  // Deletes new data and output nodes added by call `MakeExeGraph()` func in part 1
  static graphStatus RemoveDataNetoutputEdge(ComputeGraphPtr &graph);
  static NodePtr FindNode(const std::string &name, int64_t &in_index);
  static graphStatus LoadGraphFromFile(const std::map<int64_t, std::string> &options,
                                       std::vector<ComputeGraphPtr> &root_graphs,
                                       std::map<std::string, std::vector<ComputeGraphPtr>> &name_to_subgraphs);
  static NodeNametoNodeNameMap data_2_end_;
  static NodetoNodeNameMap data_node_2_end_node_;
  static NodetoNodeMap data_node_2_netoutput_node_;
  static NodeVec netoutput_nodes_;
  static NodeVec merged_graph_nodes_;
  static std::mutex mutex_;
  static std::set<std::string> reusable_weight_files_;
  static std::map<std::string, int64_t> name_to_index_;
  static std::map<size_t, std::vector<std::string>> hash_to_files_;
  // for debug
  static std::string PrintCheckLog();
  static std::string GetNodeNameByAnchor(const Anchor * const anchor);
  static std::string GenerateFileConstPath(const std::string &aoe_path, const OpDescPtr &op_desc);
  static Status GetOrSaveReusableFileConst(const GeTensorPtr &tensor, std::string &file_path);
  static Status CheckFilesSame(const std::string &file_name, const char_t *const data, const size_t data_length,
                               bool &is_content_same);
};
}
#endif // MAIN_TUNING_UTILS_H