// Copyright © Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#include "DescriptorTestUtils.hpp"
#include "HipdnnException.hpp"
#include "TestMacros.hpp"
#include "descriptors/EngineConfigDescriptor.hpp"
#include "descriptors/ExecutionPlanDescriptor.hpp"
#include "descriptors/ScopedDescriptor.hpp"
#include "hipdnn_backend.h"
#include "mocks/MockDescriptor.hpp"
#include "mocks/MockEnginePluginResourceManager.hpp"
#include "mocks/MockHandle.hpp"

#include <flatbuffers/flatbuffer_builder.h>
#include <gtest/gtest.h>
#include <hipdnn_flatbuffers_sdk/data_objects/execution_plan_generated.h>
#include <hipdnn_flatbuffers_sdk/data_objects/graph_generated.h>

#include <array>
#include <memory>

using namespace hipdnn_backend;
using namespace plugin;
using namespace hipdnn_backend::test_utilities;
using namespace ::testing;

using ::testing::Return;

class TestExecutionPlanDescriptor : public ::testing::Test
{
public:
    std::shared_ptr<ExecutionPlanDescriptor> getExecutionPlanDescriptor() const
    {
        return _planWrapper->asDescriptor<ExecutionPlanDescriptor>();
    }

    std::shared_ptr<MockGraphDescriptor> getMockGraph() const
    {
        return MockDescriptorUtility::asDescriptorUnsafe<MockGraphDescriptor>(
            _mockGraphWrapper.get());
    }

    std::shared_ptr<MockEngineDescriptor> getMockEngine() const
    {
        return MockDescriptorUtility::asDescriptorUnsafe<MockEngineDescriptor>(
            _mockEngineWrapper.get());
    }

    std::shared_ptr<MockEngineConfigDescriptor> getMockEngineConfig() const
    {
        return MockDescriptorUtility::asDescriptorUnsafe<MockEngineConfigDescriptor>(
            _mockEngineConfigWrapper.get());
    }

    std::shared_ptr<MockEngineConfigDescriptor> getMockEngineConfigBadType() const
    {
        return MockDescriptorUtility::asDescriptorUnsafe<MockEngineConfigDescriptor>(
            _mockEngineConfigBadTypeWrapper.get());
    }

    static hipdnnEnginePluginExecutionContext_t getExecutionContext()
    {
        return reinterpret_cast<hipdnnEnginePluginExecutionContext_t>(0xFFFFFFFF);
    }

    void setHandle()
    {
        EXPECT_CALL(*_mockEnginePluginResourceManager, createExecutionContext(_, _, _))
            .WillOnce(Return(getExecutionContext()));
        EXPECT_CALL(*_mockEnginePluginResourceManager, destroyExecutionContext(_, _));
        EXPECT_CALL(*_mockEnginePluginResourceManager, getWorkspaceSize(_, _))
            .WillOnce(Return(1024));

        EXPECT_CALL(*_mockHandle, getPluginResourceManager())
            .WillOnce(Return(_mockEnginePluginResourceManager));
        getExecutionPlanDescriptor()->setAttribute(
            HIPDNN_ATTR_EXECUTION_PLAN_HANDLE, HIPDNN_TYPE_HANDLE, 1, &_mockHandle);
    }

    void setEngineConfig()
    {
        EXPECT_CALL(*getMockGraph(), getHandle()).WillRepeatedly(Return(_mockHandle.get()));
        EXPECT_CALL(*getMockGraph(), getSerializedGraph()).WillRepeatedly(Invoke([this]() {
            return hipdnnPluginConstData_t{_serializedGraph.data(), _serializedGraph.size()};
        }));
        EXPECT_CALL(*getMockGraph(), isOverrideShapeEnabled()).WillRepeatedly(Return(false));
        EXPECT_CALL(*getMockEngine(), getEngineId()).WillOnce(Return(ENGINE_ID));
        EXPECT_CALL(*getMockEngine(), getGraph()).WillOnce(Return(getMockGraph()));

        EXPECT_CALL(*getMockEngineConfig(), isFinalized()).WillOnce(Return(true));
        EXPECT_CALL(*getMockEngineConfig(), getEngine()).WillOnce(Return(getMockEngine()));
        EXPECT_CALL(*getMockEngineConfig(), getSerializedEngineConfig()).WillOnce(Invoke([]() {
            return hipdnnPluginConstData_t{nullptr, 0};
        }));

        getExecutionPlanDescriptor()->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                                                   HIPDNN_TYPE_BACKEND_DESCRIPTOR,
                                                   1,
                                                   &_mockEngineConfigWrapper);
    }

    void makeExecutionPlanFinalized()
    {
        setHandle();
        setEngineConfig();
        ASSERT_NO_THROW(getExecutionPlanDescriptor()->finalize());
    }

    // Default version must match the current PLAN_SERIALIZATION_VERSION so that
    // plans built here are "supported" and reach the field-level validation.
    flatbuffers::DetachedBuffer makeSerializedPlan(uint32_t version = 2,
                                                   int64_t workspaceSize = 1024,
                                                   bool includeTensorUids = true,
                                                   bool includePluginPayload = true,
                                                   bool emptyTensorUids = false,
                                                   bool emptyPluginPayload = false,
                                                   bool isOverrideShapeEnabled = false,
                                                   bool includeTensorAlignments = true,
                                                   bool mismatchTensorAlignments = false) const
    {
        flatbuffers::FlatBufferBuilder builder;
        flatbuffers::Offset<flatbuffers::Vector<int64_t>> tensorUids;
        flatbuffers::Offset<flatbuffers::Vector<uint8_t>> pluginPayload;
        flatbuffers::Offset<flatbuffers::Vector<int64_t>> tensorAlignments;

        if(includeTensorUids)
        {
            tensorUids = emptyTensorUids ? builder.CreateVector(std::vector<int64_t>{})
                                         : builder.CreateVector(_tensorUids);
        }
        if(includePluginPayload)
        {
            auto pluginPayloadBytes
                = emptyPluginPayload ? std::vector<uint8_t>{} : std::vector<uint8_t>{4, 5, 6};
            pluginPayload = builder.CreateVector(pluginPayloadBytes);
        }
        if(includeTensorAlignments)
        {
            // A deliberately short vector exercises the size-mismatch guard.
            auto alignments
                = mismatchTensorAlignments ? std::vector<int64_t>{16} : _tensorAlignments;
            tensorAlignments = builder.CreateVector(alignments);
        }

        auto plan = hipdnn_flatbuffers_sdk::data_objects::CreateSerializedExecutionPlan(
            builder,
            version,
            ENGINE_ID,
            workspaceSize,
            tensorUids,
            pluginPayload,
            isOverrideShapeEnabled,
            tensorAlignments);
        builder.Finish(plan);
        return builder.Release();
    }

