#include "common/instance_id.hpp"
#include "common/utils.hpp"
#include "libpldmresponder/base.hpp"
#include "test/test_instance_id.hpp"

#include <libpldm/base.h>
#include <libpldm/utils.h>

#include <sdeventplus/event.hpp>

#include <algorithm>
#include <array>
#include <cstring>
#include <map>
#include <vector>

#include <gtest/gtest.h>

using namespace pldm::responder;

class TestBaseCommands : public testing::Test
{
  protected:
    TestBaseCommands() : event(sdeventplus::Event::get_default()) {}

    sdeventplus::Event event;
};

TEST_F(TestBaseCommands, testPLDMTypesGoodRequest)
{
    std::array<uint8_t, sizeof(pldm_msg_hdr)> requestPayload{};
    auto request = new (requestPayload.data()) pldm_msg;
    // payload length will be 0 in this case
    size_t requestPayloadLength = 0;
    base::Handler handler(event);
    auto response = handler.getPLDMTypes(request, requestPayloadLength);
    // Need to support OEM type.
    auto responsePtr = new (response.data()) pldm_msg;
    uint8_t* payload_ptr = responsePtr->payload;
    ASSERT_EQ(payload_ptr[0], 0);
    ASSERT_EQ(payload_ptr[1], 157); // 0b10011101 see DSP0240 table11
    ASSERT_EQ(payload_ptr[2], 0);
}

TEST_F(TestBaseCommands, testGetPLDMCommandsGoodRequest)
{
    // Need to support OEM type commands.
    std::array<uint8_t, sizeof(pldm_msg_hdr) + PLDM_GET_COMMANDS_REQ_BYTES>
        requestPayload{};
    auto request = new (requestPayload.data()) pldm_msg;
    size_t requestPayloadLength = requestPayload.size() - sizeof(pldm_msg_hdr);
    base::Handler handler(event);
    auto response = handler.getPLDMCommands(request, requestPayloadLength);
    auto responsePtr = new (response.data()) pldm_msg;
    uint8_t* payload_ptr = responsePtr->payload;
    ASSERT_EQ(payload_ptr[0], 0);
    ASSERT_EQ(payload_ptr[1], 188); // 188 = 0b10111100
    ASSERT_EQ(payload_ptr[2], 0);
}

TEST_F(TestBaseCommands, testGetPLDMCommandsBadRequest)
{
    std::array<uint8_t, sizeof(pldm_msg_hdr) + PLDM_GET_COMMANDS_REQ_BYTES>
        requestPayload{};
    auto request = new (requestPayload.data()) pldm_msg;

    request->payload[0] = 0xFF;
    size_t requestPayloadLength = requestPayload.size() - sizeof(pldm_msg_hdr);
    base::Handler handler(event);
    auto response = handler.getPLDMCommands(request, requestPayloadLength);
    auto responsePtr = new (response.data()) pldm_msg;
    uint8_t* payload_ptr = responsePtr->payload;
    ASSERT_EQ(payload_ptr[0], PLDM_ERROR_INVALID_PLDM_TYPE);
}

