#pragma once

#include <cstddef>
#include <cstdint>
#include <string>
#include <memory>
#include <unordered_map>
#include <atomic>
#include <condition_variable>
#include <mutex>

#include "calrt_utils.h"

namespace calrt
{
    class CalrtJob;

    struct TensorPkg
    {
        CalrtTensor *tensor = nullptr;
        uint64_t address[2]{0};

        // TensorPkg() = default;
        // TensorPkg(const TensorPkg&) = delete;
        // TensorPkg& operator=(const TensorPkg&) = delete;
        // TensorPkg(TensorPkg&&) noexcept = default;
        // TensorPkg& operator=(TensorPkg&&) noexcept = default;
        ~TensorPkg();
    };

    struct CALRT_API ModelHyperParameters_s
    {
        /**
         * @brief Get the Csr Value By Name object
         * 
         * @param csrName 
         * @return uint32_t 0: for error msg
         */
        uint32_t GetCsrValueByName(const std::string &csrName);
        CalrtError_e SetCsrByName(const std::string &csrName, uint32_t value);

        private:
            friend class CalrtInputBuf;
            void init(const std::string &csrName, uint32_t value);
            std::unordered_map<std::string, uint32_t> csr; // csr | value
    };

    class CALRT_API CalrtInputBuf
    {
    public:
        static std::unique_ptr<CalrtInputBuf> CreateInputBuf(CalrtBufferInfo_s &&info);
        CalrtInputBuf(CalrtBufferInfo_s &&info);
        CalrtInputBuf(const CalrtInputBuf &) = delete;
        CalrtInputBuf& operator=(const CalrtInputBuf&) = delete;
        CalrtInputBuf(CalrtInputBuf &&) noexcept = default;
        CalrtInputBuf &operator=(CalrtInputBuf &&other) noexcept = default;
        
        ~CalrtInputBuf();

        std::unordered_map<std::string, CalrtTensor*> GetTensors();
        CalrtTensor* GetTensorByName(const std::string &tensorName);

        /**
         * @brief slice tensor by offset and size for the specific data transfer
         * 
         * @param src host side buffer
         * @param tensorName 
         * @param offset 
         * @param size 
         * @return CalrtSuccess on success  
         */
        CalrtError_e SliceTensorByName(void *src, const std::string &tensorName, uint64_t offset, uint64_t size);
        CalrtError_e ResetTensorByName(const std::string &tensorName);
        void ResetAllTensors();
        std::vector<CalrtDevBuf_s>& GetTensorMemInfo() { return mDevBufs; }

        /**
         * @brief Get the current modifiable HyperParam struct
         * 
         * @return ModelHyperParameters_s& 
         */
        ModelHyperParameters_s &GetCurHyperParam() {return mHyperParam;}
        
        std::string Name();

        size_t GetTensorNum();

        float GetInputTransferTime();

    private:
        friend class CalrtJob;

        CalrtBufferInfo_s mInfo;
        std::vector<CalrtDevBuf_s> mDevBufs;
        ModelHyperParameters_s mHyperParam;
        std::vector<TensorPkg> mTensorPkgs;
        std::unordered_map<std::string, uint32_t> mTensorNameMap;
        float mTransTime;

        std::vector<TensorPkg> &GetTensorPkg() {return mTensorPkgs;}
        const std::vector< std::pair<std::string, uint32_t> > &GetCsrTable() const {return mInfo.csrTable;}
        void SetTransTime(float transTime) noexcept;
    };

    class CALRT_API CalrtOutputBuf
    {
    public:
        static std::unique_ptr<CalrtOutputBuf> CreateOutputBuf(CalrtBufferInfo_s &&info);
        CalrtOutputBuf(CalrtBufferInfo_s &&info);
        CalrtOutputBuf(const CalrtOutputBuf &) = delete;
        CalrtOutputBuf& operator=(const CalrtOutputBuf&) = delete;
        CalrtOutputBuf(CalrtOutputBuf &&) = delete;
        CalrtOutputBuf &operator=(CalrtOutputBuf &&other) = delete;

        enum class CalrtOutputBufStatus_e : int32_t {
            CALRT_OBUF_STATUS_DEFAULT = 0,
            CALRT_OBUF_STATUS_PENDING,  // in sw start queue, not started yet
            CALRT_OBUF_STATUS_RUNNING,  // in calcore, running
            CALRT_OBUF_STATUS_DONE,      // in sw done queue, done
            CALRT_OBUF_STATUS_DONE_CCU_EXCEPTION
        };

        ~CalrtOutputBuf();

        void SetStatus(CalrtOutputBufStatus_e status);
        CalrtOutputBufStatus_e GetStatus();

        std::unordered_map<std::string, CalrtTensor*> GetTensors();
        CalrtTensor* GetTensorByName(const std::string &tensorName);

        /**
         * @brief slice tensor by offset and size for the specific data transfer
         * 
         * @param src host side buffer
         * @param tensorName 
         * @param offset 
         * @param size 
         * @return CalrtSuccess on success 
         */
        CalrtError_e SliceTensorByName(void *src, const std::string &tensorName, uint64_t offset, uint64_t size);
        CalrtError_e ResetTensorByName(const std::string &tensorName);
        void ResetAllTensors();

        void Notify(CalrtOutputBuf::CalrtOutputBufStatus_e status);
        CalrtError_e Wait();

        CalrtBufferInfo_s GetBufferInfo();
        std::string Name();
        void Reset();
        size_t GetTensorNum();
        float GetWaitTime();
        float GetOutputTransferTime();

    private:
        friend class CalrtJob;

        CalrtBufferInfo_s mInfo;
        std::atomic<int32_t> mStatus {0};
        std::vector<TensorPkg> mTensorPkgs;
        std::unordered_map<std::string, uint32_t> mTensorNameMap;
        float mWaitTime;
        float mTransTime;
        std::condition_variable mDone;
        std::mutex mLock;

        std::vector<TensorPkg> &GetTensorPkg();
        void SetWaitTime(float time) noexcept;
        void SetTransTime(float transTime) noexcept;
    };
} // namespace calrt