protected:
    static constexpr int64_t ENGINE_ID = 0;

    std::unique_ptr<HipdnnBackendDescriptor> _planWrapper = nullptr;
    std::unique_ptr<HipdnnBackendDescriptor> _mockGraphWrapper = nullptr;
    std::unique_ptr<HipdnnBackendDescriptor> _mockEngineWrapper = nullptr;
    std::unique_ptr<HipdnnBackendDescriptor> _mockEngineConfigWrapper = nullptr;
    std::unique_ptr<HipdnnBackendDescriptor> _mockEngineConfigBadTypeWrapper = nullptr;
    std::unique_ptr<HipdnnBackendDescriptor> _mockWrongTypeWrapper = nullptr;
    std::unique_ptr<MockHandle> _mockHandle = nullptr;
    std::shared_ptr<MockEnginePluginResourceManager> _mockEnginePluginResourceManager = nullptr;
    flatbuffers::DetachedBuffer _serializedGraph;
    std::vector<int64_t> _tensorUids{11, 22, 33};
    // Parallel to _tensorUids: alignment[i] is the required byte alignment of
    // tensor_uids[i]. Deliberately non-default (not all 16) so tests prove the
    // values are carried through, not just defaulted.
    std::vector<int64_t> _tensorAlignments{16, 64, 256};

    void SetUp() override
    {
        flatbuffers::FlatBufferBuilder builder;
        hipdnn_flatbuffers_sdk::data_objects::GraphT graph;
        for(size_t i = 0; i < _tensorUids.size(); ++i)
        {
            auto tensor
                = std::make_unique<hipdnn_flatbuffers_sdk::data_objects::TensorAttributesT>();
            tensor->uid = _tensorUids[i];
            tensor->alignment = _tensorAlignments[i];
            graph.tensors.push_back(std::move(tensor));
        }
        builder.Finish(hipdnn_flatbuffers_sdk::data_objects::Graph::Pack(builder, &graph));
        _serializedGraph = builder.Release();

        _planWrapper = createDescriptor<ExecutionPlanDescriptor>();
        _mockGraphWrapper = createDescriptor<MockGraphDescriptor>();
        _mockEngineWrapper = createDescriptor<MockEngineDescriptor>();
        _mockEngineConfigWrapper = createDescriptor<MockEngineConfigDescriptor>();
        _mockEngineConfigBadTypeWrapper = createDescriptor<MockEngineConfigDescriptor>();
        _mockWrongTypeWrapper = createDescriptor<MockEngineDescriptor>();
        _mockHandle = std::make_unique<MockHandle>();
        _mockEnginePluginResourceManager = std::make_shared<MockEnginePluginResourceManager>();
    }

    void TearDown() override {}
};

TEST_F(TestExecutionPlanDescriptor, CreateExecutionPlanDescriptor)
{
    auto plan = getExecutionPlanDescriptor();
    ASSERT_NE(plan, nullptr);
    ASSERT_FALSE(plan->isFinalized());
    ASSERT_EQ(plan->getType(), HIPDNN_BACKEND_EXECUTION_PLAN_DESCRIPTOR);
}

