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

#pragma once

#include <memory>
#include <string>
#include <optional>
#include <vector>
#include <functional>
#include <atomic>
#include <type_traits>

#include "calrt_utils.h"
#include "calrt_calbin.h"

namespace calrt 
{
    enum CalrtDeviceType_e{
        NO_TYPE = 0,
        PCIE = 1,
        USB,
        EMU
    };

    enum class DeviceWorkingStatus : int
    {
        // 0: busy, 1: AVAILABLE, 2: taken
        BUSY = 0,
        AVAILABLE,
        TAKEN
    };

    struct CalrtDeviceId_s{
        // bool emuId = false; //true if users only want emulator which would ignore physical device. Users need to decided if they need to turn on this option.
        CalrtDeviceType_e devType;
        std::string devPath;
        int serialNum = 0; // 0 = ks1, 1 = ks2 ...

        int32_t compare(const CalrtDeviceId_s &rh) {
            int32_t result = 0;
            result += (devType == rh.devType) ? 0 : 1;
            result += (devPath == rh.devPath) ? 0 : 1;
            return result;
        }
    };

    struct CalrtDeviceInfo_s {
        CalrtDeviceId_s devId;
        std::string name;
        uint64_t ddrSize = (32UL << 30);
        float sramFreqMhz = 1000.0;
        uint64_t sramSizeMB = 18;
    };

    struct RemoteProtocol {
        uint32_t masks[3];
        char password[128];
    };

    class CALRT_API CalrtDeviceInterface {

    public:
        inline static std::function<bool(char *, const uint64_t, const uint64_t, const int32_t)> DeviceMemAccessFunc;
        static bool DeviceMemAccess(char * calrtData, const uint64_t devAddr, const uint64_t size, const int32_t write_en);
    };

    class CALRT_API  CalrtDevice
    {

    friend class VirtualDevice;

    public:
        /**
         * @brief Create a Device object. Current support [PCIe], [USB]
         *
         * @note always return valid pointer, but self-check if device is valid. Priority PCIe device.
         * @return std::unique_ptr<CalrtDevice>
         */
        static std::unique_ptr<CalrtDevice> CreateDevice();

        /**
         * @brief Create a PCIe Device
         *
         * @note always return valid pointer, but self-check if device is valid.
         * @param id
         * @return std::unique_ptr<CalrtDevice>
         */
        static std::unique_ptr<CalrtDevice> CreatePCIeDevice();

        /**
         * @brief Create a Emulator Device
         * 
         * @return std::unique_ptr<CalrtDevice> 
         */
        static std::unique_ptr<CalrtDevice> CreateEmuDevice();

        /**
         * @brief used for refresh device manager for any new hotplug devices installation
         * 
         */
        static void RefreshDevice();

        virtual ~CalrtDevice() = default;
        CalrtDevice(const CalrtDevice &) = delete;
        CalrtDevice &operator=(const CalrtDevice &) = delete;
        CalrtDevice(CalrtDevice &&) = delete;
        CalrtDevice &operator=(CalrtDevice &&other) = delete;

        typedef bool (*DeviceMemAccessPtr)(char *, const uint64_t, const uint64_t, const int32_t);

        virtual void SetDevMemAccessApi(DeviceMemAccessPtr fDeviceMemAccess);

        virtual CalrtError_e RegisterDram(int64_t startAddr, int64_t size, const char *tag="rsvd_dram") = 0;

        virtual CalrtError_e RegisterSyncUnit(int64_t startAddr, int64_t size, const char *tag="rsvd_dram") = 0;

        virtual CalrtError_e RegisterSramBuf(int64_t startAddr, int64_t size, const char *tag="rsvd_dram") = 0;

        CalrtError_e Configure(CalbinModel *pModel, Calbin *pCalbin);

        CalrtError_e Configure(CalbinModel &model, Calbin *pCalbin);

        /**
         * @brief dump all configure meta data at once
         * 
         * @param modelName 
         * @param pCalbin 
         * @param tag 
         */
        void DumpDeviceConfigureMetaData(const std::string &modelName, Calbin *pCalbin, const std::string &tag);