TEST_F(TestBaseCommands, testGetPLDMVersionGoodRequest)
{
    std::array<uint8_t, sizeof(pldm_msg_hdr) + PLDM_GET_VERSION_REQ_BYTES>
        requestPayload{};
    auto request = new (requestPayload.data()) pldm_msg;
    size_t requestPayloadLength = requestPayload.size() - sizeof(pldm_msg_hdr);

    uint8_t pldmType = PLDM_BASE;
    uint32_t transferHandle = 0x0;
    uint8_t flag = PLDM_GET_FIRSTPART;
    uint8_t retFlag = PLDM_START_AND_END;
    ver32_t version = {0x00, 0xF0, 0xF0, 0xF1};

    auto rc =
        encode_get_version_req(0, transferHandle, flag, pldmType, request);

    ASSERT_EQ(0, rc);

    base::Handler handler(event);
    auto response = handler.getPLDMVersion(request, requestPayloadLength);
    auto responsePtr = new (response.data()) pldm_msg;

    ASSERT_EQ(responsePtr->payload[0], 0);
    ASSERT_EQ(0, memcmp(responsePtr->payload + sizeof(responsePtr->payload[0]),
                        &transferHandle, sizeof(transferHandle)));
    ASSERT_EQ(0, memcmp(responsePtr->payload + sizeof(responsePtr->payload[0]) +
                            sizeof(transferHandle),
                        &retFlag, sizeof(flag)));
    ASSERT_EQ(0, memcmp(responsePtr->payload + sizeof(responsePtr->payload[0]) +
                            sizeof(transferHandle) + sizeof(flag),
                        &version, sizeof(version)));

    uint32_t expectedChecksum = pldm_edac_crc32(&version, sizeof(version));
    uint32_t responseChecksum = 0;
    memcpy(&responseChecksum,
           responsePtr->payload + PLDM_GET_VERSION_RESP_BYTES,
           sizeof(responseChecksum));
    ASSERT_EQ(expectedChecksum, le32toh(responseChecksum));
}

TEST_F(TestBaseCommands, testGetPLDMVersionBadRequest)
{
    std::array<uint8_t, sizeof(pldm_msg_hdr) + PLDM_GET_VERSION_REQ_BYTES>
        requestPayload{};
    auto request = new (requestPayload.data()) pldm_msg;
    size_t requestPayloadLength = requestPayload.size() - sizeof(pldm_msg_hdr);

    uint8_t pldmType = 6;
    uint32_t transferHandle = 0x0;
    uint8_t flag = PLDM_GET_FIRSTPART;

    auto rc =
        encode_get_version_req(0, transferHandle, flag, pldmType, request);

    ASSERT_EQ(0, rc);

    base::Handler handler(event);
    auto response = handler.getPLDMVersion(request, requestPayloadLength - 1);
    auto responsePtr = new (response.data()) pldm_msg;

    ASSERT_EQ(responsePtr->payload[0], PLDM_ERROR_INVALID_LENGTH);

    request = new (requestPayload.data()) pldm_msg;
    requestPayloadLength = requestPayload.size() - sizeof(pldm_msg_hdr);

    rc = encode_get_version_req(0, transferHandle, flag, pldmType, request);

    ASSERT_EQ(0, rc);

    response = handler.getPLDMVersion(request, requestPayloadLength);
    responsePtr = new (response.data()) pldm_msg;

    ASSERT_EQ(responsePtr->payload[0], PLDM_ERROR_INVALID_PLDM_TYPE);
}

TEST_F(TestBaseCommands, testGetTIDGoodRequest)
{
    std::array<uint8_t, sizeof(pldm_msg_hdr)> requestPayload{};
    auto request = new (requestPayload.data()) pldm_msg;
    size_t requestPayloadLength = 0;

    base::Handler handler(event);
    handler.setOemPlatformHandler(nullptr);
    auto response = handler.getTID(request, requestPayloadLength);

    auto responsePtr = new (response.data()) pldm_msg;
    uint8_t* payload = responsePtr->payload;

    ASSERT_EQ(payload[0], 0);
    ASSERT_EQ(payload[1], 1);
}

