/* Copyright (c) Advanced Micro Devices, Inc. All rights reserved.
 *
 * SPDX-License-Identifier: MIT
 */

#include "context.h"
#include "configuration.h"
#include "hipfile-warnings.h"
#include "hipfile.h"

#include "ais-capability.h"
#include "test-common.h"
#include "test-options.h"

#include <array>
#include <cstdlib>
#include <gtest/gtest.h>
#include <hip/hip_runtime_api.h>
#include <string>
#include <thread>
#include <unistd.h>

extern SystemTestOptions test_env;

using namespace hipFile;

namespace {

// Gate fastpath-only tests on AIS capability.
void
enforceFastpathGate(hipFileHandle_t handle, void *device_buffer)
{
    hipFile::test::AisCapability ais_capability{test_env.allow_skip_fastpath};

    const auto decision = ais_capability.populate(handle, device_buffer);

    if (decision == hipFile::test::AisCapability::GateDecision::Run) {
        return;
    }

    if (decision == hipFile::test::AisCapability::GateDecision::Skip) {
        // Keep this marker synchronized with test/CMakeLists.txt SKIP_REGULAR_EXPRESSION.
        GTEST_SKIP() << "fastpath not available in this environment\n" << ais_capability.report();
    }

    FAIL() << "Fastpath Validation Failed!\n" << ais_capability.report() << "\n" << ais_capability.skipHint();
}

}

HIPFILE_WARN_NO_GLOBAL_CTOR_OFF

enum class IoTestBackend {
    Fastpath,
    Fallback,
};

struct IoTestParam {
    IoTestBackend backend;
    std::string   name;
};

HIPFILE_WARN_NO_EXIT_DTOR_OFF
static std::array<IoTestParam, 2> io_test_params{
    {{IoTestBackend::Fastpath, "Fastpath"}, {IoTestBackend::Fallback, "Fallback"}}};
HIPFILE_WARN_NO_EXIT_DTOR_ON

struct HipFileIo : public testing::TestWithParam<IoTestParam> {

    Tmpfile         tmpfile;
    size_t          tmpfile_size;
    hipFileHandle_t tmpfile_handle;
    void           *unregistered_device_buffer;
    size_t          unregistered_device_buffer_size;

    HipFileIo()
        : tmpfile{test_env.ais_capable_dir}, tmpfile_size{1024 * 1024}, tmpfile_handle{nullptr},
          unregistered_device_buffer{nullptr}, unregistered_device_buffer_size{1024 * 1024}
    {
    }

    void SetUp() override
    {
        // Disable all backends
        Context<Configuration>::get()->fastpath(false);
        Context<Configuration>::get()->fallback(false);

        // Enable the desired backend
        switch (GetParam().backend) {
            case IoTestBackend::Fastpath:
                Context<Configuration>::get()->fastpath(true);
                break;

            case IoTestBackend::Fallback:
                Context<Configuration>::get()->fallback(true);
                break;

            default:
                FAIL() << "Unsupported IoTestBackend";
        }

        ASSERT_EQ(0, ftruncate(tmpfile.fd, static_cast<off_t>(tmpfile_size)));

        hipFileDescr_t descr{};
        descr.type      = hipFileHandleTypeOpaqueFD;
        descr.handle.fd = tmpfile.fd;

        ASSERT_EQ(HIPFILE_SUCCESS, hipFileHandleRegister(&tmpfile_handle, &descr));

        ASSERT_EQ(hipSuccess, hipMalloc(&unregistered_device_buffer, unregistered_device_buffer_size));

        if (GetParam().backend == IoTestBackend::Fastpath) {
            enforceFastpathGate(tmpfile_handle, unregistered_device_buffer);
        }
    }

    void TearDown() override
    {
        ASSERT_EQ(hipSuccess, hipFree(unregistered_device_buffer));

        hipFileHandleDeregister(tmpfile_handle);
    }
};

TEST_P(HipFileIo, ReadToUnregisteredBufferAtOffset)
{
    hoff_t io_buffer_offset{4096};
    size_t io_size{unregistered_device_buffer_size - static_cast<size_t>(io_buffer_offset)};

    ASSERT_EQ(io_size, hipFileRead(tmpfile_handle, unregistered_device_buffer, io_size, 0, io_buffer_offset));
}

TEST_P(HipFileIo, ReadToUnregisteredBufferAtOffsetReturnsErrorIfOverflow)
{
    hoff_t io_buffer_offset{4096};
    size_t io_size{unregistered_device_buffer_size};

    ASSERT_EQ(-hipFileInvalidValue,
              hipFileRead(tmpfile_handle, unregistered_device_buffer, io_size, 0, io_buffer_offset));
}

INSTANTIATE_TEST_SUITE_P(, HipFileIo, testing::ValuesIn(io_test_params),
                         [](const testing::TestParamInfo<HipFileIo::ParamType> &param_info) {
                             return param_info.param.name;
                         });

struct HipFileIoHipInit : public testing::Test {

    Tmpfile         tmpfile;
    size_t          tmpfile_size;
    hipFileHandle_t tmpfile_handle;
    void           *registered_device_buffer;
    size_t          registered_device_buffer_size;

    HipFileIoHipInit()
        : tmpfile{test_env.ais_capable_dir}, tmpfile_size{1024 * 1024}, tmpfile_handle{nullptr},
          registered_device_buffer{nullptr}, registered_device_buffer_size{1024 * 1024}
    {
    }

    void SetUp() override
    {
        Context<Configuration>::get()->fastpath(true);
        Context<Configuration>::get()->fallback(false);

        ASSERT_EQ(0, ftruncate(tmpfile.fd, static_cast<off_t>(tmpfile_size)));

        hipFileDescr_t descr{};
        descr.type      = hipFileHandleTypeOpaqueFD;
        descr.handle.fd = tmpfile.fd;

        ASSERT_EQ(HIPFILE_SUCCESS, hipFileHandleRegister(&tmpfile_handle, &descr));
        ASSERT_EQ(hipSuccess, hipMalloc(&registered_device_buffer, registered_device_buffer_size));
        ASSERT_EQ(HIPFILE_SUCCESS,
                  hipFileBufRegister(registered_device_buffer, registered_device_buffer_size, 0));

        enforceFastpathGate(tmpfile_handle, registered_device_buffer);
    }

    void TearDown() override
    {
        ASSERT_EQ(HIPFILE_SUCCESS, hipFileBufDeregister(registered_device_buffer));
        ASSERT_EQ(hipSuccess, hipFree(registered_device_buffer));
        hipFileHandleDeregister(tmpfile_handle);
    }
};

TEST_F(HipFileIoHipInit, spawnedThreadReadRunsWithoutSegfault)
{
    size_t  io_size{registered_device_buffer_size};
    ssize_t res{};
    std::thread([&]() { res = hipFileRead(tmpfile_handle, registered_device_buffer, io_size, 0, 0); }).join();
    ASSERT_EQ(io_size, res);
}

TEST_F(HipFileIoHipInit, spawnedThreadWriteRunsWithoutSegfault)
{
    size_t  io_size{registered_device_buffer_size};
    ssize_t res{};
    std::thread([&]() {
        res = hipFileWrite(tmpfile_handle, registered_device_buffer, io_size, 0, 0);
    }).join();
    ASSERT_EQ(io_size, res);
}

HIPFILE_WARN_NO_GLOBAL_CTOR_ON