/*
 *
 * Copyright (C) 2023 Intel Corporation
 *
 * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
 * See LICENSE.TXT
 * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 *
 * @file codegen.cpp
 *
 * @brief UR code generation and execution example for use with the Level Zero adapter.
 *
 * The codegen example demonstrates a complete flow for generating LLVM IR,
 * translating it to SPIR-V, and submitting the kernel to Level Zero Runtime via UR API.
 */

#include <iostream>
#include <vector>

#include "helpers.h"
#include "ur_api.h"

constexpr unsigned PAGE_SIZE = 4096;

void ur_check(const ur_result_t r) {
    if (r != UR_RESULT_SUCCESS) {
        urLoaderTearDown();
        throw std::runtime_error("Unified runtime error: " + std::to_string(r));
    }
}

std::vector<ur_adapter_handle_t> get_adapters() {
    uint32_t adapterCount = 0;
    ur_check(urAdapterGet(0, nullptr, &adapterCount));

    if (!adapterCount) {
        throw std::runtime_error("No adapters available.");
    }

    std::vector<ur_adapter_handle_t> adapters(adapterCount);
    ur_check(urAdapterGet(adapterCount, adapters.data(), nullptr));
    return adapters;
}

std::vector<ur_adapter_handle_t>
get_supported_adapters(std::vector<ur_adapter_handle_t> &adapters) {
    std::vector<ur_adapter_handle_t> supported_adapters;
    for (auto adapter : adapters) {
        ur_adapter_backend_t backend;
        ur_check(urAdapterGetInfo(adapter, UR_ADAPTER_INFO_BACKEND,
                                  sizeof(ur_adapter_backend_t), &backend,
                                  nullptr));

        if (backend == UR_ADAPTER_BACKEND_LEVEL_ZERO) {
            supported_adapters.push_back(adapter);
        }
    }

    return supported_adapters;
}

std::vector<ur_platform_handle_t>
get_platforms(std::vector<ur_adapter_handle_t> &adapters) {
    uint32_t platformCount = 0;
    ur_check(urPlatformGet(adapters.data(), adapters.size(), 1, nullptr,
                           &platformCount));

    if (!platformCount) {
        throw std::runtime_error("No platforms available.");
    }

    std::vector<ur_platform_handle_t> platforms(platformCount);
    ur_check(urPlatformGet(adapters.data(), adapters.size(), platformCount,
                           platforms.data(), nullptr));
    return platforms;
}

std::vector<ur_device_handle_t> get_gpus(ur_platform_handle_t p) {
    uint32_t deviceCount = 0;
    ur_check(urDeviceGet(p, UR_DEVICE_TYPE_GPU, 0, nullptr, &deviceCount));

    if (!deviceCount) {
        throw std::runtime_error("No GPUs available.");
    }

    std::vector<ur_device_handle_t> devices(deviceCount);
    ur_check(urDeviceGet(p, UR_DEVICE_TYPE_GPU, deviceCount, devices.data(),
                         nullptr));
    return devices;
}

template <typename T, size_t N> struct alignas(PAGE_SIZE) AlignedArray {
    T data[N];
};

int main() {
    ur_loader_config_handle_t loader_config = nullptr;
    ur_check(urLoaderInit(UR_DEVICE_INIT_FLAG_GPU, loader_config));

    auto adapters = get_adapters();
    auto supported_adapters = get_supported_adapters(adapters);
    auto platforms = get_platforms(supported_adapters);
    auto gpus = get_gpus(platforms.front());
    auto spv = generate_plus_one_spv();

    constexpr int a_size = 32;
    AlignedArray<int, a_size> a, b;
    for (auto i = 0; i < a_size; ++i) {
        a.data[i] = a_size - i;
        b.data[i] = i;
    }

    auto current_device = gpus.front();

    ur_context_handle_t hContext;
    ur_check(urContextCreate(1, &current_device, nullptr, &hContext));

    ur_program_handle_t hProgram;
    ur_check(urProgramCreateWithIL(hContext, spv.data(), spv.size(), nullptr,
                                   &hProgram));
    ur_check(urProgramBuild(hContext, hProgram, nullptr));

    ur_mem_handle_t dA, dB;
    ur_check(urMemBufferCreate(hContext, UR_MEM_FLAG_READ_WRITE,
                               a_size * sizeof(int), nullptr, &dA));
    ur_check(urMemBufferCreate(hContext, UR_MEM_FLAG_READ_WRITE,
                               a_size * sizeof(int), nullptr, &dB));

    ur_kernel_handle_t hKernel;
    ur_check(urKernelCreate(hProgram, "plus1", &hKernel));
    ur_check(urKernelSetArgMemObj(hKernel, 0, nullptr, dA));
    ur_check(urKernelSetArgMemObj(hKernel, 1, nullptr, dB));

    ur_queue_handle_t queue;
    ur_check(urQueueCreate(hContext, current_device, nullptr, &queue));

    ur_check(urEnqueueMemBufferWrite(queue, dA, true, 0, a_size * sizeof(int),
                                     a.data, 0, nullptr, nullptr));
    ur_check(urEnqueueMemBufferWrite(queue, dB, true, 0, a_size * sizeof(int),
                                     b.data, 0, nullptr, nullptr));

    const size_t gWorkOffset[] = {0, 0, 0};
    const size_t gWorkSize[] = {128, 1, 1};
    const size_t lWorkSize[] = {1, 1, 1};
    ur_event_handle_t event;
    ur_check(urEnqueueKernelLaunch(queue, hKernel, 3, gWorkOffset, gWorkSize,
                                   lWorkSize, 0, nullptr, &event));

    ur_check(urEnqueueMemBufferRead(queue, dB, true, 0, a_size * sizeof(int),
                                    b.data, 1, &event, nullptr));

    ur_check(urQueueFinish(queue));

    std::cout << "Input Array: ";
    for (int i = 0; i < a_size; ++i) {
        std::cout << a.data[i] << " ";
    }
    std::cout << std::endl;

    bool expectedResult = false;

    std::cout << "Output Array: ";
    for (int i = 0; i < a_size; ++i) {
        std::cout << b.data[i] << " ";
        expectedResult |= (b.data[i] == a.data[i] + 1);
    }
    std::cout << std::endl;

    if (expectedResult) {
        std::cout << "Results are correct." << std::endl;
    } else {
        std::cout << "Results are incorrect." << std::endl;
    }

    return urLoaderTearDown() == UR_RESULT_SUCCESS && expectedResult ? 0 : 1;
}