TEST_F(TestBaseCommands, testSetMultiPartSizeCallback)
{
    base::Handler handler(event);

    // Prepare a NegotiateTransferParameters request
    std::vector<uint8_t> requestMsg(
        sizeof(pldm_msg_hdr) +
        PLDM_BASE_NEGOTIATE_TRANSFER_PARAMETERS_REQ_BYTES);
    auto request = reinterpret_cast<pldm_msg*>(requestMsg.data());
    size_t requestPayloadLength = requestMsg.size() - sizeof(pldm_msg_hdr);

    struct pldm_base_negotiate_transfer_params_req req{};
    req.requester_part_size = 1024;
    // PLDM_FILE is type 7. index = 7/8 = 0. bit = 7%8 = 7.
    req.requester_protocol_support[0].byte |= (1 << 7);

    EXPECT_EQ(encode_pldm_base_negotiate_transfer_params_req(
                  /*instance_id=*/1, &req, request, &requestPayloadLength),
              0);

    // Call getNegotiateTransferParam to trigger the callback
    requestPayloadLength = requestMsg.size() - sizeof(pldm_msg_hdr);
    auto response =
        handler.getNegotiateTransferParam(1, request, requestPayloadLength);

    // Verify the response
    ASSERT_FALSE(response.empty());
    auto responsePtr = reinterpret_cast<pldm_msg*>(response.data());

    struct pldm_base_negotiate_transfer_params_resp resp;
    EXPECT_EQ(decode_pldm_base_negotiate_transfer_params_resp(
                  responsePtr, response.size() - sizeof(pldm_msg_hdr), &resp),
              0);

    EXPECT_EQ(resp.completion_code, PLDM_SUCCESS);
    EXPECT_EQ(resp.responder_part_size, PLDM_RESPONDER_PART_SIZE);
    EXPECT_TRUE(resp.responder_protocol_support[0].byte & (1 << 7));
}

TEST_F(TestBaseCommands, testNegotiateTransferParamMultipleTIDs)
{
    base::Handler handler(event);

    // Map of TIDs to their requested part sizes
    std::map<pldm_tid_t, uint16_t> tidToPartSize = {
        {10, 512},   // smaller than responder
        {20, 8192},  // larger than responder
        {30, 4096},  // equal to responder
        {40, 1024}}; // smaller than responder

    for (const auto& [tid, requester_part_size] : tidToPartSize)
    {
        // Prepare a NegotiateTransferParameters request
        std::vector<uint8_t> requestMsg(
            sizeof(pldm_msg_hdr) +
            PLDM_BASE_NEGOTIATE_TRANSFER_PARAMETERS_REQ_BYTES);
        auto request = reinterpret_cast<pldm_msg*>(requestMsg.data());
        size_t requestPayloadLength = requestMsg.size() - sizeof(pldm_msg_hdr);

        // Fill in the request payload
        auto payload =
            reinterpret_cast<pldm_base_negotiate_transfer_params_req*>(
                request->payload);
        payload->requester_part_size = htole16(requester_part_size);

        // Set the bit for PLDM_FILE_TRANSFER in requester_protocol_support
        memset(payload->requester_protocol_support, 0,
               sizeof(payload->requester_protocol_support));
        // PLDM_FILE is type 7. index = 7/8 = 0. bit = 7%8 = 7.
        payload->requester_protocol_support[0].byte |= (1 << 7);

        // Call getNegotiateTransferParam
        auto response = handler.getNegotiateTransferParam(tid, request,
                                                          requestPayloadLength);

        // Verify the response
        ASSERT_FALSE(response.empty()) << "Failed for TID " << (int)tid;
        auto responsePtr = reinterpret_cast<pldm_msg*>(response.data());

        struct pldm_base_negotiate_transfer_params_resp resp;
        EXPECT_EQ(decode_pldm_base_negotiate_transfer_params_resp(
                      responsePtr, response.size() - sizeof(pldm_msg_hdr),
                      &resp),
                  0);

        EXPECT_EQ(resp.completion_code, PLDM_SUCCESS)
            << "Failed for TID " << (int)tid;

        // Verify the negotiated part size for this TID
        uint16_t expectedNegotiatedSize =
            std::min(requester_part_size, (uint16_t)PLDM_RESPONDER_PART_SIZE);
        uint16_t negotiatedSize = handler.getNegotiatedPartSize(tid, PLDM_FILE);
        EXPECT_EQ(negotiatedSize, expectedNegotiatedSize)
            << "Failed for TID " << (int)tid;
    }
}

TEST_F(TestBaseCommands, testNegotiateTransferParamInvalidLength)
{
    std::vector<uint8_t> requestMsg(sizeof(pldm_msg_hdr) + 1);
    auto request = reinterpret_cast<pldm_msg*>(requestMsg.data());
    size_t payloadLength = requestMsg.size() - sizeof(pldm_msg_hdr);

    base::Handler handler(event);
    auto response =
        handler.getNegotiateTransferParam(1, request, payloadLength);
    auto responsePtr = reinterpret_cast<pldm_msg*>(response.data());

    struct pldm_base_negotiate_transfer_params_resp resp;
    EXPECT_EQ(decode_pldm_base_negotiate_transfer_params_resp(
                  responsePtr, response.size() - sizeof(pldm_msg_hdr), &resp),
              0);

    EXPECT_EQ(resp.completion_code, PLDM_ERROR);
}

