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

// Unit tests for IntegrationBundleVerificationHarness's core
// contract: how it translates an executor's behaviour into a GTest outcome.
//
//   executor throws (unsupported graph) -> SKIP
//   executor writes matching output     -> PASS (no failure recorded)
//   executor writes mismatching output  -> FAIL
//
// The harness reports these via GTest itself (GTEST_SKIP / EXPECT_TRUE), so we
// drive it under a ScopedFakeTestPartResultReporter and inspect the captured
// TestPartResults rather than letting them affect *this* test.

#include <gtest/gtest-spi.h>
#include <gtest/gtest.h>

#include <algorithm>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <functional>
#include <memory>
#include <optional>
#include <variant>
#include <vector>

#include <hipdnn_test_sdk/utilities/FileUtilities.hpp>
#include <hipdnn_test_sdk/utilities/FlatbufferGraphTestUtils.hpp>
#include <hipdnn_test_sdk/utilities/TestUtilities.hpp>

#include "harness/EngineNotApplicableError.hpp"
#include "harness/bundle/IntegrationBundleVerificationHarness.hpp"
#include "harness/bundle/IntegrationTestBundle.hpp"

// NOLINTBEGIN(readability-identifier-naming)

using namespace hipdnn_integration_tests::bundle;

namespace
{

// Exposes the harness's protected SetUp/TestBody so a test can drive the full
// lifecycle directly, and overrides executeGraphThroughEngine with a stub so the
// tests run on CPU-only CI without a real GPU engine.
class TestableHarness : public IntegrationBundleVerificationHarness
{
public:
    using StubFunc = std::function<void(std::unordered_map<int64_t, void*>&)>;

    explicit TestableHarness(StubFunc engineStub,
                             bool requiresDevice = false,
                             hipdnn_integration_tests::VerificationMode mode
                             = hipdnn_integration_tests::VerificationMode::AUTO)
        : IntegrationBundleVerificationHarness(requiresDevice)
        , _engineStub(std::move(engineStub))
        , _mode(mode)
    {
    }

    void SetUp() override {}

    using IntegrationBundleVerificationHarness::synthesis;
    using IntegrationBundleVerificationHarness::TestBody;

protected:
    void executeGraphThroughEngine(std::unordered_map<int64_t, void*>& variantPack) override
    {
        _engineStub(variantPack);
    }

    hipdnn_integration_tests::VerificationMode getVerificationMode() const override
    {
        return _mode;
    }

    void applyMetadataGuards() const override {}

private:
    StubFunc _engineStub;
    hipdnn_integration_tests::VerificationMode _mode;
};

class TestGoldenHarnessFixture : public ::testing::Test
{
protected:
    std::optional<hipdnn_test_sdk::utilities::ScopedDirectory> _scopedDir;
    std::filesystem::path _tempDir;

    void SetUp() override
    {
        auto path
            = std::filesystem::temp_directory_path()
              / ("golden_harness_test_"
                 + std::to_string(::testing::UnitTest::GetInstance()->current_test_info()->line()));
        std::filesystem::remove_all(path);
        _scopedDir.emplace(path);
        _tempDir = _scopedDir->path();
    }

    // The float value every output element is seeded with on disk, so "matching"
    // vs "mismatching" executor output is unambiguous (not the all-zero buffer
    // that the harness already zeroes outputs to before running).
    static constexpr float K_OUTPUT_VALUE = 3.5f;

    // uid 5 (the batchnorm y output): dims [2,3,4,5] -> 120 floats.
    static constexpr int64_t K_OUTPUT_UID = 5;
    static constexpr size_t K_OUTPUT_ELEMS = 120;