TEST_F(TestExecutionPlanDescriptor, SetAttrWhenNotFinalized)
{
    auto plan = getExecutionPlanDescriptor();
    uint64_t dummyWorkspaceSize = 0;

    ASSERT_THROW_HIPDNN_STATUS(
        plan->setAttribute(
            HIPDNN_ATTR_EXECUTION_PLAN_WORKSPACE_SIZE, HIPDNN_TYPE_INT64, 1, &dummyWorkspaceSize),
        HIPDNN_STATUS_NOT_SUPPORTED);

    makeExecutionPlanFinalized();

    ASSERT_THROW_HIPDNN_STATUS(
        plan->setAttribute(
            HIPDNN_ATTR_EXECUTION_PLAN_WORKSPACE_SIZE, HIPDNN_TYPE_INT64, 1, &dummyWorkspaceSize),
        HIPDNN_STATUS_NOT_INITIALIZED);
}

TEST_F(TestExecutionPlanDescriptor, SetHandle)
{
    auto plan = getExecutionPlanDescriptor();
    hipdnnHandle_t handle = nullptr;

    ASSERT_THROW_HIPDNN_STATUS(plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_HANDLE,
                                                  HIPDNN_TYPE_INT64,
                                                  1,
                                                  static_cast<const void*>(&handle)),
                               HIPDNN_STATUS_BAD_PARAM);

    ASSERT_THROW_HIPDNN_STATUS(plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_HANDLE,
                                                  HIPDNN_TYPE_HANDLE,
                                                  2,
                                                  static_cast<const void*>(&handle)),
                               HIPDNN_STATUS_BAD_PARAM);

    ASSERT_THROW_HIPDNN_STATUS(
        plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_HANDLE, HIPDNN_TYPE_HANDLE, 1, nullptr),
        HIPDNN_STATUS_BAD_PARAM_NULL_POINTER);

    ASSERT_THROW_HIPDNN_STATUS(plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_HANDLE,
                                                  HIPDNN_TYPE_HANDLE,
                                                  1,
                                                  static_cast<const void*>(&handle)),
                               HIPDNN_STATUS_BAD_PARAM_NULL_POINTER);

    handle = reinterpret_cast<hipdnnHandle_t>(0x12345678);
    plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_HANDLE,
                       HIPDNN_TYPE_HANDLE,
                       1,
                       static_cast<const void*>(&handle));
}

TEST_F(TestExecutionPlanDescriptor, SetEngineConfig)
{
    auto plan = getExecutionPlanDescriptor();
    auto mockEngineConfig = getMockEngineConfig();

    EXPECT_CALL(*getMockEngineConfigBadType(), isFinalized()).Times(1);
    EXPECT_CALL(*mockEngineConfig, isFinalized()).WillOnce(Return(false)).WillOnce(Return(true));

    // isFinalized()->false
    ASSERT_THROW_HIPDNN_STATUS(plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                                                  HIPDNN_TYPE_BACKEND_DESCRIPTOR,
                                                  1,
                                                  &_mockEngineConfigWrapper),
                               HIPDNN_STATUS_BAD_PARAM_NOT_FINALIZED);

    // isFinalized()->true
    plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                       HIPDNN_TYPE_BACKEND_DESCRIPTOR,
                       1,
                       &_mockEngineConfigWrapper);

    ASSERT_THROW_HIPDNN_STATUS(plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                                                  HIPDNN_TYPE_HANDLE,
                                                  1,
                                                  &_mockEngineConfigWrapper),
                               HIPDNN_STATUS_BAD_PARAM);

    ASSERT_THROW_HIPDNN_STATUS(plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                                                  HIPDNN_TYPE_BACKEND_DESCRIPTOR,
                                                  2,
                                                  &_mockEngineConfigWrapper),
                               HIPDNN_STATUS_BAD_PARAM);

    ASSERT_THROW_HIPDNN_STATUS(
        plan->setAttribute(
            HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG, HIPDNN_TYPE_BACKEND_DESCRIPTOR, 1, nullptr),
        HIPDNN_STATUS_BAD_PARAM_NULL_POINTER);

    hipdnnBackendDescriptor_t engineConfig = nullptr;
    ASSERT_THROW_HIPDNN_STATUS(plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                                                  HIPDNN_TYPE_BACKEND_DESCRIPTOR,
                                                  1,
                                                  &engineConfig),
                               HIPDNN_STATUS_BAD_PARAM_NULL_POINTER);

    ASSERT_THROW_HIPDNN_STATUS(plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                                                  HIPDNN_TYPE_BACKEND_DESCRIPTOR,
                                                  1,
                                                  &_mockEngineConfigBadTypeWrapper),
                               HIPDNN_STATUS_BAD_PARAM_NOT_FINALIZED);

    ASSERT_THROW_HIPDNN_STATUS(plan->setAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                                                  HIPDNN_TYPE_BACKEND_DESCRIPTOR,
                                                  1,
                                                  &_mockWrongTypeWrapper),
                               HIPDNN_STATUS_BAD_PARAM);
}

