/**
 * 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 "acl_resource_manager.h"
#include "framework/runtime/gert_api.h"
#include "framework/common/profiling_definitions.h"
#include "framework/memory/allocator_desc.h"
#include "framework/executor/ge_executor.h"
#include "mmpa/mmpa_api.h"
#include "common/log_inner.h"
#include "framework/runtime/subscriber/global_profiler.h"
#include "ge/ge_allocator.h"
#include "acl/acl_rt.h"

namespace {
std::atomic<std::uint64_t> atomicModelId(0UL);
}
namespace acl {
namespace {
constexpr int32_t OM_FILE_SUFFIX_LEN = 3;
constexpr int32_t OM_DIR_MAX_DEPTH = 3;
constexpr int32_t DECIMAL = 10;
const std::string ACL_MAX_OPQUEUE_NUM = "max_opqueue_num";
}

AclResourceManager::AclResourceManager() {
    GetRuntimeV2Env();
}

void AclResourceManager::AddBundleSubmodelId(const uint32_t bundleId, uint32_t modelId)
{
  const std::lock_guard<std::mutex> locker(mutex_);
  bundleInfos_[bundleId].loadedSubModelIdSet.insert(modelId);
  (void)bundleInnerIds_.insert(modelId);
}

aclError AclResourceManager::SetBundleInfo(const uint32_t bundleId, const BundleModelInfo &bundleInfos)
{
  const std::lock_guard<std::mutex> locker(mutex_);
  bundleInfos_[bundleId] = bundleInfos;
  for (const auto &modelId : bundleInfos.loadedSubModelIdSet) {
    (void)bundleInnerIds_.insert(modelId);
  }
  return ACL_SUCCESS;
}

void AclResourceManager::DeleteBundleSubmodelId(const uint32_t bundleId, uint32_t modelId)
{
  const std::lock_guard<std::mutex> locker(mutex_);
  bundleInfos_[bundleId].loadedSubModelIdSet.erase(modelId);
  (void)bundleInnerIds_.erase(modelId);
}

aclError AclResourceManager::GetBundleInfo(const uint32_t bundleId, BundleModelInfo &bundleInfos)
{
    const std::lock_guard<std::mutex> locker(mutex_);
    const auto it = bundleInfos_.find(bundleId);
    if (it == bundleInfos_.end()) {
        ACL_LOG_ERROR("This model %u is not bundle model, cannot get bundle info.", bundleId);
        return ACL_ERROR_INVALID_BUNDLE_MODEL_ID;
    }
    bundleInfos = it->second;
    return ACL_SUCCESS;
}

bool AclResourceManager::IsBundleInnerId(const uint32_t modelId)
{
    const std::lock_guard<std::mutex> locker(mutex_);
    return (bundleInnerIds_.count(modelId) > 0U);
}

void AclResourceManager::DeleteBundleInfo(const uint32_t bundleId)
{
    const std::lock_guard<std::mutex> locker(mutex_);
    const auto it = bundleInfos_.find(bundleId);
    if (it != bundleInfos_.end()) {
        for (const auto &id : it->second.loadedSubModelIdSet) {
            (void)bundleInnerIds_.erase(id);
        }
        (void)bundleInfos_.erase(it);
    }
}

void AclResourceManager::AddExecutor(uint32_t &modelId, std::unique_ptr<gert::ModelV2Executor> &&executor,
                                     const std::shared_ptr<gert::RtSession> &rtSession)
{
    const std::lock_guard<std::mutex> locker(mutex_);
    ++modelIdGenerator_;
    modelId = modelIdGenerator_.load();
    executorMap_[modelId] = std::move(executor);
    rtSessionMap_[modelId] = rtSession;
}

aclError AclResourceManager::DeleteExecutor(const uint32_t modelId)
{
    const std::lock_guard<std::mutex> locker(mutex_);
    const auto iter = executorMap_.find(modelId);
    if (iter == executorMap_.end()) {
        ACL_LOG_ERROR("model is not loaded, modelId is %u", modelId);
        return static_cast<aclError>(ACL_ERROR_GE_EXEC_MODEL_ID_INVALID);
    }
    (void)executorMap_.erase(iter);

    const auto it = rtSessionMap_.find(modelId);
    if (it == rtSessionMap_.end()) {
        ACL_LOG_ERROR("model is not loaded, modelId is %u", modelId);
        return static_cast<aclError>(ACL_ERROR_GE_EXEC_MODEL_ID_INVALID);
    }
    (void)rtSessionMap_.erase(it);
    return ACL_SUCCESS;
}


void AclResourceManager::GetRuntimeV2Env()
{
    const char_t *enableRuntimeV2Flag = nullptr;
    MM_SYS_GET_ENV(MM_ENV_ENABLE_RUNTIME_V2, enableRuntimeV2Flag);
    if (enableRuntimeV2Flag != nullptr) {
        if (enableRuntimeV2Flag[0] == '0') { // 0 both model and singleOp disable
            enableRuntimeV2ForModel_ = false;
            enableRuntimeV2ForSingleOp_ = false;
        } else if (enableRuntimeV2Flag[0] == '2') { // 2: model enable, singleOp disable
            enableRuntimeV2ForModel_ = true;
            enableRuntimeV2ForSingleOp_ = false;
        } else {
            enableRuntimeV2ForModel_ = true;
            enableRuntimeV2ForSingleOp_ = true;
        }
    }
    ACL_LOG_EVENT("runtime v2 flag : model flag = %d, singleOp flag = %d",
                  static_cast<int32_t>(enableRuntimeV2ForModel_),
                  static_cast<int32_t>(enableRuntimeV2ForSingleOp_));
}

std::shared_ptr<gert::ModelV2Executor> AclResourceManager::GetExecutor(const uint32_t modelId)
{
    const std::lock_guard<std::mutex> locker(mutex_);
    const auto iter = executorMap_.find(modelId);
    if (iter == executorMap_.end()) {
        return nullptr;
    }
    return iter->second;
}


std::shared_ptr<gert::RtSession> AclResourceManager::CreateRtSession()
{
    const std::lock_guard<std::mutex> locker(mutex_);
    ++sessionIdGenerator_;
    auto sessionId = sessionIdGenerator_.load();
    return std::make_shared<gert::RtSession>(sessionId);
}

std::shared_ptr<gert::RtSession> AclResourceManager::GetRtSession(const uint32_t rtSessionId)
{
    const std::lock_guard<std::mutex> locker(mutex_);
    const auto iter = rtSessionMap_.find(rtSessionId);
    if (iter == rtSessionMap_.cend()) {
        return nullptr;
    }
    return iter->second;
}

void *AclResourceManager::GetKeyByStreamOrDefaultStream(const aclrtStream stream)
{
    if (stream != nullptr) {
        return stream;
    }
    // get current context default stream
    aclrtStream curCtxDefaultStream = nullptr;
    const aclError aclErr = aclrtCtxGetCurrentDefaultStream(&curCtxDefaultStream);
    if (aclErr != ACL_ERROR_NONE) {
        ACL_LOG_CALL_ERROR("get current default stream failed, ret:%d", static_cast<int32_t>(aclErr));
        return nullptr;
    }
    return curCtxDefaultStream;
}

std::shared_ptr<gert::Allocators> AclResourceManager::GetAllocators(const aclrtStream stream, bool createDefaultAllocator)
{
    void *cacheKey = stream;
    if (cacheKey == nullptr) {
        cacheKey = GetKeyByStreamOrDefaultStream(stream);
    }

    if (cacheKey == nullptr) {
        return nullptr;
    }

    // If the user does not initially register an external allocator, then create a default allocator
    // and use the default allocator throughout the entire lifecycle of the stream (even if the user
    // later registers an external allocator for the stream).
    const std::unique_lock<std::recursive_mutex> lk(streamAllocatorMutex_);
    // Try using exist default allocator.
    const auto iter = streamDefaultAllocator_.find(cacheKey);
    if (iter != streamDefaultAllocator_.end()) {
        return iter->second;
    }
    // Try using exist or create external allocator.
    std::shared_ptr<gert::Allocators> externalAllocators = UpdateExternalAllocators(stream);
    if (externalAllocators != nullptr) {
        return externalAllocators;
    }
    // For model inference (dynamic shape), the default Allocator is not created to improve performance.
    if (!createDefaultAllocator) {
        ACL_LOG_INFO("The default deviceAllocator is not created.");
        return nullptr;
    }
    return CreateDefaultAllocators(cacheKey);
}

std::shared_ptr<gert::Allocators> AclResourceManager::CreateAllocators(std::shared_ptr<ge::Allocator> &deviceAllocator)
{
    std::shared_ptr<ge::Allocator> hostAllocator(gert::AllocatorFactory::Create(gert::kOnHost).release());
    if ((hostAllocator == nullptr)) {
        ACL_LOG_ERROR("hostAllocator is nullptr");
        return nullptr;
    }
    // create a new alloctors
    std::shared_ptr<gert::Allocators> allocators(new (std::nothrow) gert::Allocators());
    ACL_REQUIRES_NOT_NULL_RET_NULL(allocators);
    // only support allocator with placement kOnDeviceHbm, kOnHost, kFollowing
    for (size_t i = 0U; i < static_cast<size_t>(gert::AllocatorUsage::kEnd); ++i) {
        (void)allocators->SetAllocator(static_cast<gert::TensorPlacement>(gert::kOnDeviceHbm), i, deviceAllocator);
        (void)allocators->SetAllocator(static_cast<gert::TensorPlacement>(gert::kOnHost), i, hostAllocator);
        (void)allocators->SetAllocator(static_cast<gert::TensorPlacement>(gert::kFollowing), i, hostAllocator);
    }
    return allocators;
}

std::shared_ptr<gert::Allocators> AclResourceManager::CreateDefaultAllocators(const void * const cacheKey)
{
    std::shared_ptr<ge::Allocator> deviceAllocator(
        gert::AllocatorFactory::Create(gert::kOnDeviceHbm).release());
    if (deviceAllocator == nullptr) {
        return nullptr;
    }
    std::shared_ptr<gert::Allocators> allocators = CreateAllocators(deviceAllocator);
    ACL_REQUIRES_NOT_NULL_RET_NULL(allocators);
    (void)streamDefaultAllocator_.insert({cacheKey, allocators});
    return allocators;
}

std::shared_ptr<gert::Allocators> AclResourceManager::CreateExternalAllocators(const void * const cacheKey, ExternalAllocatorDesc &allocatorDesc)
{
    ge::AllocatorDesc geAllocatorDesc;
    geAllocatorDesc.obj = allocatorDesc.obj;
    geAllocatorDesc.alloc_func = allocatorDesc.allocFunc;
    geAllocatorDesc.free_func = allocatorDesc.freeFunc;
    geAllocatorDesc.alloc_advise_func = allocatorDesc.allocAdviseFunc;
    geAllocatorDesc.get_addr_from_block_func = allocatorDesc.getAddrFromBlockFunc;
    std::shared_ptr<ge::Allocator> deviceAllocator(gert::CreateExternalAllocator(&geAllocatorDesc).release());
    if (deviceAllocator == nullptr) {
        return nullptr;
    }
    std::shared_ptr<gert::Allocators> allocators = CreateAllocators(deviceAllocator);
    ACL_REQUIRES_NOT_NULL_RET_NULL(allocators);
    (void)streamExternalAllocator_.insert({cacheKey, std::make_pair(allocatorDesc, allocators)});
    return allocators;
}

void AclResourceManager::CleanAllocators(const void * const cacheKey)
{
    if (cacheKey == nullptr) {
        return;
    }
    const std::unique_lock<std::recursive_mutex> lk(streamAllocatorMutex_);
    (void)streamDefaultAllocator_.erase(cacheKey);
    (void)streamExternalAllocator_.erase(cacheKey);
}

std::shared_ptr<gert::Allocators> AclResourceManager::UpdateExternalAllocators(aclrtStream stream)
{
    void *cacheKey = stream;
    aclrtAllocatorDesc new_desc = nullptr;
    aclrtAllocator allocator = nullptr;
    aclrtAllocatorAllocFunc allocFunc = nullptr;
    aclrtAllocatorFreeFunc freeFunc = nullptr;
    aclrtAllocatorAllocAdviseFunc allocAdviseFunc = nullptr;
    aclrtAllocatorGetAddrFromBlockFunc getAddrFromBlockFunc = nullptr;
    bool new_desc_exist = aclrtAllocatorGetByStream(stream, &new_desc, &allocator, &allocFunc,
                            &freeFunc, &allocAdviseFunc, &getAddrFromBlockFunc) == ACL_SUCCESS;
    const auto iter_old_desc = streamExternalAllocator_.find(cacheKey);
    bool old_desc_exist = iter_old_desc != streamExternalAllocator_.end();
    // "old_desc_exist" indicates whether the external allocator stored by streamExternalAllocator_ exists (referred
    // to as the old allocator desc). "new_desc_exist" indicates whether the user has newly registered the desc
    // (referred to as the new "allocator desc"). Based on the old allocator description and whether the user has
    // newly registered the allocator description, there are four possible scenarios.
    // Scenario 1: If the user has neither registered an external allocator before nor currently, simply return nullptr.
    if (!old_desc_exist && !new_desc_exist) {
        return nullptr;
    }
    // Situation 2: The user has newly registered an allocator desc and has never registered before. We will create an
    // external allocator for the user.
    if (!old_desc_exist && new_desc_exist) {
        ExternalAllocatorDesc allocatorDesc = ExternalAllocatorDesc(allocator, allocFunc, freeFunc, allocAdviseFunc, getAddrFromBlockFunc);
        return CreateExternalAllocators(cacheKey, allocatorDesc);
    }
    // Situation 3: The user previously registered for "allocator desc", but in their most recent query, they have
    // unregistered "allocator desc". We delete the previously created external allocator for the user.
    if (old_desc_exist && !new_desc_exist) {
        (void)streamExternalAllocator_.erase(cacheKey);
        return nullptr;
    }
    // Situation 4: The user has previously registered with the "allocator desc" and the latest query has confirmed
    // the user's registration. We compare the latest allocator description retrieved with the previously registered
    // ones. If they are equal, it indicates that the user has not registered a new allocator description.In this case,
    // we can still use the previously saved allocator. If they are not equal, it indicates that the user has registered
    // a new allocator description to replace the old one. We need to delete the old allocator and create a new one for
    // the user.
    if (old_desc_exist && new_desc_exist) {
        ExternalAllocatorDesc allocatorDesc = ExternalAllocatorDesc(allocator, allocFunc, freeFunc, allocAdviseFunc, getAddrFromBlockFunc);
        if (allocatorDesc == iter_old_desc->second.first) {
            return iter_old_desc->second.second;
        }
        (void)streamExternalAllocator_.erase(cacheKey);
        return CreateExternalAllocators(cacheKey, allocatorDesc);
    }
    return nullptr;
}

AclResourceManager::~AclResourceManager()
{
    // note: op_model的executor的释放依赖allocator，所以先释放executor，再释放allocator
    streamDefaultAllocator_.clear();
    streamExternalAllocator_.clear();
}

void AclResourceManager::HandleReleaseSourceByDevice(int32_t deviceId, aclrtDeviceState state, void *args) const
{
    (void)args;
    ACL_LOG_INFO("start to execute HandleReleaseSourceByDevice, devId:%d.", deviceId);
    if (state != ACL_RT_DEVICE_STATE_RESET_PRE) {
        ACL_LOG_INFO("it's not reset pre device callback, currently do nothing.");
        return;
    }
    (void)ge::GeExecutor::ReleaseResource();
    ACL_LOG_INFO("successfully execute HandleReleaseSourceByDevice, devId:%d.", deviceId);
}

void AclResourceManager::HandleReleaseSourceByStream(aclrtStream stream, aclrtStreamState state, void *args)
{
    (void)args;
    ACL_LOG_INFO("start to execute HandleReleaseSourceByStream.");
    if (state != ACL_RT_STREAM_STATE_DESTROY_PRE) {
        ACL_LOG_INFO("it's not destroy stream callback, currently do nothing.");
        return;
    }
    (void)CleanAllocators(stream);
    ACL_LOG_INFO("successfully execute HandleReleaseSourceByStream.");
}

AclResourceManager &AclResourceManager::GetInstance()
{
  static AclResourceManager instance;
  return instance;
}
}