    // Writes a schema-valid nchw/fp32 batchnorm-inference graph, its metadata,
    // and all tensor .bin blobs. Inputs are zero-filled; the output blob (uid 5)
    // is filled with K_OUTPUT_VALUE so it is a distinct, non-zero golden.
    static void writeBundleFiles(const std::filesystem::path& dir, const std::string& name)
    {
        std::filesystem::create_directories(dir);
        std::ofstream(dir / (name + ".json"))
            << R"({"nodes": [{"inputs": {"x_tensor_uid": 0, "mean_tensor_uid": 1, )"
               R"("inv_variance_tensor_uid": 2, "scale_tensor_uid": 3, "bias_tensor_uid": 4}, )"
               R"("outputs": {"y_tensor_uid": 5}, "type": "BatchnormInferenceAttributes", )"
               R"("compute_data_type": "float", "name": ""}], "tensors": [)"
               R"({"name": "", "uid": 0, "strides": [60, 20, 5, 1], "dims": [2, 3, 4, 5], )"
               R"("data_type": "float", "virtual": false}, )"
               R"({"name": "", "uid": 1, "strides": [3, 1, 1, 1], "dims": [1, 3, 1, 1], )"
               R"("data_type": "float", "virtual": false}, )"
               R"({"name": "", "uid": 2, "strides": [3, 1, 1, 1], "dims": [1, 3, 1, 1], )"
               R"("data_type": "float", "virtual": false}, )"
               R"({"name": "", "uid": 3, "strides": [3, 1, 1, 1], "dims": [1, 3, 1, 1], )"
               R"("data_type": "float", "virtual": false}, )"
               R"({"name": "", "uid": 4, "strides": [3, 1, 1, 1], "dims": [1, 3, 1, 1], )"
               R"("data_type": "float", "virtual": false}, )"
               R"({"name": "", "uid": 5, "strides": [60, 20, 5, 1], "dims": [2, 3, 4, 5], )"
               R"("data_type": "float", "virtual": false}], "io_data_type": "float", )"
               R"("compute_data_type": "float", "intermediate_data_type": "float", "name": ""})";

        std::ofstream(dir / (name + ".meta.json"))
            << R"({"format_version": 1, "operation": "BatchnormInference"})";

        const auto basePath = (dir / name).string();
        const auto writeFloatBin = [&](int64_t uid, size_t elems, float value) {
            const std::vector<float> data(elems, value);
            std::ofstream out(basePath + ".tensor" + std::to_string(uid) + ".bin",
                              std::ios::binary);
            out.write(reinterpret_cast<const char*>(data.data()),
                      static_cast<std::streamsize>(data.size() * sizeof(float)));
        };

        writeFloatBin(0, 120, 0.0f); // x
        writeFloatBin(1, 3, 0.0f); // mean
        writeFloatBin(2, 3, 0.0f); // inv_variance
        writeFloatBin(3, 3, 0.0f); // scale
        writeFloatBin(4, 3, 0.0f); // bias
        writeFloatBin(K_OUTPUT_UID, K_OUTPUT_ELEMS, K_OUTPUT_VALUE); // y (golden)
    }

    // Loads a fully-populated bundle (graph + tensors + metadata) ready to run.
    std::shared_ptr<IntegrationTestBundle> loadRunnableBundle(const std::string& name) const
    {
        const auto dir = _tempDir / name;
        writeBundleFiles(dir, name);
        auto result = loadIntegrationTestBundle(dir / (name + ".json"));
        EXPECT_TRUE(std::holds_alternative<IntegrationTestBundle>(result));
        return std::make_shared<IntegrationTestBundle>(
            std::move(std::get<IntegrationTestBundle>(result)));
    }

    // Fills the output buffer in the variant pack with `value`. A stand-in for
    // what a real executor would compute into the output buffer.
    static void writeOutput(std::unordered_map<int64_t, void*>& variantPack, float value)
    {
        auto* ptr = static_cast<float*>(variantPack.at(K_OUTPUT_UID));
        std::fill(ptr, ptr + K_OUTPUT_ELEMS, value);
    }

    // Drives the harness's full lifecycle (SetUp then TestBody) for the given
    // bundle + stub executor, capturing every TestPartResult the harness records.
    static void runCapturing(std::shared_ptr<IntegrationTestBundle> bundle,
                             TestableHarness::StubFunc stub,
                             ::testing::TestPartResultArray* results)
    {
        TestableHarness harness(std::move(stub));
        harness.setBundle(std::move(bundle), "unit-test-bundle");

        const ::testing::ScopedFakeTestPartResultReporter reporter(
            ::testing::ScopedFakeTestPartResultReporter::INTERCEPT_ALL_THREADS, results);
        harness.SetUp();
        harness.TestBody();
    }

    static bool anySkipped(const ::testing::TestPartResultArray& results)
    {
        for(int i = 0; i < results.size(); ++i)
        {
            if(results.GetTestPartResult(i).skipped())
            {
                return true;
            }
        }
        return false;
    }