TEST_F(TestExecutionPlanDescriptor, Finalize)
{
    auto plan = getExecutionPlanDescriptor();
    ASSERT_THROW_HIPDNN_STATUS(plan->finalize(), HIPDNN_STATUS_BAD_PARAM);

    setHandle();
    setEngineConfig();

    ASSERT_NO_THROW(plan->finalize());

    ASSERT_THROW(plan->finalize(), hipdnn_backend::HipdnnException);
}

TEST_F(TestExecutionPlanDescriptor, GetAttrWhenNotFinalized)
{
    auto plan = getExecutionPlanDescriptor();
    uint64_t dummyWorkspaceSize = 0;

    ASSERT_THROW_HIPDNN_STATUS(plan->getAttribute(HIPDNN_ATTR_EXECUTION_PLAN_WORKSPACE_SIZE,
                                                  HIPDNN_TYPE_INT64,
                                                  1,
                                                  nullptr,
                                                  &dummyWorkspaceSize),
                               HIPDNN_STATUS_NOT_INITIALIZED);
}

TEST_F(TestExecutionPlanDescriptor, GetWorkspaceSize)
{
    auto plan = getExecutionPlanDescriptor();
    auto mockEngineConfig = getMockEngineConfig();
    int64_t workspaceSize = 0;

    makeExecutionPlanFinalized();
    plan->getAttribute(
        HIPDNN_ATTR_EXECUTION_PLAN_WORKSPACE_SIZE, HIPDNN_TYPE_INT64, 1, nullptr, &workspaceSize);
    ASSERT_EQ(workspaceSize, 1024);
}

TEST_F(TestExecutionPlanDescriptor, GetEngineGlobalIndexExtReturnsEngineId)
{
    auto plan = getExecutionPlanDescriptor();
    int64_t engineGlobalIndex = 0;
    int64_t count = 0;

    makeExecutionPlanFinalized();
    ASSERT_NO_THROW(plan->getAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_GLOBAL_INDEX_EXT,
                                       HIPDNN_TYPE_INT64,
                                       1,
                                       &count,
                                       &engineGlobalIndex));
    ASSERT_EQ(count, 1);
    ASSERT_EQ(engineGlobalIndex, ENGINE_ID);
}

TEST_F(TestExecutionPlanDescriptor, GetTensorUids)
{
    auto plan = getExecutionPlanDescriptor();

    makeExecutionPlanFinalized();

    int64_t count = 0;
    ASSERT_NO_THROW(plan->getAttribute(
        HIPDNN_ATTR_EXECUTION_PLAN_TENSOR_UIDS_EXT, HIPDNN_TYPE_INT64, 0, &count, nullptr));
    ASSERT_EQ(count, static_cast<int64_t>(_tensorUids.size()));

    std::vector<int64_t> tensorUids(static_cast<size_t>(count));
    ASSERT_NO_THROW(plan->getAttribute(HIPDNN_ATTR_EXECUTION_PLAN_TENSOR_UIDS_EXT,
                                       HIPDNN_TYPE_INT64,
                                       count,
                                       &count,
                                       tensorUids.data()));
    ASSERT_EQ(tensorUids, _tensorUids);
    ASSERT_EQ(plan->getTensorUids(), _tensorUids);
}

TEST_F(TestExecutionPlanDescriptor, FinalizeCollectsTensorAlignments)
{
    auto plan = getExecutionPlanDescriptor();

    ASSERT_THROW_HIPDNN_STATUS(plan->getTensorAlignments(), HIPDNN_STATUS_INTERNAL_ERROR);

    makeExecutionPlanFinalized();

    // Alignments are collected from the serialized graph, parallel to the uids.
    ASSERT_EQ(plan->getTensorAlignments(), _tensorAlignments);
    ASSERT_EQ(plan->getTensorAlignments().size(), plan->getTensorUids().size());
}

TEST_F(TestExecutionPlanDescriptor, DeserializeRejectsInvalidFlatBuffer)
{
    auto plan = getExecutionPlanDescriptor();
    std::array<uint8_t, 8> invalidPlan{0, 1, 2, 3, 4, 5, 6, 7};

    ASSERT_THROW_HIPDNN_STATUS(plan->deserializeBackendPlan(_mockEnginePluginResourceManager,
                                                            invalidPlan.data(),
                                                            invalidPlan.size()),
                               HIPDNN_STATUS_BAD_PARAM);
}

TEST_F(TestExecutionPlanDescriptor, DeserializeRejectsUnsupportedVersion)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan(99);

    ASSERT_THROW_HIPDNN_STATUS(plan->deserializeBackendPlan(_mockEnginePluginResourceManager,
                                                            serializedPlan.data(),
                                                            serializedPlan.size()),
                               HIPDNN_STATUS_NOT_SUPPORTED);
}

TEST_F(TestExecutionPlanDescriptor, DeserializeRejectsMissingTensorUids)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan(2, 1024, false);

    ASSERT_THROW_HIPDNN_STATUS(plan->deserializeBackendPlan(_mockEnginePluginResourceManager,
                                                            serializedPlan.data(),
                                                            serializedPlan.size()),
                               HIPDNN_STATUS_BAD_PARAM);
}

