/*
 * Copyright (C) 2026 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
 */

#pragma once
#include "level_zero/sysman/source/shared/linux/nl_api/sysman_drm_ras_types.h"
#include "level_zero/sysman/source/shared/linux/nl_api/sysman_drm_ras_uapi_ext.h"
#include "level_zero/sysman/source/shared/linux/nl_api/sysman_nl_api.h"
#include <level_zero/zes_api.h>

#include <linux/types.h>
#include <vector>

namespace L0 {
namespace Sysman {

class Operation;

class DrmNlApi {
  public:
    DrmNlApi();
    virtual ~DrmNlApi();

    MOCKABLE_VIRTUAL ze_result_t listNodes(std::vector<DrmRasNode> &nodeList);
    MOCKABLE_VIRTUAL ze_result_t getErrorCounter(const uint32_t &nodeId, const uint32_t &errorId, DrmErrorCounter &errorCounter);
    MOCKABLE_VIRTUAL ze_result_t getErrorsList(const uint32_t &nodeId, std::vector<DrmErrorCounter> &errorList);
    MOCKABLE_VIRTUAL ze_result_t setErrorThreshold(const uint32_t &nodeId, const uint32_t &errorId, const uint32_t &threshold);
    MOCKABLE_VIRTUAL ze_result_t getErrorThreshold(const uint32_t &nodeId, const uint32_t &errorId, DrmErrorThreshold &threshold);
    MOCKABLE_VIRTUAL ze_result_t clearErrorCounter(const uint32_t &nodeId, const uint32_t &errorId);

    MOCKABLE_VIRTUAL ze_result_t subscribeToEvents();
    MOCKABLE_VIRTUAL ze_result_t unsubscribeFromEvents();
    MOCKABLE_VIRTUAL ze_result_t pollEvent(DrmRasEvent &event);
    MOCKABLE_VIRTUAL int getEventSocketFd() const;
    MOCKABLE_VIRTUAL bool isSubscribedToEvents() const { return eventSubscribed; }

    int handleMsg(struct nl_msg *msg);
    int handleAck(struct nl_msg *msg);
    int handleEventMsg(struct nl_msg *msg);
    ze_result_t listNodesRsp(struct nl_cache_ops *ops, struct genl_cmd *cmd, struct genl_info *info);
    ze_result_t getErrorListRsp(struct nl_cache_ops *ops, struct genl_cmd *cmd, struct genl_info *info);
    ze_result_t getSingleErrorCounterRsp(struct nl_cache_ops *ops, struct genl_cmd *cmd, struct genl_info *info);
    ze_result_t getErrorThresholdRsp(struct nl_cache_ops *ops, struct genl_cmd *cmd, struct genl_info *info);

  protected:
    struct Operation {
        void *pOutput;
        ze_result_t result = ZE_RESULT_ERROR_UNKNOWN;
        uint16_t cmdOp;
        bool done = false;
        bool parseSingleErrorCounter = false;
        Operation(uint16_t cmdOp, void *pOutput) : pOutput(pOutput), cmdOp(cmdOp) {}
    };

    ze_result_t initConnection();
    ze_result_t parseEventMessage(struct nl_msg *msg, DrmRasEvent &event);
    ze_result_t initEventSocket();
    void cleanupConnection();
    void cleanupEventSocket();

    std::unique_ptr<NlApi> pNlApi;
    std::unique_ptr<Operation> currentOperation = nullptr;
    struct nl_sock *nlEventSock = nullptr;
    int familyId = 0;
    int eventGroupId = -1;
    bool eventReady = false;
    bool isInitDone = false;
    bool isDrmRasFamilyRegistered = false;
    bool eventSubscribed = false;

  private:
    ze_result_t initSocketBase(struct nl_sock *&sock, int &resolvedFamilyId);
    void freeNlSocket(struct nl_sock *&sock);
    ze_result_t allocMsg(const uint16_t &cmdOp, const bool useDumpFlags, struct nl_msg *&msg);
    ze_result_t issueRequestReadOne(const uint16_t cmdOp, const uint32_t &nodeId, const uint32_t &errorId, void *pOutput);
    ze_result_t issueRequestQueryErrors(const uint16_t &cmdOp, const uint32_t &nodeId, void *pOutput);
    ze_result_t issueRequestSetThreshold(const uint32_t &nodeId, const uint32_t &errorId, const uint32_t &threshold);
    ze_result_t issueRequestGetThreshold(const uint32_t &nodeId, const uint32_t &errorId, void *pOutput);
    ze_result_t performTransaction(const uint16_t &cmdOp, struct nl_msg *msg, void *pOutput, const bool parseSingleErrorCounter);
    ze_result_t issueRequestListNodes(void *pOutput);
    ze_result_t issueRequestClearErrorCounter(const uint32_t &nodeId, const uint32_t &errorId);
    void setupNlOperations();

    struct genl_cmd newCmds[NEO_DRM_RAS_CMD_MAX] = {};
    struct nla_policy nodePolicy[DRM_RAS_A_NODE_ATTRS_MAX + 1] = {};
    struct nla_policy errorPolicy[DRM_RAS_A_ERROR_COUNTER_ATTRS_MAX + 1] = {};
    struct nla_policy thresholdPolicy[DRM_RAS_A_ERROR_THRESHOLD_ATTRS_MAX + 1] = {};
    struct genl_ops ops = {};
    std::string deviceName = {};
    DrmRasEvent pendingEvent = {};
    struct nl_sock *nlSock = nullptr;
};

} // namespace Sysman
} // namespace L0
