#pragma once
/**
 * This file contains all cache implementation of types of Metrics
 */

#include "NVMeCache.hpp"
#include "NVMeController.hpp"
#include "NVMeVolume.hpp"

/**
 * Metric implementation for LogPage
 */

template <nvme_cmd_get_log_lid LID, class ClockType, uint8_t DA>
class ControllerLogPageMetric : public MetricBase<ClockType>
{
  public:
    ControllerLogPageMetric(std::weak_ptr<NVMeControllerEnabled> controller,
                            std::shared_ptr<Scheduler<ClockType>> scheduler,
                            ClockType::duration interval) :
        MetricBase<ClockType>(scheduler, interval),
        controller(std::move(controller))
    {}

    /* interface overrides */
    static constexpr std::string_view logpageIdentifier() noexcept
    {
        if constexpr (LID == NVME_LOG_LID_ERROR)
        {
            return "ErrorInfoMetric";
        }
        else if constexpr (LID == NVME_LOG_LID_SMART)
        {
            return "SMARTMetric";
        }
        else if constexpr (LID == NVME_LOG_LID_FW_SLOT)
        {
            return "FwSlotInfoMetric";
        }
        else if constexpr (LID == NVME_LOG_LID_DEVICE_SELF_TEST)
        {
            return "DeviceSelfTestMetric";
        }
        else if constexpr (LID == NVME_LOG_LID_TELEMETRY_HOST)
        {
            if constexpr (DA >= 1 && DA <= 4)
            {
                constexpr std::array<std::string_view, 5> daStrings = {
                    "", // DA=0 is not used in this range
                    "DriveHostInitiated_DA1", "DriveHostInitiated_DA2",
                    "DriveHostInitiated_DA3", "DriveHostInitiated_DA4"};
                return daStrings[DA];
            }
            else
            {
                return "HostInitiatedMetric";
            }
        }
        else if constexpr (LID == NVME_LOG_LID_CMD_AND_FEAT_LOCKDOWN)
        {
            return "LockdownMetric";
        }
        else
        {
            static_assert(false,
                          "LID not supported by ControllerLogPageMetric");
        }

        return {};
    }

    constexpr std::string_view getIdentifier() const noexcept override
    {
        return logpageIdentifier();
    }

    bool isCacheValid() const noexcept override
    {
        return cacheValid;
    }
    std::tuple<std::chrono::time_point<ClockType>,
               std::chrono::time_point<ClockType>, std::span<const uint8_t>>
        getCache() const noexcept override
    {
        std::chrono::time_point<ClockType> start = std::get<0>(cache);
        std::chrono::time_point<ClockType> stop = std::get<1>(cache);
        auto& data = std::get<2>(cache);
        return {start, stop,
                std::span<const uint8_t>(data.begin(), data.end())};
    }

    /* internal overrides */
  private:
    // read data from device and update the
    // implemetation should extend this virtual function and do the real
    // data fetch from device.
    void readDevice(std::function<void(std::error_code ec, size_t size,
                                       bool complete)>&& cb) noexcept override
    {
        auto startTime = ClockType::now();

        auto cntrlPtr = controller.lock();
        assert(cntrlPtr &&
               cntrlPtr->status == NVMeControllerEnabled::Status::Enabled &&
               "NVMe Metric refresh requires enabled controller");
        assert(cntrlPtr->nvmeIntf &&
               "NVMe Metric refresh requires active MI interface");

        cntrlPtr->nvmeIntf->adminGetLogPage(
            cntrlPtr->nvmeCtrl, LID, 0, 0, 0,
            [weakSelf{std::move(this->weak_from_this())}, startTime,
             path{cntrlPtr->path}, id{getIdentifier()}, cb{std::move(cb)}](
                const std::error_code& ec, std::span<uint8_t> data) {
            auto self = std::dynamic_pointer_cast<
                ControllerLogPageMetric<LID, ClockType, DA>>(weakSelf.lock());
            if (!self)
            {
                lg2::info("metric released before callback", "PATH", path, "ID",
                          id);
                // The Scheduler will clear the invalid cache Task
                cb({}, 4096, true);
                return;
            }

            if (ec)
            {
                if (self->errorCount <
                    ControllerLogPageMetric<LID, ClockType, DA>::errorCountMax)
                {
                    lg2::warning("fail to fresh log page at retry({RETRY})",
                                 "PATH", path, "ID", id, "RETRY",
                                 self->errorCount++);

                    // Schedule the retry after 1s
                    // TODO: costomize the retry delay for each log pages.
                    cb(std::make_error_code(std::errc::io_error), 4096, false);
                    return;
                }
                lg2::error("retry exhausted to fresh log page", "PATH", path,
                           "ID", id);
                self->cacheValid = false;

                // Notify the RefreshCB with error code; the task will
                // be reschuduled after this->interval
                cb(std::make_error_code(std::errc::io_error), 4096, true);
                return;
            }
            self->errorCount = 0;
            self->cacheValid = true;
            std::get<0>(self->cache) = startTime;
            std::get<1>(self->cache) = ClockType::now();

            auto& cacheData = std::get<2>(self->cache);
            cacheData.clear();
            cacheData.insert(cacheData.begin(), data.begin(), data.end());

            cb({}, data.size(), true);
        },
            DA);
    }
    /* member variables */
    std::weak_ptr<NVMeControllerEnabled> controller;

