#pragma once

#include "NVMeMi.hpp"

#include <boost/asio/spawn.hpp>
#include <sdbusplus/asio/object_server.hpp>

#include <cstdint>
#include <map>
#include <memory>
#include <string>
#include <vector>

enum class FeatureScope : std::uint8_t
{
    Controller,
    Namespace,
    Subsystem
};

class FeatureStore : public std::enable_shared_from_this<FeatureStore>
{
  public:
    FeatureStore(boost::asio::io_context& io,
                 sdbusplus::asio::object_server& objServer,
                 const std::shared_ptr<sdbusplus::asio::connection>& conn,
                 const std::string& objectPath,
                 std::shared_ptr<NVMeMiIntf> nvmeIntf, nvme_mi_ctrl_t ctrl,
                 uint32_t nsid, FeatureScope scope);

    virtual ~FeatureStore()
    {
        if (dbusIntf)
        {
            objServer.remove_interface(dbusIntf);
        }
    }

    // Initialize the feature collection by querying the device
    void init(std::function<void(const std::error_code&)>&& cb);

  private:
    boost::asio::io_context& io;
    sdbusplus::asio::object_server& objServer;
    std::shared_ptr<sdbusplus::asio::connection> conn;
    std::string objectPath;
    std::shared_ptr<NVMeMiIntf> nvmeIntf;
    nvme_mi_ctrl_t ctrl;
    uint32_t nsid;
    FeatureScope scope;

    bool vwcSupported = false;
    uint16_t oncs = 0;

    // Map of FID to its supported capabilities (from SEL=011b)
    std::map<uint8_t, uint32_t> featureCaps;

    std::shared_ptr<sdbusplus::asio::dbus_interface> dbusIntf;
    std::vector<std::tuple<std::string, std::string, std::string>>
        featureCollection;

    // Helper to determine feature collection
    void discoverFeatures(const boost::asio::yield_context& yield);
    void fetchIdentifyData(const boost::asio::yield_context& yield);

    void registerDbusInterface();

    std::vector<uint8_t> getFeature(boost::asio::yield_context yield,
                                    const std::string& featureName);
    void setFeature(boost::asio::yield_context yield,
                    const std::string& featureName,
                    const std::vector<uint8_t>& data);
};