TEST_F(TestExecutionPlanDescriptor, DeserializeRejectsEmptyTensorUids)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan(2, 1024, true, true, true);

    ASSERT_THROW_HIPDNN_STATUS(plan->deserializeBackendPlan(_mockEnginePluginResourceManager,
                                                            serializedPlan.data(),
                                                            serializedPlan.size()),
                               HIPDNN_STATUS_BAD_PARAM);
}

TEST_F(TestExecutionPlanDescriptor, DeserializeRejectsMissingPluginPayload)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan(2, 1024, true, false);

    ASSERT_THROW_HIPDNN_STATUS(plan->deserializeBackendPlan(_mockEnginePluginResourceManager,
                                                            serializedPlan.data(),
                                                            serializedPlan.size()),
                               HIPDNN_STATUS_BAD_PARAM);
}

TEST_F(TestExecutionPlanDescriptor, DeserializeRejectsEmptyPluginPayload)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan(2, 1024, true, true, false, true);

    ASSERT_THROW_HIPDNN_STATUS(plan->deserializeBackendPlan(_mockEnginePluginResourceManager,
                                                            serializedPlan.data(),
                                                            serializedPlan.size()),
                               HIPDNN_STATUS_BAD_PARAM);
}

TEST_F(TestExecutionPlanDescriptor, DeserializeRejectsNegativeWorkspaceSize)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan(2, -1);

    ASSERT_THROW_HIPDNN_STATUS(plan->deserializeBackendPlan(_mockEnginePluginResourceManager,
                                                            serializedPlan.data(),
                                                            serializedPlan.size()),
                               HIPDNN_STATUS_BAD_PARAM);
}

TEST_F(TestExecutionPlanDescriptor, DeserializeRejectsMissingTensorAlignments)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan(2,
                                             1024,
                                             /*includeTensorUids=*/true,
                                             /*includePluginPayload=*/true,
                                             /*emptyTensorUids=*/false,
                                             /*emptyPluginPayload=*/false,
                                             /*isOverrideShapeEnabled=*/false,
                                             /*includeTensorAlignments=*/false);

    ASSERT_THROW_HIPDNN_STATUS(plan->deserializeBackendPlan(_mockEnginePluginResourceManager,
                                                            serializedPlan.data(),
                                                            serializedPlan.size()),
                               HIPDNN_STATUS_BAD_PARAM);
}

TEST_F(TestExecutionPlanDescriptor, DeserializeRejectsMismatchedTensorAlignments)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan(2,
                                             1024,
                                             /*includeTensorUids=*/true,
                                             /*includePluginPayload=*/true,
                                             /*emptyTensorUids=*/false,
                                             /*emptyPluginPayload=*/false,
                                             /*isOverrideShapeEnabled=*/false,
                                             /*includeTensorAlignments=*/true,
                                             /*mismatchTensorAlignments=*/true);

    ASSERT_THROW_HIPDNN_STATUS(plan->deserializeBackendPlan(_mockEnginePluginResourceManager,
                                                            serializedPlan.data(),
                                                            serializedPlan.size()),
                               HIPDNN_STATUS_BAD_PARAM);
}

TEST_F(TestExecutionPlanDescriptor, DeserializeRestoresSerializedExecutionPlan)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan();

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                createExecutionContextFromSerialized(ENGINE_ID, _))
        .WillOnce([](int64_t, const hipdnnPluginConstData_t* serializedContext) {
            EXPECT_EQ(serializedContext->size, 3);
            return getExecutionContext();
        });
    EXPECT_CALL(*_mockEnginePluginResourceManager, destroyExecutionContext(_, _));

    ASSERT_NO_THROW(plan->deserializeBackendPlan(
        _mockEnginePluginResourceManager, serializedPlan.data(), serializedPlan.size()));

    ASSERT_TRUE(plan->isFinalized());
    ASSERT_EQ(plan->getTensorUids(), _tensorUids);
    ASSERT_EQ(plan->getTensorAlignments(), _tensorAlignments);
    ASSERT_FALSE(plan->isOverrideShapeEnabled());
    ASSERT_EQ(plan->getExecutionContext(), getExecutionContext());

    int64_t workspaceSize = 0;
    ASSERT_NO_THROW(plan->getAttribute(
        HIPDNN_ATTR_EXECUTION_PLAN_WORKSPACE_SIZE, HIPDNN_TYPE_INT64, 1, nullptr, &workspaceSize));
    ASSERT_EQ(workspaceSize, 1024);
}

TEST_F(TestExecutionPlanDescriptor, GetEngineGlobalIndexExtReturnsDeserializedEngineId)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan();

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                createExecutionContextFromSerialized(ENGINE_ID, _))
        .WillOnce([](int64_t, const hipdnnPluginConstData_t*) { return getExecutionContext(); });
    EXPECT_CALL(*_mockEnginePluginResourceManager, destroyExecutionContext(_, _));

    ASSERT_NO_THROW(plan->deserializeBackendPlan(
        _mockEnginePluginResourceManager, serializedPlan.data(), serializedPlan.size()));

    int64_t engineGlobalIndex = 0;
    int64_t count = 0;
    ASSERT_NO_THROW(plan->getAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_GLOBAL_INDEX_EXT,
                                       HIPDNN_TYPE_INT64,
                                       1,
                                       &count,
                                       &engineGlobalIndex));
    ASSERT_EQ(count, 1);
    ASSERT_EQ(engineGlobalIndex, ENGINE_ID);
}

