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

#include <compare>
#include <fstream>

#include "CustomMatchers.hpp"
#include "CustomSections.hpp"
#include "ExpressionMatchers.hpp"
#include "TestContext.hpp"

#include <catch2/catch_test_macros.hpp>
#include <catch2/generators/catch_generators.hpp>
#include <catch2/matchers/catch_matchers_exception.hpp>
#include <catch2/matchers/catch_matchers_string.hpp>

#include <common/CommonGraphs.hpp>
#include <common/SourceMatcher.hpp>
#include <common/Utilities.hpp>

#include <rocRoller/AssemblyKernel.hpp>
#include <rocRoller/KernelGraph/ControlGraph/ControlFlowArgumentTracer.hpp>
#include <rocRoller/KernelGraph/ControlGraph/ControlGraph.hpp>
#include <rocRoller/KernelGraph/CoordinateGraph/CoordinateGraph.hpp>
#include <rocRoller/KernelGraph/KernelGraph.hpp>
#include <rocRoller/KernelGraph/Transforms/All.hpp>

namespace ArgumentTracerTest
{
    using namespace rocRoller;
    namespace KG = KernelGraph;
    namespace CT = KG::CoordinateGraph;
    namespace CG = KG::ControlGraph;

    TEST_CASE("AddDeallocateArguments transform works.", "[kernel-graph]")
    {
        auto context = TestContext::ForDefaultTarget();

        auto example           = rocRollerTest::Graphs::VectorAddNegSquare<int>();
        auto command           = example.getCommand();
        auto commandParameters = example.getCommandParameters();

        auto one = Expression::literal(1);
        context->kernel()->setWorkgroupSize({64, 1, 1});
        context->kernel()->setWorkitemCount({one, one, one});

        auto findDeallocate
            = [](KG::KernelGraph const& kg, std::string const& arg) -> std::optional<int> {
            for(auto nodeID : kg.control.getNodes<CG::Deallocate>())
            {
                auto const& node = kg.control.getNode<CG::Deallocate>(nodeID);
                if(std::find(node.arguments.begin(), node.arguments.end(), arg)
                   != node.arguments.end())
                    return nodeID;
            }

            return {};
        };

        auto findPtrComArg = [](CommandPtr               command,
                                KG::KernelGraph const&   kg,
                                Operations::OperationTag tag) -> CommandArgumentPtr {
            auto args = command->getArguments();

            auto match = fmt::format("{}_pointer", static_cast<int>(tag));

            for(auto const& arg : args)
            {
                if(arg->name().find(match) != std::string::npos)
                    return arg;
            }

            return nullptr;
        };

        auto findKernarg
            = [&context](CommandArgumentPtr const& comArg) -> AssemblyKernelArgumentPtr {
            auto expr   = std::make_shared<Expression::Expression>(comArg);
            auto argExp = context->kernel()->findArgumentForExpression(expr);
            if(argExp == nullptr)
                return nullptr;

            return std::get<AssemblyKernelArgumentPtr>(*argExp);
        };

        auto findDimForKernarg
            = [](KG::KernelGraph const& kg, std::string const& argName) -> std::optional<int> {
            for(auto dimId : kg.coordinates.getNodes<CT::User>())
            {
                auto dim = kg.coordinates.getNode<CT::User>(dimId);

                if(dim.argumentName == argName)
                    return dimId;
            }

            return {};
        };

        auto kgraph = KG::translate(command);

        kgraph = transform<KG::UpdateParameters>(kgraph, commandParameters);
        kgraph = transform<KG::LowerLinear>(kgraph, context.get());
        kgraph = transform<KG::CleanArguments>(kgraph, context.get(), command);
        kgraph = transform<KG::UpdateWavefrontParameters>(kgraph, commandParameters);
        kgraph = transform<KG::SetWorkitemCount>(kgraph, context.get());
        kgraph = transform<KG::AddDeallocateArguments>(kgraph, context.get());
        kgraph = transform<KG::SetWorkitemCount>(kgraph, context.get());
        kgraph = transform<KG::RemoveSetCoordinate>(kgraph);

        auto hasUserMapping = [](KG::KernelGraph const& kg, int userDim) {
            auto pred
                = [&kg, userDim](int node) { return kg.mapper.get<CT::User>(node) == userDim; };

            return pred;
        };

        SECTION("A pointer access")
        {
            auto aPtrArg = findKernarg(findPtrComArg(command, kgraph, example.aTag));
            REQUIRE(aPtrArg != nullptr);

            auto aDim = findDimForKernarg(kgraph, aPtrArg->getName());
            REQUIRE(aDim != std::nullopt);
            auto aPtrDeallocate = findDeallocate(kgraph, aPtrArg->getName());
            REQUIRE(aPtrDeallocate != std::nullopt);

            auto aLoad = kgraph.control.getNodes<CG::LoadVGPR>()
                             .filter(hasUserMapping(kgraph, *aDim))
                             .only();
            REQUIRE(aLoad != std::nullopt);
            CHECK(CG::NodeOrdering::LeftFirst
                  == kgraph.control.compareNodes(UpdateCache, *aLoad, *aPtrDeallocate));
        }

        SECTION("B pointer access")
        {
            auto bPtrArg = findKernarg(findPtrComArg(command, kgraph, example.bTag));
            REQUIRE(bPtrArg != nullptr);

            auto bDim = findDimForKernarg(kgraph, bPtrArg->getName());
            REQUIRE(bDim != std::nullopt);
            auto bPtrDeallocate = findDeallocate(kgraph, bPtrArg->getName());
            REQUIRE(bPtrDeallocate != std::nullopt);

            auto bLoad = kgraph.control.getNodes<CG::LoadVGPR>()
                             .filter(hasUserMapping(kgraph, *bDim))
                             .only();
            REQUIRE(bLoad != std::nullopt);
            CHECK(CG::NodeOrdering::LeftFirst
                  == kgraph.control.compareNodes(UpdateCache, *bLoad, *bPtrDeallocate));
        }

        SECTION("D pointer access")
        {
            auto dPtrArg = findKernarg(findPtrComArg(command, kgraph, example.resultTag));
            REQUIRE(dPtrArg != nullptr);

            auto dDim = findDimForKernarg(kgraph, dPtrArg->getName());
            REQUIRE(dDim != std::nullopt);
            auto dPtrDeallocate = findDeallocate(kgraph, dPtrArg->getName());
            REQUIRE(dPtrDeallocate != std::nullopt);

            auto dLoad = kgraph.control.getNodes<CG::StoreVGPR>()
                             .filter(hasUserMapping(kgraph, *dDim))
                             .only();
            REQUIRE(dLoad != std::nullopt);
            CHECK(CG::NodeOrdering::LeftFirst
                  == kgraph.control.compareNodes(UpdateCache, *dLoad, *dPtrDeallocate));
        }

        SECTION("Regular codegen")
        {
            KG::ControlFlowArgumentTracer argTracer(kgraph, context->kernel());

            auto generate = [&]() -> Generator<Instruction> {
                co_yield context->kernel()->preamble();
                co_yield context->kernel()->prolog();

                co_yield KernelGraph::generate(kgraph, context->kernel(), std::move(argTracer));

                co_yield context->kernel()->postamble();
                co_yield context->kernel()->amdgpu_metadata();
            };

            CHECK_NOTHROW(context->schedule(generate()));
        }

        SECTION("Codegen with added kernel argument access: LowerFromKernelGraph's auditing of "
                "ControlFlowArgumentTracer.")
        {
            KG::ControlFlowArgumentTracer argTracer(kgraph, context->kernel());

            // Create an additional kernel argument after we've made the arg tracer.
            std::string const& newArgName = "newArg";
            context->kernel()->addArgument({newArgName, DataType::Int64, DataDirection::ReadOnly});
            auto newArg = std::make_shared<AssemblyKernelArgument>(
                context->kernel()->findArgument(newArgName));

            // Modify a random Assign node to access an additional kernel
            // argument after we've made the arg tracer.
            {
                auto assignId = kgraph.control.getNodes<CG::Assign>().take(1).only().value();
                auto assign   = kgraph.control.getNode<CG::Assign>(assignId);

                REQUIRE_FALSE(argTracer.referencedArguments(assignId).contains(newArgName));
                assign.expression
                    = assign.expression
                      + convert(DataType::Int32, std::make_shared<Expression::Expression>(newArg));

                CAPTURE(assignId, assign.expression, (argTracer.referencedArguments(assignId)));

                kgraph.control.setElement(assignId, std::move(assign));
            }

            auto generate = [&]() -> Generator<Instruction> {
                co_yield context->kernel()->preamble();
                co_yield context->kernel()->prolog();

                co_yield KernelGraph::generate(kgraph, context->kernel(), std::move(argTracer));

                co_yield context->kernel()->postamble();
                co_yield context->kernel()->amdgpu_metadata();
            };

            namespace m = Catch::Matchers;

            CHECK_THROWS_MATCHES(context->schedule(generate()),
                                 FatalError,
                                 m::MessageMatches(m::ContainsSubstring(newArgName)));
        }
    }

