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

#ifndef __HLAI_MODEL_MANAGER_H__
#define __HLAI_MODEL_MANAGER_H__

#include <cstddef>

// define the related macro to shield the related function symbols
#ifndef HLAI_AP_VISIBILITY
    #if !(defined _WIN32) && !(defined __CYGWIN__)
        #define HLAI_AP_VISIBILITY __attribute__ ((visibility ("default")))
    #else
        #define HLAI_AP_VISIBILITY
    #endif
#endif

#define HLAI_API_EXPORT HLAI_AP_VISIBILITY

/**
This is the HLAI ModelManager C API:
*/
#ifdef __cplusplus
extern "C" {
#endif

typedef enum {
    HLAI_MODELTYPE_ONLINE = 0,
    HLAI_MODELTYPE_OFFLINE
} HLAI_ModelType;

typedef enum {
    HLAI_FRAMEWORK_NONE = 0,
    HLAI_FRAMEWORK_TENSORFLOW,
    HLAI_FRAMEWORK_KALDI,
    HLAI_FRAMEWORK_CAFFE,
    HLAI_FRAMEWORK_TENSORFLOW_8BIT,
    HLAI_FRAMEWORK_CAFFE_8BIT,
    HLAI_FRAMEWORK_OFFLINE,
    HLAI_FRAMEWORK_IR,
    HLAI_FRAMEWORK_INVALID,
} HLAI_Framework;

typedef enum {
    HLAI_DEVPERF_UNSET = 0,
    HLAI_DEVPREF_LOW,
    HLAI_DEVPREF_NORMAL,
    HLAI_DEVPREF_HIGH,
} HLAI_DevPerf;

typedef enum {
    HLAI_DATATYPE_UINT8 = 0,
    HLAI_DATATYPE_FLOAT32 = 1,
    HLAI_DATATYPE_FLOAT16 = 2,
    HLAI_DATATYPE_INT32 = 3,
} HLAI_DataType;

typedef enum {
    HLAI_YUV420SP_U8 = 0,
    HLAI_XRGB8888_U8,
    HLAI_YUV400_U8,
    HLAI_ARGB8888_U8,
    HLAI_YUYV_U8,
    HLAI_YUV422SP_U8,
    HLAI_AYUV444_U8,
} HLAI_ImageFormat;

typedef enum {
    HLAI_DYNAMIC_BATCH = 0,
    HLAI_DYNAMIC_HW = 1,
    HLAI_NO_DYNAMIC = 2,
} HLAI_DynamicType;

#define HLAI_DEVPERF_DEFAULT (HLAI_DEVPREF_HIGH)

typedef struct {
    const char* modelNetName;

    const char* modelNetPath;
    bool isModelNetEncrypted;
    const char* modelNetKey;

    const char* modelNetParamPath;
    bool isModelNetParamEncrypted;
    const char* modelNetParamKey;

    HLAI_ModelType modelType;
    HLAI_Framework frameworkType;
    HLAI_DevPerf   perf;
} HLAI_ModelDescription;

#define HLAI_MODEL_DESCRIPTION_INIT { \
    "", \
    "", false, "", \
    "", false, "", \
    HLAI_MODELTYPE_OFFLINE, \
    HLAI_FRAMEWORK_CAFFE, \
    HLAI_DEVPERF_DEFAULT \
}

typedef struct {
    int number;
    int channel;
    int height;
    int width;
    HLAI_DataType dataType;  /* optional */
} HLAI_TensorDescription;

struct HLAI_MemBuffer {
    unsigned int size;
    void* data;
};

#define HLAI_TENSOR_DESCRIPTION_INIT {0, 0, 0, 0, HLAI_DATATYPE_FLOAT32}

typedef struct HLAI_ModelManagerListener_struct {
    void (*onLoadDone)(void* userdata, int taskStamp);
    void (*onRunDone)(void* userdata, int taskStamp);
    void (*onUnloadDone)(void* userdata, int taskStamp);
    void (*onTimeout)(void* userdata, int taskStamp);
    void (*onError)(void* userdata, int taskStamp, int errCode);
    void (*onServiceDied)(void* userdata);

    void* userdata;
} HLAI_ModelManagerListener;

#define HLAI_MODEL_MANAGER_LISTENER_INIT {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr}

struct HLAI_TensorBuffer {
    HLAI_TensorDescription desc;
    int size;
    void* data;
    void* impl; /* DON'T MODIFY */
    const char* name;
};

struct HLAI_DynamicPara {
    int batch;
    int height;
    int width;
    HLAI_DynamicType dynamicType;
};