    static bool anyFailed(const ::testing::TestPartResultArray& results)
    {
        for(int i = 0; i < results.size(); ++i)
        {
            if(results.GetTestPartResult(i).failed())
            {
                return true;
            }
        }
        return false;
    }
};

// uids: x=1, y=2, scale=3, bias=4, epsilon=5, prev_mean=8, prev_variance=9, momentum=10
std::shared_ptr<IntegrationTestBundle> makeRuntimePbvSynthesisBundle()
{
    auto builder = hipdnn_test_sdk::utilities::createValidBatchnormFwdTrainingGraph(
        {3, 1},
        {2, 3},
        /*withMeanVariance=*/false,
        /*overrideShapeEnabled=*/false,
        /*runtimeEpsilon=*/true,
        /*withRunningStatsAndMomentum=*/true,
        /*runtimeMomentum=*/true);
    auto bundle = std::make_shared<IntegrationTestBundle>();
    bundle->graphBuffer = builder.Release();
    bundle->outputTensorUids = {2};
    return bundle;
}

std::shared_ptr<IntegrationTestBundle> makeRuntimePassByValueBundle()
{
    using namespace hipdnn_flatbuffers_sdk::data_objects;

    flatbuffers::FlatBufferBuilder builder;
    const std::vector<int64_t> scalarDims = {1};
    const std::vector<int64_t> scalarStrides = {1};
    const std::vector<int64_t> outputDims = {1};
    const std::vector<int64_t> outputStrides = {1};

    std::vector<flatbuffers::Offset<TensorAttributes>> tensors;
    tensors.push_back(CreateTensorAttributesDirect(builder,
                                                   1,
                                                   "epsilon",
                                                   DataType::FLOAT,
                                                   &scalarStrides,
                                                   &scalarDims,
                                                   false,
                                                   TensorValue::NONE,
                                                   0,
                                                   true));
    // Ordinary (non-runtime-PBV) input, to prove buildVariantPack still
    // routes normal tensors to device memory alongside a PBV sibling.
    tensors.push_back(CreateTensorAttributesDirect(
        builder, 3, "scale", DataType::FLOAT, &scalarStrides, &scalarDims));
    tensors.push_back(CreateTensorAttributesDirect(
        builder, 2, "output", DataType::FLOAT, &outputStrides, &outputDims));

    const std::vector<flatbuffers::Offset<Node>> nodes;
    const auto graph = CreateGraphDirect(builder,
                                         "runtime_pbv",
                                         DataType::FLOAT,
                                         DataType::FLOAT,
                                         DataType::FLOAT,
                                         &tensors,
                                         &nodes);
    builder.Finish(graph);

    auto bundle = std::make_shared<IntegrationTestBundle>();
    bundle->graphBuffer = builder.Release();
    bundle->outputTensorUids = {2};
    return bundle;
}

TEST(TestBundleVerificationHarness, DeviceVariantPackUsesHostPointerForRuntimePassByValue)
{
    SKIP_IF_NO_DEVICES();
    auto bundle = makeRuntimePassByValueBundle();
    const auto wrapper = bundle->graphWrapper();
    const auto& tensorAttributes = wrapper.getTensorMap();

    TensorMap inputs;
    inputs.emplace(1, hipdnn_test_sdk::detail::createTensorFromAttribute(*tensorAttributes.at(1)));
    inputs.at(1)->fillTensorWithValue(0.01f);
    auto* expectedHostPointer = inputs.at(1)->rawHostData();

    // Ordinary (non-PBV) input: must still route to device even though a PBV
    // sibling is present in the same variant pack.
    inputs.emplace(3, hipdnn_test_sdk::detail::createTensorFromAttribute(*tensorAttributes.at(3)));

    // Uid with no entry in tensorAttributes at all: buildVariantPack must
    // default isRuntimePassByValue to false (device pointer) rather than
    // treating an unknown uid as runtime-PBV.
    static constexpr int64_t K_UNKNOWN_UID = 99;
    inputs.emplace(K_UNKNOWN_UID,
                   hipdnn_test_sdk::detail::createTensorFromAttribute(*tensorAttributes.at(3)));

    OutputTensors outputs;
    outputs.emplace(2, hipdnn_test_sdk::detail::createTensorFromAttribute(*tensorAttributes.at(2)));
    auto variantPack = detail::buildVariantPack(
        inputs, outputs, tensorAttributes, bundle->outputTensorUids, /*useDevice=*/true);

    ASSERT_EQ(variantPack.at(1), expectedHostPointer);
    EXPECT_FLOAT_EQ(*static_cast<const float*>(variantPack.at(1)), 0.01f);
    EXPECT_EQ(variantPack.at(3), inputs.at(3)->rawDeviceData());
    EXPECT_EQ(variantPack.at(K_UNKNOWN_UID), inputs.at(K_UNKNOWN_UID)->rawDeviceData());
    EXPECT_EQ(variantPack.at(2), outputs.at(2)->rawDeviceData());
}
} // namespace