TEST_F(TestExecutionPlanDescriptor, GetIsOverrideShapeEnabledExtReturnsDeserializedFlagFalse)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan(1, 1024, true, true, false, false, false);

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                createExecutionContextFromSerialized(ENGINE_ID, _))
        .WillOnce([](int64_t, const hipdnnPluginConstData_t*) { return getExecutionContext(); });
    EXPECT_CALL(*_mockEnginePluginResourceManager, destroyExecutionContext(_, _));

    ASSERT_NO_THROW(plan->deserializeBackendPlan(
        _mockEnginePluginResourceManager, serializedPlan.data(), serializedPlan.size()));

    bool isOverrideShapeEnabled = true;
    int64_t count = 0;
    ASSERT_NO_THROW(plan->getAttribute(HIPDNN_ATTR_EXECUTION_PLAN_IS_OVERRIDE_SHAPE_ENABLED_EXT,
                                       HIPDNN_TYPE_BOOLEAN,
                                       1,
                                       &count,
                                       &isOverrideShapeEnabled));
    ASSERT_EQ(count, 1);
    ASSERT_FALSE(isOverrideShapeEnabled);
}

TEST_F(TestExecutionPlanDescriptor, GetIsOverrideShapeEnabledExtReturnsDeserializedFlagTrue)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan(1, 1024, true, true, false, false, true);

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                createExecutionContextFromSerialized(ENGINE_ID, _))
        .WillOnce([](int64_t, const hipdnnPluginConstData_t*) { return getExecutionContext(); });
    EXPECT_CALL(*_mockEnginePluginResourceManager, destroyExecutionContext(_, _));

    ASSERT_NO_THROW(plan->deserializeBackendPlan(
        _mockEnginePluginResourceManager, serializedPlan.data(), serializedPlan.size()));

    bool isOverrideShapeEnabled = false;
    int64_t count = 0;
    ASSERT_NO_THROW(plan->getAttribute(HIPDNN_ATTR_EXECUTION_PLAN_IS_OVERRIDE_SHAPE_ENABLED_EXT,
                                       HIPDNN_TYPE_BOOLEAN,
                                       1,
                                       &count,
                                       &isOverrideShapeEnabled));
    ASSERT_EQ(count, 1);
    ASSERT_TRUE(isOverrideShapeEnabled);
}

TEST_F(TestExecutionPlanDescriptor, SerializeRejectsUnfinalizedPlan)
{
    auto plan = getExecutionPlanDescriptor();
    size_t planByteSize = 0;

    ASSERT_THROW_HIPDNN_STATUS(plan->serializeBackendPlan(0, &planByteSize, nullptr),
                               HIPDNN_STATUS_NOT_INITIALIZED);
}

TEST_F(TestExecutionPlanDescriptor, SerializeRejectsNullSizePointer)
{
    auto plan = getExecutionPlanDescriptor();
    makeExecutionPlanFinalized();

    ASSERT_THROW_HIPDNN_STATUS(plan->serializeBackendPlan(0, nullptr, nullptr),
                               HIPDNN_STATUS_BAD_PARAM_NULL_POINTER);
}

TEST_F(TestExecutionPlanDescriptor, SerializeRoundTripsFlatBufferEnvelope)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan();
    const std::vector<uint8_t> pluginPayload{9, 8, 7, 6};

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                createExecutionContextFromSerialized(ENGINE_ID, _))
        .WillOnce(Return(getExecutionContext()));
    EXPECT_CALL(*_mockEnginePluginResourceManager, destroyExecutionContext(_, _));
    ASSERT_NO_THROW(plan->deserializeBackendPlan(
        _mockEnginePluginResourceManager, serializedPlan.data(), serializedPlan.size()));

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                serializeExecutionContext(ENGINE_ID, getExecutionContext(), _))
        .Times(1)
        .WillOnce([&pluginPayload](int64_t,
                                   hipdnnEnginePluginExecutionContext_t,
                                   std::vector<uint8_t>& serializedContext) {
            serializedContext = pluginPayload;
        });

    size_t planByteSize = 0;
    ASSERT_NO_THROW(plan->serializeBackendPlan(0, &planByteSize, nullptr));
    ASSERT_GT(planByteSize, 0);

    std::vector<uint8_t> serializedOutput(planByteSize);
    ASSERT_NO_THROW(
        plan->serializeBackendPlan(planByteSize, &planByteSize, serializedOutput.data()));
    serializedOutput.resize(planByteSize);

    flatbuffers::Verifier verifier(serializedOutput.data(), serializedOutput.size());
    ASSERT_TRUE(
        hipdnn_flatbuffers_sdk::data_objects::VerifySerializedExecutionPlanBuffer(verifier));

    auto executionPlan
        = hipdnn_flatbuffers_sdk::data_objects::GetSerializedExecutionPlan(serializedOutput.data());
    ASSERT_EQ(executionPlan->version(), 2);
    ASSERT_EQ(executionPlan->engine_id(), ENGINE_ID);
    ASSERT_EQ(executionPlan->workspace_size(), 1024);
    ASSERT_FALSE(executionPlan->is_override_shape_enabled());
    ASSERT_EQ(std::vector<int64_t>(executionPlan->tensor_uids()->begin(),
                                   executionPlan->tensor_uids()->end()),
              _tensorUids);
    ASSERT_NE(executionPlan->tensor_alignments(), nullptr);
    ASSERT_EQ(std::vector<int64_t>(executionPlan->tensor_alignments()->begin(),
                                   executionPlan->tensor_alignments()->end()),
              _tensorAlignments);
    ASSERT_EQ(std::vector<uint8_t>(executionPlan->plugin_payload()->begin(),
                                   executionPlan->plugin_payload()->end()),
              pluginPayload);
}