    std::tuple<std::chrono::time_point<ClockType>,
               std::chrono::time_point<ClockType>, std::vector<uint8_t>>
        cache;
    bool cacheValid{false};
    unsigned errorCount{0};
    static constexpr unsigned errorCountMax = 5;
};

/**
 * Metric implementation for Identify info
 */
template <nvme_identify_cns CNS, class ClockType>
class IdentifyMetric : public MetricBase<ClockType>
{
  public:
    template <class T>
        requires((CNS == NVME_IDENTIFY_CNS_CTRL &&
                  std::is_base_of_v<NVMeControllerEnabled, T>) ||
                 (CNS == NVME_IDENTIFY_CNS_NS &&
                  std::is_base_of_v<NVMeVolume, T>))
    IdentifyMetric(std::shared_ptr<T> resource,
                   std::shared_ptr<Scheduler<ClockType>> scheduler,
                   ClockType::duration interval) :
        MetricBase<ClockType>(scheduler, interval)
    {
        assert(resource &&
               "Fail to construct Identify Metric on empty resource");
        path = resource->path;
        if constexpr (CNS == NVME_IDENTIFY_CNS_CTRL)
        {
            controller = resource;
            namespaceId = 0;
        }
        else if constexpr (CNS == NVME_IDENTIFY_CNS_NS)
        {
            auto subsys = resource->subsys.lock();
            assert(subsys &&
                   "Fail to construct NS Identify Metric on expired subsystem");
            controller = subsys->getPrimaryController();
            namespaceId = resource->namespaceId();
        }
        else
        {
            static_assert(false, "unsupported CNS for IdentifyMetric");
        }
    }

    /* interface overrides */
    inline static constexpr std::string_view identifyIdentifier() noexcept
    {
        return "IdentifyMetric";
    }
    constexpr std::string_view getIdentifier() const noexcept override
    {
        return identifyIdentifier();
    }

    bool isCacheValid() const noexcept override
    {
        return cacheValid;
    }
    std::tuple<std::chrono::time_point<ClockType>,
               std::chrono::time_point<ClockType>, std::span<const uint8_t>>
        getCache() const noexcept override
    {
        std::chrono::time_point<ClockType> start = std::get<0>(cache);
        std::chrono::time_point<ClockType> stop = std::get<1>(cache);
        auto& data = std::get<2>(cache);
        return {start, stop,
                std::span<const uint8_t>(data.begin(), data.end())};
    }

    /* internal overrides */
  private:
    // read data from device and update the
    // implemetation should extend this virtual function and do the real
    // data fetch from device.
    void readDevice(std::function<void(std::error_code ec, size_t size,
                                       bool complete)>&& cb) noexcept override
    {
        auto startTime = ClockType::now();

        auto cntrlPtr = controller.lock();
        assert(cntrlPtr &&
               cntrlPtr->status == NVMeControllerEnabled::Status::Enabled &&
               "NVMe Metric refresh requires enabled controller");
        assert(cntrlPtr->nvmeIntf &&
               "NVMe Metric refresh requires active MI interface");

        cntrlPtr->nvmeIntf->adminIdentify(
            cntrlPtr->nvmeCtrl, CNS, namespaceId,
            0 /* CNTID will not be used for CNS = 00h/01h */,
            [weakSelf{std::move(this->weak_from_this())}, startTime, path{path},
             id{getIdentifier()}, cb{std::move(cb)}](const nvme_ex_ptr& ex,
                                                     std::span<uint8_t> data) {
            auto self =
                std::dynamic_pointer_cast<IdentifyMetric<CNS, ClockType>>(
                    weakSelf.lock());
            if (!self)
            {
                lg2::info("metric released before callback", "PATH", path, "ID",
                          id);
                // The Scheduler will clear the invalid cache Task
                cb({}, 4096, true);
                return;
            }
            if (ex)
            {
                if (self->errorCount <
                    IdentifyMetric<CNS, ClockType>::errorCountMax)
                {
                    lg2::warning(
                        "fail to fresh identify at retry({RETRY}): {ERROR}",
                        "PATH", path, "ID", id, "RETRY", self->errorCount++,
                        "ERROR", ex->description());

                    // Schedule the retry after 1s
                    cb(std::make_error_code(std::errc::io_error), 4096, false);
                    return;
                }
                lg2::error("retry exhausted to fresh identify: {ERROR}", "PATH",
                           path, "ID", id, "ERROR", ex->description());
                self->cacheValid = false;

                // Notify the RefreshCB with error code; the task will
                // be reschuduled after this->interval
                cb(std::make_error_code(std::errc::io_error), 4096, true);
                return;
            }
            self->errorCount = 0;
            self->cacheValid = true;
            std::get<0>(self->cache) = startTime;
            std::get<1>(self->cache) = ClockType::now();

            auto& cacheData = std::get<2>(self->cache);
            cacheData.clear();
            cacheData.insert(cacheData.begin(), data.begin(), data.end());

            cb({}, data.size(), true);
        });
    }

    /* member variables */
    std::string path;
    std::weak_ptr<NVMeControllerEnabled> controller;
    uint32_t namespaceId;

    std::tuple<std::chrono::time_point<ClockType>,
               std::chrono::time_point<ClockType>, std::vector<uint8_t>>
        cache;
    bool cacheValid{false};
    unsigned errorCount{0};
    static constexpr unsigned errorCountMax = 5;
};