#pragma once

#include "NVMeIntf.hpp"
#include "Utils.hpp"

#include <boost/asio.hpp>
#include <sdbusplus/bus.hpp>

#include <cstdint>
#include <thread>
#include <utility>

// A worker thread for calling NVMeMI cmd.
class NVMeMiWorker
{
  private:
    bool workerStop = false;
    std::mutex workerMtx;
    std::condition_variable workerCv;
    boost::asio::io_context workerIO;
    bool workerIsNotified = false;
    std::thread thread;

    // private to enforce the use of static create
    NVMeMiWorker();

  public:
    NVMeMiWorker(const NVMeMiWorker&) = delete;
    ~NVMeMiWorker();
    void post(std::function<void(void)>&& func);

    /**
       will create a shared_ptr of NVMeMiWorker asserting we are in the same
       thread that is running the passed io context and the shared_ptr will have
       custom deleter which will ensure the NVMeMiWorker will be destroyed in
       the passing io. The main purpose of this is to ensure the worker is being
       destroyed in same thread it was created from following single thread
       arch.

       @param io the io context of the caller thread, assertion will fail if the
       io context is not from the caller thread
    */
    template <typename... Args>
    static std::shared_ptr<NVMeMiWorker> create(boost::asio::io_context& io,
                                                Args&&... args)
    {
        assert(io.get_executor().running_in_this_thread() &&
               "NVMeMiWorker::create called from wrong io_context thread and "
               "should be called from the same thread that is running the io "
               "context.");
        // Note: NVMeMiWorker constructor is private
        auto* worker = new NVMeMiWorker(std::forward<Args>(args)...);
        return {worker, [&io](NVMeMiWorker* objPtr) {
            boost::asio::post(io, [objPtr]() { delete objPtr; });
        }};
    }
};

class NVMeMi : public NVMeMiIntf, public std::enable_shared_from_this<NVMeMi>
{
  public:
    NVMeMi(boost::asio::io_context& io,
           const std::shared_ptr<sdbusplus::asio::connection>& conn,
           const std::shared_ptr<const MctpDevice>& device,
           const std::shared_ptr<NVMeMiWorker>& worker,
           PowerState readState = PowerState::always);
    ~NVMeMi() override;

    bool flushOperations(std::function<void()>&& cb) override;
    void miSubsystemHealthStatusPoll(
        std::function<void(const std::error_code&,
                           nvme_mi_nvm_ss_health_status*)>&& cb) override;
    void miScanCtrl(std::function<void(const std::error_code&,
                                       const std::vector<nvme_mi_ctrl_t>&)>
                        cb) override;
    void adminIdentify(
        nvme_mi_ctrl_t ctrl, nvme_identify_cns cns, uint32_t nsid,
        uint16_t cntid,
        std::function<void(nvme_ex_ptr, std::span<uint8_t>)>&& cb) override;
    void
        adminGetFeatures(nvme_mi_ctrl_t ctrl, const GetFeaturesRequest& req,
                         std::function<void(nvme_ex_ptr, uint32_t response,
                                            std::span<uint8_t>)>&& cb) override;
    void adminSetFeatures(
        nvme_mi_ctrl_t ctrl, const SetFeaturesRequest& req,
        std::function<void(nvme_ex_ptr, uint32_t response)>&& cb) override;
    void adminGetLogPage(nvme_mi_ctrl_t ctrl, nvme_cmd_get_log_lid lid,
                         uint32_t nsid, uint8_t lsp, uint16_t lsi,
                         std::function<void(const std::error_code&,
                                            std::span<uint8_t>)>&& cb) override;

    void adminFwCommit(
        nvme_mi_ctrl_t ctrl, nvme_fw_commit_ca action, uint8_t slot, bool bpid,
        std::function<void(const std::error_code&, nvme_status_field)>&& cb)
        override;

    void adminFwDownload(nvme_mi_ctrl_t ctrl, std::string firmwarefile,
                         std::function<void(const std::error_code&,
                                            nvme_status_field)>&& cb) override;

    void adminXfer(nvme_mi_ctrl_t ctrl, const nvme_mi_admin_req_hdr& adminReq,
                   std::span<uint8_t> data, unsigned int timeoutMs,
                   std::function<void(const std::error_code&,
                                      const nvme_mi_admin_resp_hdr&,
                                      std::span<uint8_t>)>&& cb) override;

    void adminLockdownInband(
        nvme_mi_ctrl_t ctrl, uint8_t prohibit,
        const std::vector<uint8_t>& adminCmds,
        const std::vector<uint8_t>& features,
        const std::vector<uint8_t>& logPages,
        std::function<
            void(const std::error_code&, int nvmeStatus, uint32_t comptionDw0,
                 const std::string& failureScope, uint32_t failureId)>&& cb)
        override;

    void adminSecuritySend(nvme_mi_ctrl_t ctrl, uint8_t proto,
                           uint16_t protoSpecific, std::span<uint8_t> data,
                           std::function<void(const std::error_code&,
                                              int nvmeStatus)>&& cb) override;

    void adminSecurityReceive(
        nvme_mi_ctrl_t ctrl, uint8_t proto, uint16_t protoSpecific,
        uint32_t transferLength,
        std::function<void(const std::error_code&, int nvmeStatus,
                           std::span<uint8_t> data)>&& cb) override;

    void adminNonDataCmd(
        nvme_mi_ctrl_t ctrl, uint8_t opcode, uint32_t cdw1, uint32_t cdw2,
        uint32_t cdw3, uint32_t cdw10, uint32_t cdw11, uint32_t cdw12,
        uint32_t cdw13, uint32_t cdw14, uint32_t cdw15,
        std::function<void(const std::error_code&, int nvmeStatus,
                           uint32_t comptionDw0)>&& cb) override;

