#include "model_process.h"
#include <iostream>
#include "utils.h"
#include "sy_errorinfo.h"
#include <cstring>

using namespace std;

ModelProcess::ModelProcess():loadFlag_(false), modelId_(0), modelMemPtr_(nullptr), modelMemSize_(0),
modelWeightPtr_(nullptr),modelWeightSize_(0), modelDesc_(nullptr), input_(nullptr), output_(nullptr) {
}

ModelProcess::~ModelProcess() {
    Unload();
    DestroyDesc();
    DestroyInput();
    DestroyOutput();
}

int ModelProcess::LoadModelFromFileWithMem(const char *modelPath) {
    if (loadFlag_) {
        return SY_FAILED;
    }

    ACL_CALL(aclmdlQuerySize(modelPath, &modelMemSize_, &modelWeightSize_),
        ACL_SUCCESS, SY_FAILED);
    
    ACL_CALL(aclrtMalloc(&modelMemPtr_, modelMemSize_, ACL_MEM_MALLOC_HUGE_FIRST),
        ACL_SUCCESS, SY_FAILED);

    ACL_CALL(aclrtMalloc(&modelWeightPtr_, modelWeightSize_, ACL_MEM_MALLOC_HUGE_FIRST),
        ACL_SUCCESS, SY_FAILED);

    ACL_CALL(aclmdlLoadFromFileWithMem(modelPath, &modelId_, modelMemPtr_,
        modelMemSize_, modelWeightPtr_, modelWeightSize_), ACL_SUCCESS, SY_FAILED);

    loadFlag_ = true;

    return SY_SUCCESS;
}

int ModelProcess::CreateDesc() {
    modelDesc_ = aclmdlCreateDesc();
    if (modelDesc_ == nullptr) {
        return SY_FAILED;
    }

    ACL_CALL(aclmdlGetDesc(modelDesc_, modelId_), ACL_SUCCESS, SY_FAILED);

    return SY_SUCCESS;
}

void ModelProcess::DestroyDesc() {
    if (modelDesc_ != nullptr) {
        (void)aclmdlDestroyDesc(modelDesc_);
        modelDesc_ = nullptr;
    }
}

int ModelProcess::CreateOutput(vector<vector<int>>& dims) {
    if (modelDesc_ == nullptr) {
        return SY_FAILED;
    }

    output_ = aclmdlCreateDataset();
    if (output_ == nullptr) {
        return SY_FAILED;
    }

    aclError ret;
    size_t outputSize = aclmdlGetNumOutputs(modelDesc_);
    for (size_t i = 0; i < outputSize; ++i) {
        size_t buffer_size = aclmdlGetOutputSizeByIndex(modelDesc_, i);

        void *outputBuffer = nullptr;
        ACL_CALL(aclrtMalloc(&outputBuffer, buffer_size, ACL_MEM_MALLOC_NORMAL_ONLY),
            ACL_SUCCESS, SY_FAILED);

        aclDataBuffer* outputData = aclCreateDataBuffer(outputBuffer, buffer_size);
        if (outputData == nullptr) {
            aclrtFree(outputBuffer);
            return SY_FAILED;
        }

        ret = aclmdlAddDatasetBuffer(output_, outputData);
        if (ret != ACL_SUCCESS) {
            aclrtFree(outputBuffer);
            aclDestroyDataBuffer(outputData);
            return SY_FAILED;
        }

        // get output dims
        aclmdlIODims outDims;
        vector<int> dims_;
        ACL_CALL(aclmdlGetOutputDims(modelDesc_, i, &outDims), ACL_SUCCESS, SY_FAILED);
        for (int j = 0; j < outDims.dimCount; j++) {
            dims_.emplace_back(outDims.dims[j]);
        }
        dims.emplace_back(dims_);

    }

    return SY_SUCCESS;
}

void ModelProcess::DestroyOutput() {
    if (output_ == nullptr) {
        return;
    }

    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(output_); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(output_, i);
        void* data = aclGetDataBufferAddr(dataBuffer);
        (void)aclrtFree(data);
        (void)aclDestroyDataBuffer(dataBuffer);
        dataBuffer = nullptr;
    }

    (void)aclmdlDestroyDataset(output_);
    output_ = nullptr;
}

// get input dims
int ModelProcess::GetInputDims(vector<vector<int>>& dims) {
    if (modelDesc_ == nullptr) {
        return SY_FAILED;
    }

    aclError ret;
    size_t inputSize = aclmdlGetNumInputs(modelDesc_);
    //debug===================================
    // uint32_t modelInputSize = aclmdlGetInputSizeByIndex(modelDesc_, 0);
    // cout << "modelInputSize:" << modelInputSize << endl;
    //debug end===============================
    for (size_t i = 0; i < inputSize; ++i) {
        // get output dims
        aclmdlIODims inDims;
        vector<int> dims_;
        ACL_CALL(aclmdlGetInputDims(modelDesc_, i, &inDims), ACL_SUCCESS, SY_FAILED);
        for (int j = 0; j < inDims.dimCount; j++) {
            dims_.emplace_back(inDims.dims[j]);
        }
        dims.emplace_back(dims_);
    }

    return SY_SUCCESS;
}

int ModelProcess::CreateInput(void *input, size_t inputsize) {
    input_ = aclmdlCreateDataset();
    if (input_ == nullptr) {
        return SY_FAILED;
    }

    aclDataBuffer* inputData = aclCreateDataBuffer(input, inputsize);
    if (inputData == nullptr) {
        return SY_FAILED;
    }

    aclmdlAddDatasetBuffer(input_, inputData);
    if (inputData == nullptr) {
        aclDestroyDataBuffer(inputData);
        inputData = nullptr;
        return SY_FAILED;
    }

    return SY_SUCCESS;
}

