/**
 * @file calrt_utils.h
 * @author your name (you@domain.com)
 * @brief 
 * @version 0.1
 * @date 2024-11-11
 * 
 * @copyright Copyright (c) 2024
 * 
 */

#pragma once

#include "calrt_platform.h"

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

#include "calrt_error.h"

namespace calrt {

enum CALRT_API PrimitiveType 
{
  U1=0,     //1bit
  PRED=1,   //bool, int8
  U8=2,
  S8=3,
  FP8_143=4, // exp_bias = 8; has subnormal, no inf/nan. when biased_exp=0, it is fixed explained as unbiased_exp=-7;
  FP8_152=5, // exp_bias = 16; has subnormal, no inf/nan.  when biased_exp=0, it is fixed explained as unbiased_exp=-15;
  U16=6,
  S16=7,    //16b
  BF16=8,   //16b
  U32=9,    //32b
  S32=10,
  F32=11,
  U4=12,
  S4=13,
  F35=14,
  U64=15,   // not support
  S64=16,   // u64和s64需要保留，仅做类型转换，因为pytorch要求tensor做index时必须为long或byte或bool。

  TF32=17,//19b    // not support
  F16=18,
  F64=19,
  C64,
  C128,

  TUPLE,//0
  TOKEN,//32b
  OPAQUE_TYPE,//32b
  INVALID, //0b
  TYPE_SIZE
};

PrimitiveType PrimitiveTypeFromInt(const int32_t& type);

std::string PrimitiveTypeToString(const PrimitiveType& type);
std::string PrimitiveTypeToString(const int32_t& type);

CALRT_API int32_t PrimitiveTypeBitSize(const PrimitiveType& type);

struct CALRT_API CalbinTensorInfo_s {
    std::string             name;
    std::vector<uint64_t>    shape;
    PrimitiveType           dataType;
    int64_t                 devAddr0 = -1;   // ping
    int64_t                 devAddr1 = -1;   // pong
    int64_t                 size = -1;
};

enum CalbinSectionType_e {
    CALBIN_SECTION_NULL     = 0,
    CALBIN_SECTION_CPU_CMD  ,
    CALBIN_SECTION_CPU_SO   ,
    CALBIN_SECTION_CCU_ELF  ,
    CALBIN_SECTION_RODATA   ,
    CALBIN_SECTION_IBUF     ,
    CALBIN_SECTION_OBUF     ,
    CALBIN_SECTION_MEMRSVD  ,
    CALBIN_SECTION_KV       ,
    CALBIN_SECTION_CSR      ,
    CALBIN_SECTION_NUM
};

constexpr const char * CalbinSecTypeToString(CalbinSectionType_e e)
{
    switch (e)
    {
        case CALBIN_SECTION_CPU_CMD: return "CPU command";
        case CALBIN_SECTION_CCU_ELF: return "CCU ELF";
        case CALBIN_SECTION_CPU_SO: return ".so file";
        case CALBIN_SECTION_RODATA: return "parameter";
        case CALBIN_SECTION_IBUF: return "input buffer";
        case CALBIN_SECTION_OBUF: return "output buffer";
        case CALBIN_SECTION_MEMRSVD: return "reserved memory";
        case CALBIN_SECTION_KV: return "kv cache";
        case CALBIN_SECTION_CSR: return "csr";
        default: return "null";
    }
}

enum CalbinSectionPlace_e{
    DRAM = 0,
    SRAM,
    UNKNOWN
};

/**
 * @brief 
 * @note CalbinSectionType_e == CPU_SO. only need offset(used for prgIdx) devAddr, size,
 */
struct CALRT_API CalbinSection_s {
    CalbinSectionPlace_e            place       = UNKNOWN;
    CalbinSectionType_e             type        = CALBIN_SECTION_NULL; // cpu txt, ccu elf, data
    std::string                     name        = "";
    std::string                     filePath    = ""; // must include elf/cpu file path
    int64_t                         offset      = -1; // [progbit & .bss] entry address
    int64_t                         devAddr     = -1;    // -1 means calrt set memory, for memory and program
    int64_t                         size        = -1;
    int64_t                         progEntry   = -1; // program header entry address

    std::vector<uint32_t>           chipMask;     // mask2, mask1, mask0, small-endian
    std::vector<CalbinTensorInfo_s>   tensorInfos;        // for i/o buf

    bool operator<(const CalbinSection_s& other) const;
    bool operator==(const CalbinSection_s& other) const;
    bool operator!=(const CalbinSection_s& other) const;
    friend std::ostream& operator<< (std::ostream& os, const CalbinSection_s&);
};

struct CALRT_API KV_Cache_s
{
    PrimitiveType dtype = INVALID;
    PrimitiveType sin_cos_table_data_type = INVALID;
    uint32_t layer;
    uint64_t kv_base[2]; //[0] k base address, [1] v base address
    uint64_t size;
    uint32_t n_head;
    uint32_t n_dim;

    operator bool () const noexcept
    {
        return dtype != INVALID;
    }
};

struct CALRT_API CalbinLLM_s
{
    uint32_t max_batch_size = 0; // 16
    uint32_t max_seq_len = 0;