TEST_F(TestBaseCommands, testNegotiateTransferParamInvalidPartSizeZero)
{
    std::vector<uint8_t> requestMsg(
        sizeof(pldm_msg_hdr) +
        PLDM_BASE_NEGOTIATE_TRANSFER_PARAMETERS_REQ_BYTES);
    auto request = reinterpret_cast<pldm_msg*>(requestMsg.data());
    size_t payloadLength = requestMsg.size() - sizeof(pldm_msg_hdr);

    auto payload = reinterpret_cast<pldm_base_negotiate_transfer_params_req*>(
        request->payload);
    payload->requester_part_size = 0;

    base::Handler handler(event);
    auto response =
        handler.getNegotiateTransferParam(1, request, payloadLength);
    auto responsePtr = reinterpret_cast<pldm_msg*>(response.data());

    struct pldm_base_negotiate_transfer_params_resp resp;
    EXPECT_EQ(decode_pldm_base_negotiate_transfer_params_resp(
                  responsePtr, response.size() - sizeof(pldm_msg_hdr), &resp),
              0);

    EXPECT_EQ(resp.completion_code, PLDM_ERROR_INVALID_DATA);
}

TEST_F(TestBaseCommands, testNegotiateTransferParamInvalidPartSizeNotPowerOfTwo)
{
    std::vector<uint8_t> requestMsg(
        sizeof(pldm_msg_hdr) +
        PLDM_BASE_NEGOTIATE_TRANSFER_PARAMETERS_REQ_BYTES);
    auto request = reinterpret_cast<pldm_msg*>(requestMsg.data());
    size_t payloadLength = requestMsg.size() - sizeof(pldm_msg_hdr);

    auto payload = reinterpret_cast<pldm_base_negotiate_transfer_params_req*>(
        request->payload);
    payload->requester_part_size = 3000;

    base::Handler handler(event);
    auto response =
        handler.getNegotiateTransferParam(1, request, payloadLength);
    auto responsePtr = reinterpret_cast<pldm_msg*>(response.data());
    struct pldm_base_negotiate_transfer_params_resp resp;
    EXPECT_EQ(decode_pldm_base_negotiate_transfer_params_resp(
                  responsePtr, response.size() - sizeof(pldm_msg_hdr), &resp),
              0);

    EXPECT_EQ(resp.completion_code, PLDM_ERROR_INVALID_DATA);
}

TEST_F(TestBaseCommands, testNegotiateTransferParamGoodRequest)
{
    std::vector<uint8_t> requestMsg(
        sizeof(pldm_msg_hdr) +
        PLDM_BASE_NEGOTIATE_TRANSFER_PARAMETERS_REQ_BYTES);
    auto request = reinterpret_cast<pldm_msg*>(requestMsg.data());
    size_t payloadLength = requestMsg.size() - sizeof(pldm_msg_hdr);

    auto payload = reinterpret_cast<pldm_base_negotiate_transfer_params_req*>(
        request->payload);
    payload->requester_part_size = htole16(4096);

    base::Handler handler(event);
    auto response =
        handler.getNegotiateTransferParam(1, request, payloadLength);
    auto responsePtr = reinterpret_cast<pldm_msg*>(response.data());

    struct pldm_base_negotiate_transfer_params_resp resp;
    EXPECT_EQ(decode_pldm_base_negotiate_transfer_params_resp(
                  responsePtr, response.size() - sizeof(pldm_msg_hdr), &resp),
              0);

    EXPECT_EQ(resp.completion_code, PLDM_SUCCESS);
    EXPECT_EQ(resp.responder_part_size, PLDM_RESPONDER_PART_SIZE);
}