TEST_F(TestExecutionPlanDescriptor, SerializeCachesPlanAndSerializesContextOnce)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan();
    const std::vector<uint8_t> pluginPayload{9, 8, 7, 6};

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                createExecutionContextFromSerialized(ENGINE_ID, _))
        .WillOnce(Return(getExecutionContext()));
    EXPECT_CALL(*_mockEnginePluginResourceManager, destroyExecutionContext(_, _));
    ASSERT_NO_THROW(plan->deserializeBackendPlan(
        _mockEnginePluginResourceManager, serializedPlan.data(), serializedPlan.size()));

    // The descriptor caches the serialized blob after the first build, so the
    // plugin context is serialized exactly once across all size/fill calls.
    EXPECT_CALL(*_mockEnginePluginResourceManager,
                serializeExecutionContext(ENGINE_ID, getExecutionContext(), _))
        .Times(1)
        .WillOnce([&pluginPayload](int64_t,
                                   hipdnnEnginePluginExecutionContext_t,
                                   std::vector<uint8_t>& serializedContext) {
            serializedContext = pluginPayload;
        });

    size_t firstSize = 0;
    ASSERT_NO_THROW(plan->serializeBackendPlan(0, &firstSize, nullptr));
    ASSERT_GT(firstSize, 0);
    std::vector<uint8_t> firstOutput(firstSize);
    ASSERT_NO_THROW(plan->serializeBackendPlan(firstSize, &firstSize, firstOutput.data()));

    // A second size/fill round-trip must reuse the cache and produce
    // byte-identical output.
    size_t secondSize = 0;
    ASSERT_NO_THROW(plan->serializeBackendPlan(0, &secondSize, nullptr));
    ASSERT_EQ(secondSize, firstSize);
    std::vector<uint8_t> secondOutput(secondSize);
    ASSERT_NO_THROW(plan->serializeBackendPlan(secondSize, &secondSize, secondOutput.data()));

    EXPECT_EQ(firstOutput, secondOutput);
}

TEST_F(TestExecutionPlanDescriptor, SerializeRoundTripsOverrideShapeEnabledFlag)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan
        = makeSerializedPlan(2, 1024, true, true, false, false, /*isOverrideShapeEnabled=*/true);
    const std::vector<uint8_t> pluginPayload{9, 8, 7, 6};

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                createExecutionContextFromSerialized(ENGINE_ID, _))
        .WillOnce(Return(getExecutionContext()));
    EXPECT_CALL(*_mockEnginePluginResourceManager, destroyExecutionContext(_, _));
    ASSERT_NO_THROW(plan->deserializeBackendPlan(
        _mockEnginePluginResourceManager, serializedPlan.data(), serializedPlan.size()));
    ASSERT_TRUE(plan->isOverrideShapeEnabled());

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                serializeExecutionContext(ENGINE_ID, getExecutionContext(), _))
        .Times(1)
        .WillOnce([&pluginPayload](int64_t,
                                   hipdnnEnginePluginExecutionContext_t,
                                   std::vector<uint8_t>& serializedContext) {
            serializedContext = pluginPayload;
        });

    size_t planByteSize = 0;
    ASSERT_NO_THROW(plan->serializeBackendPlan(0, &planByteSize, nullptr));

    std::vector<uint8_t> serializedOutput(planByteSize);
    ASSERT_NO_THROW(
        plan->serializeBackendPlan(planByteSize, &planByteSize, serializedOutput.data()));

    auto executionPlan
        = hipdnn_flatbuffers_sdk::data_objects::GetSerializedExecutionPlan(serializedOutput.data());
    EXPECT_TRUE(executionPlan->is_override_shape_enabled());
}

