/**
 * Copyright (c) Hisilicon Technologies Co., Ltd.. 2020-2020. All rights reserved.
 *
 */

#ifndef __AI_MODEL_MANAGER_IMPL_H__
#define __AI_MODEL_MANAGER_IMPL_H__

#include <utils/Mutex.h>
#include <utils/RefBase.h>
#include <utils/Singleton.h>
#include <vendor/huanglong/hardware/hwhlai/1.0/IAiEngineService.h>
#include <vendor/huanglong/hardware/hwhlai/1.0/IAiMMListenerProxy.h>
#include <vendor/huanglong/hardware/hwhlai/1.0/IAiModelMngr.h>
#include <vendor/huanglong/hardware/hwhlai/1.0/types.h>
#include <hwbinder/IPCThreadState.h>
#include <atomic>
#include "AiConfig.h"
#include "AiMMListenerProxy.h"

static const unsigned MAX_INSTANCE_NUM = 10;

static const int PROCESS_NAME_MAX_LENGTH = 256;
std::string ConcatModelProcessName(std::string modelName);
// ??׮????λ
#define PILING 1
#define DEBUG0 1
#define REMOVE30 1
#define SUPPORT_AIPP 1

namespace ai {
using namespace android;

using ::vendor::huanglong::hardware::hwhlai::V1_0::IAiEngineService;
using ::vendor::huanglong::hardware::hwhlai::V1_0::IAiModelMngr;
using ::vendor::huanglong::hardware::hwhlai::V1_0::IAiMMListenerProxy;
using ::android::hardware::Return;
using ::android::hardware::hidl_death_recipient;
using ::android::hardware::IPCThreadState;
using ::android::hidl::base::V1_0::IBase;
using ::vendor::huanglong::hardware::hwhlai::V1_0::ModelInfo;
using ::vendor::huanglong::hardware::hwhlai::V1_0::ModelDescription;
using ::vendor::huanglong::hardware::hwhlai::V1_0::ModelBuffer;
using ::vendor::huanglong::hardware::hwhlai::V1_0::AippPosition;
using ::vendor::huanglong::hardware::hwhlai::V1_0::DynamicConfig;
using ::vendor::huanglong::hardware::hwhlai::V1_0::TensorDescriptionV2;
using ::vendor::huanglong::hardware::hwhlai::V1_0::DataInfo;
// instaceID是任务标记,如模型兼容性检查,加载运行模型都是独立的instance
// modelID是模型标记,由ge加载模型成功后返回
class AiModelManagerImpl : virtual public RefBase, public Singleton<AiModelManagerImpl> {
public:
    typedef enum {
        NO_INIT = 0,
        INITED,
    } AiStatus;

    class InstanceInfo {
    public:
        InstanceInfo() = default;
        InstanceInfo(int32_t instanceID, bool isSync)
            : mInstanceID(instanceID), mIsSync(isSync)
        {
        }
        ~InstanceInfo() {};

        int32_t getInstanceID()
        {
            return mInstanceID;
        }

        bool isSync()
        {
            return mIsSync;
        }
    private:
        int32_t mInstanceID {0};
        bool mIsSync {true};
    };

    struct TensorAippInfo {
        uint32_t batchCount = 0;
        uint32_t aippNodesCount = 0;
    };

public:
    AiModelManagerImpl();
    ~AiModelManagerImpl();

public:
    bool _register(void* instance);
    bool _register(void* instance, const sp<AiModelManagerListener>& listener);
    void _unregister(void* instance);

public:
    /* QUERY FUNCS */
    int32_t checkModelNameValid(const string& modelName, bool emtpy = false);

    int32_t checkModelDescValid(ModelDescription& modelPara);
    int32_t checkModelDescVecValid(vector<ModelDescription>& modelParaVec);

    int32_t checkModelBufferValid(ModelBuffer& modelBuffer);
    int32_t checkModelBufferVecValid(vector<ModelBuffer>& modelBufferVec);

    AiModelManager::ModelTensorInfo getModelTensor(void* instance, const string& modelName);
    AiModelManager::ModelTensorName getModelTensorName(void* instance, const string& modelName);

    AiModelManager::ModelDynamicBatch getDynamicBatch(void* instance, const string& modelName);
    AiModelManager::ModelDynamicHW getDynamicHW(void* instance, const string& modelName);

    TensorAippInfo GetTensorAippInfo(void* instance, const string& modelName, uint32_t index);
    int32_t GetTensorAippParas(void* instance, const string& modelName, uint32_t index,
                               vector<native_handle_t*>& buffers);

#if REMOVE30

#else
    std::pair<vector<TensorDescriptionV2>, vector<TensorDescriptionV2>> getModelTensorV2(void* instance,
                                                                                        const string& modelName);
#endif
    int32_t getMaxUsedMemory(void* instance);

    bool isServiceDied();
public:
    /* OPERATION FUNCS */
    int32_t startModel(void* instance, int32_t* taskStamp, ModelDescription& modelPara);

    int32_t startModel(void* instance, int32_t* taskStamp, vector<ModelDescription>& modelParaVec);