int ModelProcess::CreateInputV2(void *input, size_t inputsize) {
    input_ = aclmdlCreateDataset();
    if (input_ == nullptr) {
        return SY_FAILED;
    }
    aclError ret;
    for(size_t index = 0; index < aclmdlGetNumInputs(modelDesc_); ++index)
    {   
        const char* name = aclmdlGetInputNameByIndex(modelDesc_, index);
        size_t inputLen = aclmdlGetInputSizeByIndex(modelDesc_, index);
        if(strcmp(name, ACL_DYNAMIC_TENSOR_NAME) == 0)
        {   
            void *data = nullptr;
            ret = aclrtMalloc(&data, inputLen, ACL_MEM_MALLOC_HUGE_FIRST); 
            aclDataBuffer* batchBuffer = aclCreateDataBuffer(data, inputLen);
            if (batchBuffer == nullptr) {
                (void)aclrtFree(data);
                data = nullptr;
                return SY_FAILED;
            }
            ret = aclmdlAddDatasetBuffer(input_, batchBuffer);
            if (ret != ACL_SUCCESS) {
                ERROR_LOG("add input dataset buffer failed, errorCode = %d.", static_cast<int32_t>(ret));
                aclDestroyDataBuffer(batchBuffer);
                batchBuffer = nullptr;
                (void)aclrtFree(data);
                data = nullptr;
                return SY_FAILED;
            }
        }
        else
        {   
            aclDataBuffer* inputData = aclCreateDataBuffer(input, inputsize);
            if (inputData == nullptr) {
                return SY_FAILED;
            }
            ret = aclmdlAddDatasetBuffer(input_, inputData);
            if (inputData == nullptr) {
                aclDestroyDataBuffer(inputData);
                inputData = nullptr;
                return SY_FAILED;
            }
        }
    }

    return SY_SUCCESS;
}

int ModelProcess::AddInputBuff(void *input, size_t inputsize) {
    aclDataBuffer* inputData = aclCreateDataBuffer(input, inputsize);
    if (inputData == nullptr) {
        return SY_FAILED;
    }

    aclmdlAddDatasetBuffer(input_, inputData);
    if (inputData == nullptr) {
        aclDestroyDataBuffer(inputData);
        inputData = nullptr;
        return SY_FAILED;
    }
    
    return SY_SUCCESS;
}


void ModelProcess::DestroyInput() {
    if (input_ == nullptr) {
        return;
    }

    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(input_); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(input_, i);
        aclDestroyDataBuffer(dataBuffer);
        dataBuffer = nullptr;
    }
    aclmdlDestroyDataset(input_);
    input_ = nullptr;
}

void ModelProcess::DestroyInputV2() {
    if (input_ == nullptr) {
        return;
    }

    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(input_); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(input_, i);
        if (dataBuffer == nullptr) {
            continue;
        }
        void *data = aclGetDataBufferAddr(dataBuffer);
        if (data == nullptr){
            (void)aclDestroyDataBuffer(dataBuffer);
            dataBuffer = nullptr;
            continue;
        }
        const char* name = aclmdlGetInputNameByIndex(modelDesc_, i);
        if(strcmp(name, ACL_DYNAMIC_TENSOR_NAME) == 0) { 
            (void)aclrtFree(data);
            data = nullptr; 
        }
        
        (void)aclDestroyDataBuffer(dataBuffer);
        dataBuffer = nullptr; 
    }
    (void)aclmdlDestroyDataset(input_);
    input_ = nullptr;
}

int ModelProcess::Execute() {
    aclError ret = aclmdlExecute(modelId_, input_, output_);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("execute model failed, modelId is %u %d", modelId_, ret);
        return SY_FAILED;
    }
    // ACL_CALL(aclmdlExecute(modelId_, input_, output_), ACL_SUCCESS, SY_FAILED);

    return SY_SUCCESS;
}

int ModelProcess::Execute(int batchsize) {
    size_t index;
    aclError ret = aclmdlGetInputIndexByName(modelDesc_, ACL_DYNAMIC_TENSOR_NAME, &index);
    ret = aclmdlSetDynamicBatchSize(modelId_, input_, index, batchsize); 
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("aclmdlSetDynamicBatchSize failed, modelId is %u %d", modelId_, ret);
    }
    ret = aclmdlExecute(modelId_, input_, output_);
    if (ret != ACL_SUCCESS) {
        ERROR_LOG("execute model failed, modelId is %u %d", modelId_, ret);
        return SY_FAILED;
    }
    return SY_SUCCESS;
}

void ModelProcess::Unload() {
    if (!loadFlag_) {
        return;
    }

    aclError ret = aclmdlUnload(modelId_);
    if (ret != ACL_SUCCESS) {
        return;
    }

    if (modelDesc_ != nullptr) {
        (void)aclmdlDestroyDesc(modelDesc_);
        modelDesc_ = nullptr;
    }

    if (modelMemPtr_ != nullptr) {
        aclrtFree(modelMemPtr_);
        modelMemPtr_ = nullptr;
        modelMemSize_ = 0;
    }

    if (modelWeightPtr_ != nullptr) {
        aclrtFree(modelWeightPtr_);
        modelWeightPtr_ = nullptr;
        modelWeightSize_ = 0;
    }

    loadFlag_ = false;
}

aclmdlDataset *ModelProcess::GetModelOutputData() {
    return output_;
}