    TEST_CASE("AddDeallocateArguments removes unused kernel arguments", "[kernel-graph]")
    {
        auto context = TestContext::ForDefaultTarget();

        auto example           = rocRollerTest::Graphs::VectorAddNegSquare<int>();
        auto command           = example.getCommand();
        auto commandParameters = example.getCommandParameters();

        auto one = Expression::literal(1);
        context->kernel()->setWorkgroupSize({64, 1, 1});
        context->kernel()->setWorkitemCount({one, one, one});

        auto kgraph = KG::translate(command);

        kgraph = transform<KG::UpdateParameters>(kgraph, commandParameters);
        kgraph = transform<KG::LowerLinear>(kgraph, context.get());
        kgraph = transform<KG::UpdateWavefrontParameters>(kgraph, commandParameters);
        kgraph = transform<KG::CleanArguments>(kgraph, context.get(), command);

        // Add a new unused argument
        context->kernel()->addArgument(
            {"unusedArg", {DataType::Int32, PointerType::PointerGlobal}, DataDirection::WriteOnly});

        // Ensure the unused arg has been added to the kernel successfully
        CHECK_NOTHROW(context->kernel()->findArgument("unusedArg"));

        kgraph = transform<KG::AddDeallocateArguments>(kgraph, context.get());

        // Verify unused argument is removed after AddDeallocateArguments by
        // checking an error is thrown when finding it in the kernel.
        CHECK_THROWS_AS(context->kernel()->findArgument("unusedArg"), FatalError);
    }
}