    KV_Cache_s kv_cache;
    operator bool () const noexcept
    {
        return max_batch_size != 0 && max_seq_len != 0;
    }
};

enum ModelChipMode_e
{
    SCHIP=0,
    MCHIP=1,
    UNDEFINED_CHIP_MODE
};

struct CALRT_API CalbinModel {
    enum CalDevAccType_e {
        CAL_DEV_ACC_NULL    = 0,
        CAL_DEV_ACC_CPU     = 1,
        CAL_DEV_ACC_CCU     = 2,
        CAL_DEV_ACC_CPU_CCU = 3
    };

    bool                    isConfigured = false;
    CalDevAccType_e         devAccType;
    ModelChipMode_e         chipMode = UNDEFINED_CHIP_MODE;
    std::string             modelType;
    std::string             modelArch;
    std::string             modelName;

    std::unordered_map<CalbinSectionType_e, std::vector<CalbinSection_s>> sections; //idx follow above
    std::vector<int32_t>    tarChipN;
};

struct CALRT_API CalrtCalbin {
    uint64_t name; // md5 value or hash value
    std::vector<CalbinModel>  models;
};

enum CalrtBufferDirection_e{
    HostToDevice = 0,
    DeviceToHost = 1,
};

enum class PowerMode : uint8_t
{
    BALANCE = 1,
    HIGH_PERFORMANCE = 2
};

class CALRT_API CalrtTensor {

public:
    CalrtTensor() = delete;

    CalrtTensor(const std::vector<uint64_t> &shape, PrimitiveType elemType, CalrtBufferDirection_e direction, std::string name = "");

    template<typename T>
    CalrtError_e CheckTensorByDataSize(const std::vector<T> &datas_or_shape, PrimitiveType dtype)
    {

        bool success = mElemBitSize == PrimitiveTypeBitSize(dtype);
        if(!success) return CalrtErrorInvalidDataType;
        success = success && ((int64_t)datas_or_shape.size() == mElemSize);

        if(!success) return CalrtErrorInvalidDataShape;

        return CalrtSuccess;
    }

    CalrtError_e CheckTensor(const std::vector<uint64_t> &shape, PrimitiveType dtype);

    /**
     * @brief safe to copy date to tensor in byte
     * @note for special custom data type, Using GetDataPtr() to manuelly fill data into tensor.
     * 
     * @tparam T 
     * @param datas input data
     * @param dtype data type
     */
    template<typename T>
    void Fill(std::vector<T> &datas, PrimitiveType dtype)
    {
        FillImpl(static_cast<void*>(datas.data()), CheckTensorByDataSize(datas, dtype));
    }

    /**
     * @brief map host buffer address to device address
     * @note unsafety for non-tensor check
     * 
     * @param src 
     */
    void MapBuf(void *src);

    void UnMapBuf();

    /**
     * @brief Get the Data Ptr object. User must guarantee a safety copy action
     * 
     * @return char* 
     */
    uint8_t* GetDataPtr();

    PrimitiveType Type() {return mElemType;}
    const std::vector<uint64_t> &Shape() {return mShape;}

    /**
     * @brief Get the current tensor byte size
     * 
     * @return const int64_t 
     */
    uint64_t ByteSize() {return mByteSize;}
    
    /**
     * @brief Get the number of element from the current tensor based on dtype
     * 
     * @return const int64_t 
     */
    uint64_t Size() {return mElemSize;}

    const std::string &Name() {return mName;}

    CalrtBufferDirection_e Direction();

    CalrtError_e SliceTensor(uint64_t offset, uint64_t size);
    void UndoSlice();
    uint64_t Offset();
    uint64_t TransSize();

    friend std::ostream& operator<<(std::ostream& os, const CalrtTensor& obj);

    private:
        void FillImpl(void* src, CalrtError_e err);
        
        std::vector<uint64_t> mShape;
        PrimitiveType mElemType;    // tensor data elem type
        CalrtBufferDirection_e mDirection;
        std::string mName;
        uint32_t mElemBitSize;   // each elem bit size
        uint64_t mElemSize;      // elem size
        uint64_t mByteSize;      // momory byte size
        uint64_t mOffset;
        uint64_t mTransDataSize;
        std::vector<uint8_t> mData;
        uint8_t *mRawPtr = nullptr;
};

struct CalrtDevBuf_s {
    int64_t addr0 = -1;
    int64_t addr1 = -1;
    int64_t size = -1;

    friend std::ostream& operator<<(std::ostream& os, const CalrtDevBuf_s& obj);
};

struct CalrtBufferInfo_s{
    CalrtBufferDirection_e direction;
    std::string name;
    std::vector<CalbinTensorInfo_s> tensorInfos;
    std::vector< std::pair<std::string, uint32_t>> csrTable;    
};

enum class TaskType_e : uint32_t
{
    TASK_PING = 0,
    TASK_PONG,
    UNDEFINED_TASK
};

enum class ChipArch_e : uint32_t
{
    SINGLE_CHIP_TASK = 0,
    MULTIPLE_CHIP_TASK,
    UNDEFINED
};

struct TaskInfo_s
{
    TaskType_e m_taskType;
    uint32_t m_jobId;
    uint32_t m_enableCcu;
    ChipArch_e m_taskChipArch;
    uint32_t m_syncRegIdx;
    uint32_t m_syncExpectedVal;
};

// struct CalrtVersion
// {
//     uint32_t major; // API or ABI change. E.G. remove/change function, user may need to adapt their code
//     uint32_t minor; // added new feature
//     uint32_t patch; // fix bug
// };

CALRT_API const char *calrt_version();

CALRT_API void SetFullDebug(bool enable);

CALRT_API bool isFullDebug();

}