typedef struct HLAI_ModelBuffer HLAI_ModelBuffer;

HLAI_API_EXPORT HLAI_ModelBuffer* HLAI_ModelBuffer_create_from_file(
    const char* name, const char* path, HLAI_DevPerf perf);
HLAI_API_EXPORT HLAI_ModelBuffer* HLAI_ModelBuffer_create_from_buffer(
    const char* name, void* modelBuf, int size, HLAI_DevPerf perf);

HLAI_API_EXPORT const char* HLAI_ModelBuffer_getName(HLAI_ModelBuffer* b);
HLAI_API_EXPORT const char* HLAI_ModelBuffer_getPath(HLAI_ModelBuffer* b);
HLAI_API_EXPORT HLAI_DevPerf HLAI_ModelBuffer_getPerf(HLAI_ModelBuffer* b);

HLAI_API_EXPORT int HLAI_ModelBuffer_getSize(HLAI_ModelBuffer* b);

HLAI_API_EXPORT void HLAI_ModelBuffer_destroy(HLAI_ModelBuffer* b);

typedef struct HLAI_TensorBuffer HLAI_TensorBuffer;

HLAI_API_EXPORT HLAI_TensorBuffer* HLAI_TensorBuffer_create(int n, int c, int h, int w);
HLAI_API_EXPORT HLAI_TensorBuffer* HLAI_TensorBuffer_CreateWithDataType(
    int n, int c, int h, int w, HLAI_DataType DataType);
HLAI_API_EXPORT HLAI_TensorBuffer* HLAI_CreateByFd(int fd, size_t offset, size_t size);
HLAI_API_EXPORT HLAI_TensorBuffer* HLAI_ImageBuffer_create(int n, int h, int w, HLAI_ImageFormat imageFormat);

HLAI_API_EXPORT HLAI_TensorBuffer* HLAI_TensorBuffer_createFromTensorDesc(HLAI_TensorDescription* tensor);

HLAI_API_EXPORT HLAI_TensorDescription HLAI_TensorBuffer_getTensorDesc(HLAI_TensorBuffer* b);

HLAI_API_EXPORT void* HLAI_TensorBuffer_getRawBuffer(HLAI_TensorBuffer* b);

HLAI_API_EXPORT int HLAI_TensorBuffer_getBufferSize(HLAI_TensorBuffer* b);

HLAI_API_EXPORT void HLAI_TensorBuffer_destroy(HLAI_TensorBuffer* b);

HLAI_API_EXPORT HLAI_TensorBuffer* HLAI_TensorBuffer_create_v2(int n, int c, int h, int w, int size);

typedef struct {
    int input_cnt;
    int output_cnt;
    int *input_shape;
    int *output_shape;
} HLAI_ModelTensorInfo;

typedef struct {
    const char* name;
} HLAI_TensorName;

typedef struct {
    int inputCnt;
    int outputCnt;
    HLAI_TensorName* inputName;
    HLAI_TensorName* outputName;
} HLAI_ModelTensorName;

typedef struct {
    int batchCount;
    int* batch;
} HLAI_DynamicBatch;

typedef struct {
    int hwCount;
    int* height;
    int* width;
} HLAI_DynamicHW;

typedef struct HLAI_ModelManager HLAI_ModelManager;

HLAI_API_EXPORT HLAI_ModelTensorInfo* HLAI_ModelManager_getModelTensorInfo(
    HLAI_ModelManager* manager, const char* modelName);

HLAI_API_EXPORT void HLAI_ModelManager_releaseModelTensorInfo(HLAI_ModelTensorInfo* modelTensor);
typedef enum {
    HLAI_IO_INPUT    = 0,
    HLAI_IO_OUTPUT   = 1,

    HLAI_IO_INVALID  = 255,
} HLAI_IO_TYPE;

typedef struct HLAI_TensorDescriptionV2 HLAI_TensorDescriptionV2;
typedef struct HLAI_ModelTensorInfoV2 HLAI_ModelTensorInfoV2;

HLAI_API_EXPORT HLAI_ModelTensorInfoV2* HLAI_ModelManager_getModelTensorInfoV2(
    HLAI_ModelManager* manager, const char* modelName); // 查询模型输入输出V2版
HLAI_API_EXPORT void HLAI_ModelManager_releaseModelTensorInfoV2(HLAI_ModelTensorInfoV2* modelTensor); // C语言需要释放接口

HLAI_API_EXPORT HLAI_ModelTensorName* HLAI_ModelManager_getModelTensorName(
    HLAI_ModelManager* manager, const char* modelName); // 查询模型输入输出名称