    void createNamespace(
        nvme_mi_ctrl_t ctrl, uint64_t size, size_t lbaFormat,
        bool metadataAtEnd, std::function<void(nvme_ex_ptr ex)>&& submittedCb,
        std::function<void(nvme_ex_ptr ex, NVMeNSIdentify newid)>&& finishedCb)
        override;

    void adminDeleteNamespace(
        nvme_mi_ctrl_t ctrl, uint32_t nsid,
        std::function<void(const std::error_code&, int nvmeStatus)>&& cb)
        override;

    void adminListNamespaces(
        nvme_mi_ctrl_t ctrl,
        std::function<void(nvme_ex_ptr ex, std::vector<uint32_t> ns)>&& cb)
        override;

    void adminAttachDetachNamespace(
        nvme_mi_ctrl_t ctrl, uint16_t ctrlid, uint32_t nsid, bool attach,
        std::function<void(const std::error_code&, int nvmeStatus)>&& cb)
        override;

    void adminSanitize(nvme_mi_ctrl_t ctrl, enum nvme_sanitize_sanact sanact,
                       uint8_t passes, uint32_t pattern, bool invertPattern,
                       std::function<void(nvme_ex_ptr ex)>&& cb) override;

    void start(const std::shared_ptr<MctpEndpoint>& ep) override;
    void stop() override;
    void recover() override;

  private:
    // the transfer size for nvme mi messages.
    // define in github.com/linux-nvme/libnvme/blob/master/src/nvme/mi.c
    static constexpr size_t nvmeMiXferSize = 4096;

    static nvme_root_t nvmeRoot;

    boost::asio::io_context& io;
    std::shared_ptr<const MctpDevice> device;

    // power state
    std::unique_ptr<PowerCallbackEntry> powerCallback;
    PowerState readState;

    /*
     * A state machine to represent the current status of the MCTP connection.
     * In Reset state, the MCTP endpoint (EP) is not setup with the device.
     * In event of successful setup we move from Reset to Configured. If
     * opening the EP is successful from Configured the status will change to
     * Initiated. The status will change to Connected once the MTU of local
     * and device side MTU and frequency is optimized. In an event of connection
     * EP closure, the status will move back to Reset via Terminating.
     *
     * Transitions to the terminal state indicate a logic error.
     *
     * stateDiagram
     *   [*] --> Reset
     *
     *   Reset --> Reset: epReset()
     *   Reset --> [*]: epConnect()
     *   Reset --> [*]: epOptimize()
     *
     *   Initiated --> Terminating: epReset()
     *   Initiated --> Initiated: epConnect()
     *   Initiated --> Connected: epOptimize()
     *
     *   Connected --> Terminating: epReset()
     *   Connected --> Connected: epConnect()
     *   Connected --> Connected: epOptimize()
     *
     *   Terminating --> Reset: Reset close job executes
     *   Terminating --> Terminating: epReset()
     *   Terminating --> Terminating: epConnect()
     *   Terminating --> [*]: epOptimize()
     */
    enum class Status : uint8_t
    {
        Reset,
        Initiated,
        Connected,
        Terminating,
    };

    void epReset();
    bool epConnect(int lnid, uint8_t leid);
    void epOptimize();

    Status mctpStatus;
    std::shared_ptr<MctpEndpoint> endpoint;
    uint16_t mtu;
    nvme_mi_ep_t nvmeEP;
    // Handle a start() while in Status::Terminating on entry to Status::Reset.
    bool restart;
    std::shared_ptr<boost::asio::steady_timer> optimizeTimer;

    std::shared_ptr<NVMeMiWorker> worker;
    void post(std::function<void(void)>&& func);

    void
        miConfigureRemoteMCTP(uint8_t port, uint16_t mtu,
                              uint8_t maxSupportedFreq,
                              std::function<void(const std::error_code&)>&& cb);

    void miConfigureSMBusFrequency(
        uint8_t portId, uint8_t maxSupportedFreq,
        std::function<void(const std::error_code&)>&& cb);

    void miSetMCTPConfiguration(
        std::function<void(const std::error_code&)>&& cb);

    void configureLocalRouteMtu(
        std::function<void(const std::error_code& ec)>&& completed,
        int retries = 5);

    std::optional<std::error_code> isEndpointDegraded() const;

    bool readingStateGood() const
    {
        return ::readingStateGood(readState);
    }

    std::error_code tryPost(std::function<void(void)>&& func);

    void adminFwDownloadChunk(
        nvme_mi_ctrl_t ctrl, const std::string& firmwarefile, size_t size,
        size_t offset, int attemptCount,
        std::function<void(const std::error_code&, nvme_status_field)>&& cb);

    void getTelemetryLogChunk(
        nvme_mi_ctrl_t ctrl, bool host, uint64_t offset,
        std::vector<uint8_t>&& data,
        std::function<void(const std::error_code&, std::span<uint8_t>)>&& cb);

    void processNextLockdownCommand(
        nvme_mi_ctrl_t ctrl, uint8_t prohibit,
        std::deque<std::pair<uint8_t, uint8_t>> lockdownQueue,
        const std::function<void(const std::error_code&, int, uint32_t,
                                 const std::string&, uint32_t)>& finalCb,
        int lastNvmeStatus, uint32_t lastCompletionDw0, int attempt,
        const std::string& lastFailureScope, uint32_t lastFailureId);

    static size_t getBlockSize(nvme_mi_ctrl_t ctrl, size_t lbaFormat);
};