/*
 * Copyright (C) 2021-2024 Advanced Micro Devices, Inc. All rights reserved.
 */

#ifndef RML_AMODEL_API_H_
#define RML_AMODEL_API_H_

#include <string>
#include <map>

namespace amd {
namespace rmlAmodel {

enum EngineType : int8_t {
  EngineType_TVM = 0,
  EngineType_IPU1 = 1,
  EngineType_IPU2 = 2,
  EngineType_IPU3 = 3,
  EngineType_IPU4 = 4,
  EngineType_ONNX = 5,
  EngineType_TENSORFLOW_PB = 6,
  EngineType_AKS = 7,
  EngineType_TVM2 = 8,
  EngineType_TVM3 = 9,
  EngineType_TVM4 = 10,
  EngineType_MIN = EngineType_TVM,
  EngineType_MAX = EngineType_TVM4
};

struct rmlAmodelPkg;

/**
 * RadeonML packager amodel container interface
 */
class rmlAmodel {
  public:
    /**
     * Load amodel from an .amodel container file
     */
    rmlAmodel (const char *amodel_path);

    /**
     * Load amodel from a buffer
     */
    rmlAmodel (const void* buff_ptr, size_t size);

    virtual ~rmlAmodel();

    /**
     * Get metadata map<meta_key, meta_value> list from the amodel container
     * the amodel meta_key should include: "filename", "object_count"
     *     may also include: "signature", "QoS_param", etc
     */
    virtual std::map<std::string, std::string> GetAmodelMetadataList();

    /**
     * Get total count of model objects included in the amodel container
     */
    virtual int GetTotalObjectsCount();

    /**
     * Get total count of buffers in a given model object
     * @param obj_index: model object index in the models vector
     */
    virtual int GetObjectTotalBufferCount(uint32_t obj_index);

    /**
     * Get one model object data from amodel container
     * @param obj_index: model object index in the models vector
     * @param buf_index: data buffer index in the specified model object.
     *                   the buffer at index 0 is the model's data
     */
    virtual const uint8_t* GetObjectDataBuffer(uint32_t obj_index, uint32_t buf_index = 0);

    /**
     * Get size of one model object data from the amodel container
     * @param obj_index: model object index in the models vector
     * @param buf_index: data buffer index in the specified model object
     *                   the buffer at index 0 is the model's data
     */
    virtual int GetObjectDataSize(uint32_t obj_index, uint32_t buf_index = 0);

    /**
     * Get one model object name from the amodel container
     * @param obj_index: model object index in the models vector
     */
    virtual const char* GetObjectName(uint32_t obj_index);

    /**
     * Get target device of one model object name from the amodel container
     * @param obj_index: model object index in the models vector
     */
    virtual EngineType GetObjectTargetDev(uint32_t obj_index);

    /**
     * Get one model object metadata map<meta_key, meta_value> list
     * @param obj_index: model object index in models vector
     * the object meta_key should include: "filename", "target"
     */
    virtual std::map<std::string, std::string> GetObjectMetadataList(uint32_t obj_index);

  private:
    /**
     *  data buffer read from amodel container file
     */
    uint8_t *amodel_buffer;

    /**
     * amodel structure retrieved from an amodel container file
     */
    rmlAmodelPkg *amodel;

    /**
     * metadata map<meta_key, meta_value> list of amodel
     */
    std::map<std::string, std::string> amodelMetadataList;

    /**
     * Generate metadata map<meta_key, meta_value> list
     * @param meta_str: input metadata string
     * @param meta_map: metadata map list pointer
     */
    void GenMetadataList(char* meta_str, std::map<std::string, std::string> *meta_map);
};

}  // namespace rmlAmodel
}  // namespace amd

#endif  // RML_AMODEL_API_H_