TEST_F(TestExecutionPlanDescriptor, SerializeRejectsInsufficientBuffer)
{
    auto plan = getExecutionPlanDescriptor();
    auto serializedPlan = makeSerializedPlan();
    const std::vector<uint8_t> pluginPayload{9, 8, 7, 6};

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                createExecutionContextFromSerialized(ENGINE_ID, _))
        .WillOnce(Return(getExecutionContext()));
    EXPECT_CALL(*_mockEnginePluginResourceManager, destroyExecutionContext(_, _));
    ASSERT_NO_THROW(plan->deserializeBackendPlan(
        _mockEnginePluginResourceManager, serializedPlan.data(), serializedPlan.size()));

    EXPECT_CALL(*_mockEnginePluginResourceManager,
                serializeExecutionContext(ENGINE_ID, getExecutionContext(), _))
        .WillOnce([&pluginPayload](int64_t,
                                   hipdnnEnginePluginExecutionContext_t,
                                   std::vector<uint8_t>& serializedContext) {
            serializedContext = pluginPayload;
        });

    size_t planByteSize = 0;
    std::array<uint8_t, 1> serializedOutput{0};
    ASSERT_THROW_HIPDNN_STATUS(
        plan->serializeBackendPlan(serializedOutput.size(), &planByteSize, serializedOutput.data()),
        HIPDNN_STATUS_BAD_PARAM_SIZE_INSUFFICIENT);
}

TEST_F(TestExecutionPlanDescriptor, GetEngineConfig)
{
    auto plan = getExecutionPlanDescriptor();

    ScopedDescriptor returnedEngineConfig;
    ScopedDescriptor nullCountEngineConfig;
    int64_t count = 0;

    makeExecutionPlanFinalized();

    ASSERT_NO_THROW(plan->getAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                                       HIPDNN_TYPE_BACKEND_DESCRIPTOR,
                                       1,
                                       &count,
                                       static_cast<void*>(returnedEngineConfig.getPtr())));

    ASSERT_EQ(count, 1);
    ASSERT_EQ(*returnedEngineConfig.get(), *(_mockEngineConfigWrapper.get()));

    ASSERT_NO_THROW(plan->getAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                                       HIPDNN_TYPE_BACKEND_DESCRIPTOR,
                                       1,
                                       nullptr,
                                       static_cast<void*>(nullCountEngineConfig.getPtr())));

    ASSERT_EQ(*nullCountEngineConfig.get(), *(_mockEngineConfigWrapper.get()));
}

TEST_F(TestExecutionPlanDescriptor, GetEngineConfigErrors)
{
    auto plan = getExecutionPlanDescriptor();
    hipdnnBackendDescriptor_t returnedEngineConfig = nullptr;
    int64_t count = 0;
    void* dummy = static_cast<void*>(&returnedEngineConfig);

    makeExecutionPlanFinalized();

    ASSERT_THROW_HIPDNN_STATUS(
        plan->getAttribute(
            HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG, HIPDNN_TYPE_INT64, 1, &count, &dummy),
        HIPDNN_STATUS_BAD_PARAM);

    ASSERT_THROW_HIPDNN_STATUS(plan->getAttribute(HIPDNN_ATTR_EXECUTION_PLAN_ENGINE_CONFIG,
                                                  HIPDNN_TYPE_BACKEND_DESCRIPTOR,
                                                  0, // Too small (needs to be at least 1)
                                                  &count,
                                                  &dummy),
                               HIPDNN_STATUS_BAD_PARAM);
}

TEST_F(TestExecutionPlanDescriptor, GetUnsupportedAttr)
{
    auto plan = getExecutionPlanDescriptor();
    int64_t count = 0;
    std::array<char, 256> dummyBuffer{};

    makeExecutionPlanFinalized();

    ASSERT_THROW_HIPDNN_STATUS(plan->getAttribute(HIPDNN_ATTR_EXECUTION_PLAN_KERNEL_CACHE,
                                                  HIPDNN_TYPE_INT64,
                                                  1,
                                                  &count,
                                                  dummyBuffer.data()),
                               HIPDNN_STATUS_NOT_SUPPORTED);
}

TEST_F(TestExecutionPlanDescriptor, GetEngineConfigThrowsIfNotFinalized)
{
    auto plan = getExecutionPlanDescriptor();
    ASSERT_THROW_HIPDNN_STATUS(plan->getEngineConfig(), HIPDNN_STATUS_INTERNAL_ERROR);
}

TEST_F(TestExecutionPlanDescriptor, GetEngineConfigReturnsPointerIfFinalized)
{
    auto plan = getExecutionPlanDescriptor();
    makeExecutionPlanFinalized();
    auto engineConfigPtr = plan->getEngineConfig();
    ASSERT_NE(engineConfigPtr, nullptr);
    ASSERT_EQ(static_cast<const IBackendDescriptor*>(engineConfigPtr.get()),
              static_cast<const IBackendDescriptor*>(getMockEngineConfig().get()));
}

TEST_F(TestExecutionPlanDescriptor, GetExecutionContext)
{
    auto plan = getExecutionPlanDescriptor();
    makeExecutionPlanFinalized();
    auto context = plan->getExecutionContext();
    ASSERT_NE(context, nullptr);
    ASSERT_EQ(context, getExecutionContext());
}