HLAI_API_EXPORT void HLAI_ModelManager_releaseModelTensorName(HLAI_ModelTensorName* modelTensorName);

HLAI_API_EXPORT HLAI_DynamicBatch* HLAI_ModelManager_getDynamicBatch(HLAI_ModelManager* manager, const char* modelName);
HLAI_API_EXPORT void HLAI_ModelManager_releaseDynamicBatch(HLAI_DynamicBatch* info);

HLAI_API_EXPORT HLAI_DynamicHW* HLAI_ModelManager_getDynamicHW(HLAI_ModelManager* manager, const char* modelName);
HLAI_API_EXPORT void HLAI_ModelManager_releaseDynamicHW(HLAI_DynamicHW* info);

HLAI_API_EXPORT int  HLAI_ModelTensorInfoV2_getIOCount(
    const HLAI_ModelTensorInfoV2* tensorInfo, HLAI_IO_TYPE type); // 获取输入或输出的个数
HLAI_API_EXPORT HLAI_TensorDescriptionV2*  HLAI_ModelTensorInfoV2_getTensorDescription(
    const HLAI_ModelTensorInfoV2* tensorInfo, HLAI_IO_TYPE type, int index); // 获取某个输入或输出的信息

HLAI_API_EXPORT HLAI_TensorBuffer* HLAI_TensorBuffer_createFromTensorDescV2(HLAI_TensorDescriptionV2* tensor);

HLAI_API_EXPORT int HLAI_TensorDescriptionV2_getDimensions(
    const HLAI_TensorDescriptionV2* tensorDesc, int *n, int *c, int *h, int *w); // 获取某个tensor的nchw
HLAI_API_EXPORT const char* HLAI_TensorDescriptionV2_getName(
    const HLAI_TensorDescriptionV2* tensorDesc); // 获取某个tensor的名称

#ifdef DYNAMIC_BATCH
HLAI_API_EXPORT int HLAI_TensorBuffer_setDynamicBatch(
    HLAI_TensorBuffer* input[], int nInput, int inputIndex, int batchNumber);
#endif

typedef struct HLAI_TensorAippPara HLAI_TensorAippPara;

// aipp related
HLAI_API_EXPORT HLAI_TensorAippPara* HLAI_TensorAipp_create(unsigned int batchCount);

HLAI_API_EXPORT void* HLAI_TensorAipp_getRawBuffer(HLAI_TensorAippPara* tensorAipp);

HLAI_API_EXPORT int HLAI_TensorAipp_getInputIndex(HLAI_TensorAippPara* tensorAipp);
HLAI_API_EXPORT void HLAI_TensorAipp_setInputIndex(HLAI_TensorAippPara* tensorAipp, int inputIndex);

HLAI_API_EXPORT int HLAI_TensorAipp_getInputAippIndex(HLAI_TensorAippPara* tensorAipp);
HLAI_API_EXPORT void HLAI_TensorAipp_setInputAippIndex(HLAI_TensorAippPara* tensorAipp, int inputAippIndex);

HLAI_API_EXPORT void HLAI_TensorAipp_destroy(HLAI_TensorAippPara* aippParas);

HLAI_API_EXPORT int HLAI_ModelManager_runAippModel(HLAI_ModelManager* manager, HLAI_TensorBuffer* input[], int nInput,
    HLAI_TensorAippPara* aipp[], int nAipp, HLAI_TensorBuffer* output[], int nOutput, int ulTimeout,
    const char* modelName);

HLAI_API_EXPORT HLAI_ModelManager* HLAI_ModelManager_create(HLAI_ModelManagerListener* listener);

HLAI_API_EXPORT void HLAI_ModelManager_destroy(HLAI_ModelManager* manager);

/**
 * @return >0 means load success, and the return value is a taskStamp, < 0 means error occured
**/
HLAI_API_EXPORT int HLAI_ModelManager_loadFromModelBuffers(
    HLAI_ModelManager* manager, HLAI_ModelBuffer* bufferArray[], int nBuffers);

HLAI_API_EXPORT int HLAI_ModelManager_loadFromModelDescriptions(
    HLAI_ModelManager* manager, HLAI_ModelDescription descsArray[], int nDescs);

HLAI_API_EXPORT bool HLAI_ModelManager_containDangerousOp(HLAI_ModelManager* manager, const char* modelName);

HLAI_API_EXPORT int HLAI_ModelManager_runModel(HLAI_ModelManager* manager, HLAI_TensorBuffer* input[],
    int nInput, HLAI_TensorBuffer* output[], int nOutput, int ulTimeout, const char* modelName);

