//
//  OnnxMLProcessorNode.hpp
//  SwitchboardOnnx
//
//  Created by Tom Dodds on 24/04/2024.
//

#pragma once

#include <switchboard_onnx/export.h>
#include "ModelInstance.h"
#include <switchboard_core/NodeTypeInfo.hpp>
#include <switchboard_core/SingleBusAudioProcessorNode.hpp>

#include <atomic>
#include <map>
#include <vector>

namespace switchboard::extensions::onnx {

/**
 *  OnnxMLProcessorNode class.
 *  @brief Run inference on the incoming audio using an Onnx model for audio-in audio-out use-case.
 */
class SWITCHBOARDONNX_EXPORT OnnxMLProcessorNode : public SingleBusAudioProcessorNode {
public:
    static NodeTypeInfo getNodeTypeInfo() {
        return NodeTypeInfo {
            "Onnx",
            "OnnxMLProcessorNode",
            "Onnx ML Processor",
            "Runs inference on incoming audio using an Onnx model.",
            { NODE_CATEGORY_AI, NODE_CATEGORY_AUDIO_PROCESSING },
            SBAnyMap({
                { "modelFilePath",
                  SBAnyMap({
                      { CONFIGURATION_FIELD_NAME, "Model File Path" },
                      { CONFIGURATION_FIELD_TYPE, CONFIGURATION_TYPE_STRING },
                      { CONFIGURATION_FIELD_DESCRIPTION, "Relative path of the model file to load" },
                  }) },
            }),
            SBAnyMap({
                { "enabled",
                  SBAnyMap({
                      { PROPERTY_FIELD_DESCRIPTION, "Enable or disable the inference." },
                      { PROPERTY_FIELD_TYPE, PROPERTY_TYPE_BOOLEAN },
                      { PROPERTY_FIELD_READ_ONLY, false },
                      { PROPERTY_FIELD_DEFAULT_VALUE, true },
                  }) },
            }),
            SBAnyMap {},
            SBAnyMap {},
        };
    }

    /**
     *  @brief Creates a OnnxMLProcessorNode instance.
     */
    OnnxMLProcessorNode();

    /**
     * @brief Creates a OnnxMLProcessorNode instance.
     *
     * @param config Configuration map for the node.
     */
    OnnxMLProcessorNode(const SBAnyMap& config);

    /**
     * OnnxMLProcessorNode destructor.
     */
    ~OnnxMLProcessorNode();

    /**
     *  @brief Checks if the inference is enabled.
     *
     *  @returns True if the inference is enabled.
     */
    bool getIsEnabled() const;

    /**
     *  @brief Enables or disables the inference.
     *
     *  @param isEnabled True if the inference is enabled.
     */
    void setIsEnabled(const bool isEnabled);

    /**
     *  @brief Loads Model via string path
     *
     *  @param path String pointer to the model
     */
    ModelInstance* loadModel(std::string path);

    /**
     *  @brief Loads Model via pointer to model data
     *
     *  @param modelData Pointer to the model data
     *  @param modelSize Size of model data in bytes
     */
    void loadModelFromStringPointer(const char* modelData, const unsigned int modelSize);

    /**
     *  @brief A method that should be overriden to provide model data initialization
     *
     *  @param modelInstance Pointer to the ModelInstance
     */
    virtual void initialize(ModelInstance* modelInstance);

    /**
     *  @brief A method that should be overriden and implementation should call predict method on modelIstance
     *
     *  @param audioBufferIn Input AudioBuffer
     *  @param audioBufferOut Output AudioBuffer
     *  @param modelInstance Pointer to the ModelInstance
     */
    virtual void predict(AudioBuffer<float>& audioBufferIn, AudioBuffer<float>& audioBufferOut, ModelInstance* modelInstance);

    std::string getModelFilePath() const;
    void setInputFrameSize(unsigned int size);
    void setInputName(std::string& name);
    void setOutputName(std::string& name);

#pragma mark Overridden methods

    bool setBusFormat(AudioBusFormat& inputBusFormat, AudioBusFormat& outputBusFormat) override;
    bool process(AudioBus& inBus, AudioBus& outBus) override;

private:
    void setModelFilePath(const std::string& modelFilePath);

    std::string modelFilePath;
    std::atomic<bool> isEnabled;
    uint sampleRate;
    uint numberOfFrames;
    std::unique_ptr<ModelInstance> modelInstance;
    std::vector<float> inputBuffer;
    std::vector<float> outputBuffer;

    uint inputFrameSize;
    std::string inputName;
    std::string outputName;
};

}