// Full harness seam: start with graph metadata only, then prove allocation,
// fixed/random synthesis, and host-pointer delivery all happen before execute.
TEST_F(TestGoldenHarnessFixture, GraphOnlyRuntimePbvValuesAreSynthesizedEndToEnd)
{
    const auto runHarness = [&](float& epsilon, float& momentum) {
        ::testing::TestPartResultArray results;
        const auto execute = [&](std::unordered_map<int64_t, void*>& variantPack) {
            epsilon = *static_cast<const float*>(variantPack.at(5));
            momentum = *static_cast<const float*>(variantPack.at(10));
            EXPECT_GE(momentum, 0.0f);
            EXPECT_LE(momentum, 1.0f);
            throw hipdnn_integration_tests::EngineNotApplicableError(
                "value-capture stub completed");
        };
        TestableHarness harness(execute);

        auto bundle = makeRuntimePbvSynthesisBundle();
        harness.setBundle(std::move(bundle), "runtime-pbv-synthesis");
        harness.synthesis().setGlobalSeed(42);

        const ::testing::ScopedFakeTestPartResultReporter reporter(
            ::testing::ScopedFakeTestPartResultReporter::INTERCEPT_ALL_THREADS, &results);
        harness.TestBody();
        EXPECT_TRUE(anySkipped(results));
        EXPECT_FALSE(anyFailed(results));
    };

    float firstEpsilon = 0.0f;
    float firstMomentum = 0.0f;
    runHarness(firstEpsilon, firstMomentum);
    EXPECT_FLOAT_EQ(firstEpsilon, 1e-5f);

    float secondEpsilon = 0.0f;
    float secondMomentum = 0.0f;
    runHarness(secondEpsilon, secondMomentum);
    EXPECT_FLOAT_EQ(secondEpsilon, 1e-5f);
    EXPECT_FLOAT_EQ(secondMomentum, firstMomentum);
}

// An executor that throws ("unsupported graph") must yield a SKIP, not a FAIL —
// the harness translates the throw into GTEST_SKIP.
TEST_F(TestGoldenHarnessFixture, ExecutorThrowsYieldsSkip)
{
    ::testing::TestPartResultArray results;
    runCapturing(
        loadRunnableBundle("throws"),
        [](std::unordered_map<int64_t, void*>&) {
            throw hipdnn_integration_tests::EngineNotApplicableError(
                "engine does not support this graph");
        },
        &results);

    EXPECT_TRUE(anySkipped(results));
    EXPECT_FALSE(anyFailed(results));
}

// An executor that writes output matching the golden reference must yield a
// PASS: no failure and no skip recorded.
TEST_F(TestGoldenHarnessFixture, MatchingOutputYieldsPass)
{
    ::testing::TestPartResultArray results;
    runCapturing(
        loadRunnableBundle("match"),
        [](std::unordered_map<int64_t, void*>& vp) { writeOutput(vp, K_OUTPUT_VALUE); },
        &results);

    EXPECT_FALSE(anyFailed(results));
    EXPECT_FALSE(anySkipped(results));
}

// An executor that writes output differing from the golden reference must yield
// a FAIL.
TEST_F(TestGoldenHarnessFixture, MismatchingOutputYieldsFail)
{
    ::testing::TestPartResultArray results;
    runCapturing(
        loadRunnableBundle("mismatch"),
        [](std::unordered_map<int64_t, void*>& vp) { writeOutput(vp, K_OUTPUT_VALUE + 100.0f); },
        &results);

    EXPECT_TRUE(anyFailed(results));
    EXPECT_FALSE(anySkipped(results));
}

// NOLINTEND(readability-identifier-naming)