    int32_t startModel(void* instance, int32_t* taskStamp, ModelBuffer& modelBuffer);

    int32_t startModel(void* instance, int32_t* taskStamp, vector<ModelBuffer>& modelBufferVec);

    int32_t runModel(void* instance, int32_t* taskStamp, const vector<DataInfo> srcDataVec,
        const vector<TensorDescription> srcTensorVec,
        vector<DataInfo> destDataVec, vector<TensorDescription> destTensorVec,
        uint32_t ulTimeout, string modelName = "");

    int32_t RunAippModel(void* instance, int32_t* taskStamp,
        const vector<DataInfo>& srcDataVec, const vector<TensorDescription>& srcTensorVec,
        const vector<native_handle_t*>& aippDataVec, const vector<AippPosition>& aippPostionVec,
        const vector<DataInfo>& destDataVec, const vector<TensorDescription> &destTensorVec,
        uint32_t ulTimeout, string modelName = "");

    int32_t stopModel(void* instance, int32_t* taskStamp);

    int32_t runDynamicModel(void* instance, int32_t* taskStamp, const vector<DataInfo> srcDataVec,
        const vector<TensorDescription> srcTensorVec,
        vector<DataInfo> destDataVec, vector<TensorDescription> destTensorVec,
        uint32_t ulTimeout, DynamicConfig dynamicPara, string modelName = "");

    int32_t RunAippDynamicModel(void* instance, int32_t* taskStamp,
        const vector<DataInfo>& srcDataVec, const vector<TensorDescription>& srcTensorVec,
        const vector<native_handle_t*>& aippDataVec, const vector<AippPosition>& aippPostionVec,
        const vector<DataInfo>& destDataVec, const vector<TensorDescription> &destTensorVec,
        uint32_t ulTimeout, DynamicConfig dynamicPara, string modelName = "");

public:
    /* FOR MICROSOFT */
    int32_t setInputsAndOutputs(void* instance, string modelname,
                                vector<AINeuralNetworkBuffer>& nn_inputs,
                                vector<AINeuralNetworkBuffer>& nn_outputs);

    int32_t startCompute(void* instance, string modelname);

    int32_t buildModel(uint8_t frameworkType,
                       vector<void *> inputModelBuffers, vector<uint32_t> inputModelBufferSizes,
                       void* outputModelBuffer, uint32_t outputModelBufferSize,
                       uint32_t* outModelSize);

    int32_t checkModelValid(void* instance);

    int32_t checkModelCompatibility(void* instance, ModelBuffer& modelBuffer);

    int32_t allocMemory(const void* instance, string config, int size, native_handle_t* &handle);
    int32_t freeMemory(const void* instance, string config, native_handle_t* handle);
    int32_t createByFd(const void* instance, size_t offset, size_t size, native_handle_t* handle);

    int createAshmemFd(const char* name, size_t size);
private:
    void init();
    void deinit();

    int32_t createInstanceID();
    void freeInstanceID(int32_t id);
    void resetInstanceIDPool();
    bool isTensorDescValid(const TensorDescription& tensor);

    int32_t checkDataVecValid(const vector<DataInfo>& dataVec, bool inFlag);
    int32_t checkTensorDescVecValid(const vector<TensorDescription>& tensorVec, bool inFlag);

    int32_t checkNNBufferValid(AINeuralNetworkBuffer& nn_buffer, bool inFlag);
    int32_t checkNNBufferVecValid(vector<AINeuralNetworkBuffer>& nn_bufferVec, bool inFlag);
    int32_t CheckRunModelVecInfoValid(void* instance, int32_t* taskStamp,
        const vector<DataInfo> srcDataVec, const vector<TensorDescription> srcTensorVec,
        vector<DataInfo> destDataVec, vector<TensorDescription> destTensorVec,
        uint32_t ulTimeout, string modelName);

private:
    sp<IAiEngineService> connectAiEngineService();
    void setServiceDiedStatus(bool status);

private:
    android::Mutex mStatusLock;
    android::Mutex mLock;
    bool mIsBinderDied{false};

    class MMRecipientClient : public hidl_death_recipient {
    public:
        MMRecipientClient(AiModelManagerImpl* client) : mClient(client) {}
        ~MMRecipientClient() {}
        virtual void serviceDied(uint64_t cookie, const wp<IBase>& who);

    private:
        wp<AiModelManagerImpl> mClient;
    };

    sp<MMRecipientClient> mMMRecipientClient;
    sp<IAiModelMngr> mIAiModelMngrService;
    sp<IAiEngineService> mAiEngineService;
    sp<AiMMListenerProxy> mMMListenerProxy;

    uint32_t mMaxSizeCheck{0};  // to limit the number of inputs and outputs

    std::atomic<int32_t> mTaskCounter;
    std::map<void*, InstanceInfo> mInstanceMap;
    int32_t mInstanceIDPool[MAX_INSTANCE_NUM * 2];
};
};
#endif  // __AI_MODEL_MANAGER_IMPL_H__