/**************************************************************************************************
 * ZipVoice AXERA C++ Port
 *
 * EngineWrapper: wraps AX_ENGINE API for loading and running .axmodel files.
 * Supports both index-based and name-based I/O access.
 *
 * Adapted from melotts.axera-main
 **************************************************************************************************/

#pragma once

#include <string>
#include <vector>
#include <unordered_map>
#include <cstring>
#include <cstdint>

#include "ax_engine_api.h"

class EngineWrapper {
public:
    EngineWrapper();
    ~EngineWrapper();

    int Init(const char* strModelPath, uint32_t nNpuType = 0);

    // Index-based I/O (simple interface)
    int SetInput(void* pInput, int index);
    int RunSync();
    int GetOutput(void* pOutput, int index);

    // Name-based I/O (for multi-input/output models)
    int SetInputByName(const char* name, void* pInput);
    int GetOutputByName(const char* name, void* pOutput);

    int GetInputSize(int index);
    int GetOutputSize(int index);
    int GetInputSizeByName(const char* name);
    int GetOutputSizeByName(const char* name);

    int GetInputIndex(const char* name);
    int GetOutputIndex(const char* name);
    int GetInputCount() const { return m_input_num; }
    int GetOutputCount() const { return m_output_num; }
    const char* GetInputName(int index) const;
    const char* GetOutputName(int index) const;
    const char* GetInputDtypeStr(int index) const;
    const char* GetOutputDtypeStr(int index) const;

    bool HasInit() const { return m_hasInit; }
    int Release();

private:
    bool m_hasInit;
    AX_ENGINE_HANDLE m_handle;
    AX_ENGINE_IO_INFO_T *m_io_info;
    AX_ENGINE_IO_T m_io;
    int m_input_num, m_output_num;

    std::unordered_map<std::string, int> m_input_name_to_idx;
    std::unordered_map<std::string, int> m_output_name_to_idx;
};