        /**
         * @brief allow to dump multiple section at once in "ini" file
         * 
         * @tparam CalbinSectionType_e
         */
        template <typename... SecType>
        inline void DumpDeviceConfigToIniFileBySections(const std::string &outputPath, SecType... secType)
        {
            static_assert((std::is_same_v<SecType, CalbinSectionType_e> && ...), "Must pass CalbinSectionType_e arguments");
            std::unordered_map<std::string, std::string> iniFile;

            (RecordDeviceConfig(iniFile, secType), ...);

            DumpMapMirror(outputPath, iniFile, 2);
        }

        /**
         * @brief copy host memory to device memory
         *
         * @param srcAddr host mem address
         * @param devAddr device mem address
         * @param size bytes
         */
        virtual CalrtError_e WriteToDevice(void *srcAddr, const uint64_t devAddr, const uint64_t size) = 0;

        /**
         * @brief copy device memory to host memory
         *
         * @param srcAddr host mem address
         * @param devAddr device mem address
         * @param size bytes
         */
        virtual CalrtError_e ReadFromDevice(void *srcAddr, const uint64_t devAddr, const uint64_t size) = 0;

        virtual CalrtError_e WriteReg(const uint64_t regAddr, uint32_t data) = 0; // only support bar2

        virtual CalrtError_e ReadReg(const uint64_t regAddr, uint32_t &data) = 0; // only support bar2

        // only support bar4
        virtual CalrtError_e WriteReg64(const uint64_t regAddr, uint64_t data)
        {
            return CalrtSuccess;
        }

        virtual CalrtError_e ReadReg64(const uint64_t regAddr, uint64_t &data)
        {
            return CalrtSuccess;
        }

        virtual uint64_t CreateBuf(uint64_t size, uint32_t alignLog2Byte=0) = 0;

        virtual uint64_t CreateSramBuf(uint64_t size, uint32_t alignLog2Byte=0) = 0;

        virtual uint64_t ApplySyncUnit() = 0;

        virtual void FreeBuf(uint64_t addr) = 0;

        virtual void FreeSramBuf(uint64_t addr) = 0;

        virtual void FreeSyncUnitByAddress(uint64_t addr) = 0;

        virtual void FreeSyncUnitByIndex(uint32_t idx) = 0;

        virtual void ClearAllDevMem() = 0;

        virtual void ClearDynamicMem() = 0;

        //virtual void WaitJobDone(std::shared_ptr<CalrtJob> pJob);

        virtual CalrtError_e Reset() = 0;

        void ResetCCU();

        //TODO it should return a list or struct. TBD
        virtual CalrtError_e GetFirmwareInfo();

        virtual std::string GetDevInfo() = 0;

        virtual std::unordered_map<int64_t, int64_t> GetReservedMem() = 0;

        virtual void SetPowerMode(PowerMode power) = 0;

        /**
         * @brief Get the Memory Usage object
         * 
         * @return std::pair<float,float> {usedDram, DramTotal} size in MiB
         */
        virtual std::pair<float,float> GetMemoryUsage() = 0;

        void SetConfigModel(Calbin* pCalbin);

        const uint64_t GetCurrentCalbinOnDevice();

        int32_t GetNumPendingJob(); // CLI impl func

        int32_t GetNumFinishedJob(); // CLI impl func

        /**
         * @brief indicating if device is created successfully
         * 
         * @return const CalrtError_e 
         */
        virtual const CalrtError_e Status() const = 0;

        /**
         * @brief Set the Status object. indicating if device is created successfully
         * 
         */
        virtual void SetStatus(CalrtError_e status) = 0;

        /**
         * @brief Get the Device Running Status.
         * 
         * @return int 
         */
        virtual DeviceWorkingStatus GetDeviceRunningStatus() = 0;

        const std::string &Name() {return mInfo.name;}

        void EnableTraceDevice(bool enable);

        virtual uint64_t GetDeviceLeftJobs() = 0;

        int GetDeviceSerialNum() { return mInfo.devId.serialNum; }

        virtual void Release() = 0;

