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

#include "rank_table_generator.h"
#include <atomic>
#include <set>
#include "nlohmann/json.hpp"
#include "acl/acl.h"
#include "llm_datadist/llm_datadist.h"
#include "common/llm_utils.h"
#include "common/mem_utils.h"
#include "rank_table_generator_v1.h"
#include "rank_table_generator_v2.h"

namespace llm {
namespace {
const std::string kConfigVersion = "version";
constexpr const char kConfigVersionV1[] = "1.0";
constexpr const char kConfigVersionV2[] = "1.2";
constexpr uint32_t kVersionMaxLen = 32U;
}  // namespace

std::unique_ptr<RankTableGenerator> RankTableGeneratorFactory::Create(const std::string &local_comm_res,
                                                                      const std::string &peer_comm_res) {
  nlohmann::json local_res;
  nlohmann::json peer_res;
  std::string local_version;
  std::string peer_version;
  try {
    if (!rank_table_json::IsCommResJsonSizeValid(local_comm_res) ||
        !rank_table_json::IsCommResJsonSizeValid(peer_comm_res)) {
      LLMLOGE(ge::LLM_PARAM_INVALID, "comm_res size exceeds limit");
      return nullptr;
    }
    local_res = rank_table_json::ParseCommResJson(local_comm_res);
    peer_res = rank_table_json::ParseCommResJson(peer_comm_res);
    LLMUtils::AssignRequired(local_version, kConfigVersion, local_res);
    LLMUtils::AssignRequired(peer_version, kConfigVersion, peer_res);
  } catch (const nlohmann::json::exception &e) {
    LLMLOGE(ge::LLM_PARAM_INVALID, "Invalid json, exception:%s", e.what());
    return nullptr;
  }

  if ((local_version != kConfigVersionV1 && local_version != kConfigVersionV2) ||
      (peer_version != kConfigVersionV1 && peer_version != kConfigVersionV2)) {
    LLMLOGE(ge::LLM_PARAM_INVALID, "version in option:%s only support %s and %s, local version:%s, peer version:%s.",
            llm_datadist::OPTION_LOCAL_COMM_RES, kConfigVersionV1, kConfigVersionV2, local_version.c_str(),
            peer_version.c_str());
    return nullptr;
  }

  auto version = std::min(local_version, peer_version);
  if (version == kConfigVersionV1) {
    return MakeUnique<RankTableGeneratorV1>(local_comm_res, peer_comm_res);
  }
  return MakeUnique<RankTableGeneratorV2>(local_comm_res, peer_comm_res);
}

ge::Status LocalCommResGenerator::Generate(const std::string &server_id, int32_t device_id, std::string &local_comm_res,
                                           std::optional<uint32_t> device_port) {
  const static std::set<std::string> kV2Version = {"Ascend910_9391", "Ascend910_9381", "Ascend910_9392",
                                                   "Ascend910_9382", "Ascend910_9372", "Ascend910_9362"};
  const char *version = aclrtGetSocName();
  LLM_CHECK_NOTNULL(version, "aclrt get soc name");
  const auto &it = kV2Version.find(version);
  if (it != kV2Version.cend()) {
    LLM_CHK_STATUS_RET(RankTableGeneratorV2::GenerateLocalCommRes(server_id, device_id, local_comm_res, device_port),
                       "Failed to generate local comm res, server_id:%s, device_id:%d.", server_id.c_str(), device_id);
  } else {
    LLM_CHK_STATUS_RET(RankTableGeneratorV1::GenerateLocalCommRes(server_id, device_id, local_comm_res, device_port),
                       "Failed to generate local comm res, server_id:%s, device_id:%d.", server_id.c_str(), device_id);
  }
  return ge::SUCCESS;
}

}  // namespace llm