HLAI_API_EXPORT int HLAI_ModelManager_setInputsAndOutputs(HLAI_ModelManager* manager,
    const char* modelname, HLAI_TensorBuffer* input[], int nInput, HLAI_TensorBuffer* output[], int nOutput);

HLAI_API_EXPORT int HLAI_ModelManager_startCompute(HLAI_ModelManager* manager, const char* modelname);

HLAI_API_EXPORT int HLAI_ModelManager_unloadModel(HLAI_ModelManager* manager);
HLAI_API_EXPORT char* HLAI_GetVersion();
HLAI_API_EXPORT char* HLAI_GetApiVersion();

HLAI_API_EXPORT int HLAI_ModelManager_runDynamicModel(HLAI_ModelManager* manager, HLAI_TensorBuffer* input[],
    int nIput, HLAI_TensorBuffer* output[], int nOutput,
    int ulTimeout, const char* modelName, HLAI_DynamicPara dynamicPara);

HLAI_API_EXPORT int HLAI_ModelManager_runAippDynamicModel(HLAI_ModelManager* manager, HLAI_TensorBuffer* input[],
    int nInput, HLAI_TensorAippPara* aipp[], int nAipp, HLAI_TensorBuffer* output[], int nOutput, int ulTimeout,
    const char* modelName, HLAI_DynamicPara dynamicPara);

// Description
//      read an onlineModel file into a MemBuffer
// Params
//      @path: an onlineModel file. (eg: xxx.prototxt / xxx.caffemodel / tf_param.txt)
HLAI_API_EXPORT HLAI_MemBuffer* HLAI_MemBuffer_create_from_file(const char *path);


// Description
//      read an onlineModel buffer into a MemBuffer
// Params
//      @buffer: a buffer stores an onlineModel file content.
//      @size: buffer size.
HLAI_API_EXPORT HLAI_MemBuffer* HLAI_MemBuffer_create_from_buffer(void* buffer, const unsigned int size);

// Description
//      read a tensorflow pb file, covnert it to a GraphDef-format MemBuffer
// Params
//      @path: tensorflow pb file. (eg: xxx.pb)
HLAI_API_EXPORT HLAI_MemBuffer* HLAI_ReadBinaryProto_from_file(const char *path);

// Description
//      read a tensorflow pb buffer, covnert it to a GraphDef-format MemBuffer
// Params
//      @buffer: a buffer stores a tensorflow pb file content.
//      @size: buffer size
HLAI_API_EXPORT HLAI_MemBuffer* HLAI_ReadBinaryProto_from_buffer(void* buffer, const unsigned int size);

// Description
//      create an offlineModel MemBuffer through an onlineModel
// Params
//      @frameworkType: frameworkType. (eg: CAFFE / TENSORFLOW)
//      @inputModelBuffers: a full-onlineModel
//      @inputModelBuffersNum: the number of inputModelBuffers
HLAI_API_EXPORT HLAI_MemBuffer* HLAI_MemBuffer_create(HLAI_Framework frameworkType,
    HLAI_MemBuffer* inputModelBuffers[], const unsigned int inputModelBuffersNum);

// Description
//      destroy a MemBuffer
// Params
//      @membuf: MemBuffer
HLAI_API_EXPORT void HLAI_MemBuffer_destroy(HLAI_MemBuffer* membuf);


// Description
//      export a offlineModel file through offlineModel MemBuffer
// Params
//      @membuf: offlineModel MemBuffer
//      @buildSize: the actual offlineModel size.
//      @buildPath: the path of to-export offlineModel.
HLAI_API_EXPORT bool HLAI_MemBuffer_export_file(HLAI_MemBuffer* membuf,
    const unsigned int buildSize, const char *buildPath);

// Description
//      read an offlineModel file, check the compatibility.
// Params
//      @manager: modelManager instance
//      @modelPath: the path of the offlineModel file.
HLAI_API_EXPORT bool HLAI_CheckModelCompatibility_from_file(
    HLAI_ModelManager* manager, const char *modelPath);

// Description
//      read an offlineModel buffer, check the compatibility.
// Params
//      @manager: modelManager instance
//      @buffer: a buffer stores an offlineModel file content.
//      @size: buffer size.
HLAI_API_EXPORT bool HLAI_CheckModelCompatibility_from_buffer(HLAI_ModelManager* manager,
    void* buffer, const unsigned int size);

#ifdef __cplusplus
}  // extern "C"
#endif

#endif  // __HLAI_MODEL_MANAGER_H__