        uint32_t ReadRemoteReg(uint64_t regAddr, uint32_t chipId);

        CalrtError_e ReadRemoteChipMem(void *srcAddr, const uint64_t devAddr, const uint64_t size, uint32_t chipId);

        CalrtError_e WriteRemoteChipMem(void *srcAddr, const uint64_t devAddr, const uint64_t size, uint32_t chipId);

        void WriteRemoteReg(uint64_t regAddr, uint32_t regValue, uint32_t chipId);

        CalrtDeviceType_e Type() noexcept {return mInfo.devId.devType;}

        void ClearDevLib();

        virtual uint32_t GetExpectedEndQueueEntries() {return 0;}

#ifndef RELRT
        std::string GetCalbinPath() noexcept;
#endif

    protected:
        CalrtDevice(const CalrtDeviceInfo_s &info);

        void StartJob(uint64_t jobAddr, uint64_t jobSize);

        /**
         * @brief wait job done
         * 
         * @return uint32_t ccu work status
         */
        uint32_t WaitJobDone(bool singleEntry = false);

        /**
         * @brief wait job done
         * 
         * @param jobEndInfo record [Hi32: infer time consume, Lo32: finished job id]
         * @return uint32_t ccu work status
         */
        uint32_t WaitJobDone(uint64_t& jobEndInfo, bool singleEntry = false);
        virtual void PollingJobDonePowerUp() = 0;
        /**
         * @brief 
         * 
         * @param status 0: busy, 1: ready, 2: taken 3: error
         */
        virtual void UpdateDeviceRunningStatus(uint32_t status) {};
        virtual void UpdateDeviceJobs(bool add) {};
        void CleanJobQ();
        CalbinSectionPlace_e GetStackPlace(const std::string &modelName);

        CalrtError_e DeployGlobalMem(Calbin *pCalbin);
        CalrtError_e DeployCPUSection (std::vector<CalbinSection_s> &sectionList, std::vector<int32_t> &tempChips, Calbin *pCalbin);
        CalrtError_e DeployCCUSection (std::vector<CalbinSection_s> &sectionList, std::vector<int32_t> &tempChips, CalbinModel &model, Calbin *pCalbin);
        CalrtError_e DeployRodataSection (std::vector<CalbinSection_s> &sectionList, Calbin *pCalbin);
        CalrtError_e DeployLibSection (std::vector<CalbinSection_s> &sectionList, std::vector<int32_t> &tempChips, CalbinModel &model, Calbin *pCalbin);

        virtual std::vector<std::unordered_map<int64_t, int64_t>> GetDeviceMemoryLayout() = 0;
        virtual void UpdateBandWidth(float bandwidth) = 0;

        void RecordDeviceConfig(std::unordered_map<std::string, std::string> &iniFile, const CalbinSectionType_e sectionType);

        void DumpMapMirror(const std::string &path, const std::unordered_map<std::string, std::string> &iniFile, int flag);

        /**
         * @brief 
         * 
         * @param num 
         * @param flag 0: minus, 1: add 2: set 3: reset
         */
        virtual void UpdateDevRuntimeInfo (uint32_t num, int flag) = 0;

    private:
        CalrtDeviceInfo_s mInfo;
        Calbin *pmCurrentCalbin = nullptr;
        uint64_t mCurCalbinHash = 0;
        std::unordered_map<std::string, CalbinSectionPlace_e> mStackPlace;
        std::unordered_map<std::string, uint64_t> mWorkloadsPerModel;
        std::atomic<int32_t> mDoneJobs = 0;

        /**
         * @brief deploy lib for chip 0 and chipx if chip mask contains both chip0 and chipx
         * 
         * @param section 
         * @param modelName 
         */
        void DeployLibToX(CalbinSection_s &section, const std::string &modelName);

        /**
         * @brief deploy lib for chip0 only
         * 
         * @param section 
         * @param modelName 
         */
        void DeployLibToHead(CalbinSection_s &section, const std::string &modelName);

        // class DeviceInternal;
        // std::unique_ptr<DeviceInternal> pmInternal;
    };
}