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

#pragma once

#include <cstdint>
#include <initializer_list>
#include <string>
#include <vector>

#include <rocRoller/CodeGen/BufferInstructionOptions.hpp>
#include <rocRoller/Expression_fwd.hpp>
#include <rocRoller/InstructionValues/Register_fwd.hpp>
#include <rocRoller/KernelGraph/ControlGraph/Operation_fwd.hpp>
#include <rocRoller/KernelGraph/CoordinateGraph/Dimension.hpp>
#include <rocRoller/KernelGraph/RegisterTagManager.hpp>
#include <rocRoller/KernelGraph/StructUtils.hpp>
#include <rocRoller/Operations/BlockScale_fwd.hpp>
#include <rocRoller/Utilities/Utils.hpp>

namespace rocRoller
{
    namespace KernelGraph::ControlGraph
    {
        /**
         * @brief Selects how a ConditionalOp is lowered to GPU instructions.
         *
         * Branch:        Scalar branch-based conditional. Suitable for uniform conditions where
         *                all lanes take the same path.
         * Exec:          Exec-mask-based conditional for per-lane (VGPR) conditions. Both the
         *                true and else bodies are always entered; only active lanes satisfying
         *                the condition execute Body, and only active lanes not satisfying it
         *                execute Else. EXEC is restored afterward.
         * BranchAndExec: Like Exec, but additionally branches over each body when EXECZ is set
         *                (i.e. no lanes are active), avoiding unnecessary work.
         */
        enum class ConditionalMode
        {
            Branch = 0,
            Exec,
            BranchAndExec,
            Count
        };

        std::string   toString(ConditionalMode m);
        std::ostream& operator<<(std::ostream& stream, ConditionalMode m);

        /*
         * Control flow graph nodes.
         * Represent operations done on the input.
         */

        /**
         * Kernel - represents the start of a kernel.
         */
        RR_EMPTY_STRUCT_WITH_NAME(Kernel);

        /**
         * Scope - represents a register scope.
         */
        RR_EMPTY_STRUCT_WITH_NAME(Scope);

        /**
         * SetCoordinate - Sets the value of a Coordinate
         */
        struct SetCoordinate
        {
            SetCoordinate();
            explicit SetCoordinate(Expression::ExpressionPtr value);

            Expression::ExpressionPtr value;

            std::string name() const;
        };

        /**
         * DoWhileLoopOp - Represents a do-while loop.
         *
         * Must have nodes connected via the following outgoing edges:
         *
         * - Body: The loop body. The loop body must cause a change in the condition, this body will also be emitted at least once.
         *
         * There may be multiple outgoing edges for any of these.  Code that follows the for loop should be connected via a Sequence edge.
         *
         * condition is a scalar or vector condition and is executed before each iteration to determine if we must exit the loop.
         *
         * Currently generates code that behaves like:
         *
         * while_top:
         * <Body>
         * if(condition) goto while_top
         * <Sequence>
         */
        struct DoWhileOp
        {
            Expression::ExpressionPtr condition;

            std::string loopName;

            std::string name() const;
            std::string toString() const;
        };

        /**
         * ForLoopOp - Represents a for loop.
         *
         * Must have nodes connected via the following outgoing edges:
         *
         * - Initialize: Always executed once, when entering the for loop
         * - Body: The loop body.
         * - ForLoopIncrement: Executed after each iteration.
         *
         * There may be multiple outgoing edges for any of these.  Code that follows the for loop should be connected via a Sequence edge.
         *
         * condition is a scalar condition and is executed before each iteration to determine if we must exit the for loop.
         *
         * Currently generates code that behaves like:
         *
         * <Initialize>
         * if(!condition) goto for_bottom
         * for_top:
         * <Body>
         * <ForLoopIncrement>
         * if(condition) goto for_top
         * for_bottom:
         * <Sequence>
         */
        struct ForLoopOp
        {
            Expression::ExpressionPtr condition;

            std::string loopName;

            std::string name() const;
            std::string toString() const;
        };

        /**
         * ConditionalOp - Represents a conditional with one of three execution modes.
         *
         * Outgoing edges:
         *   - Body     : true branch (required)
         *   - Else     : false branch (optional)
         *   - Sequence : code that runs after both branches, regardless of the condition
         *
         * The `mode` field selects how the conditional is lowered:
         *
         * ConditionalMode::Branch
         *   Scalar branch-based conditional. Suitable for uniform conditions where all
         *   lanes take the same path.
         *   if (condition) { <Body> } else { <Else> }
         *
         * ConditionalMode::Exec
         *   Exec-mask-based conditional for per-lane (VGPR) conditions. Both the true
         *   and else sections are always executed: only active lanes satisfying the
         *   condition execute <Body>, and only active lanes not satisfying the condition
         *   execute <Else>. EXEC is restored afterward.
         *
         * ConditionalMode::BranchAndExec
         *   Like Exec (including EXEC restore), but checks EXECZ after masking: if no
         *   active lanes satisfy the condition (EXECZ is set), the true body is skipped
         *   and execution jumps to the else section (where EXECZ is checked again upon
         *   masking with ~condition). If some active lanes satisfy the condition, the
         *   true body executes for those lanes and the else body is skipped for all lanes.
        */
        struct ConditionalOp
        {
            Expression::ExpressionPtr condition;
            ConditionalMode           mode;
            std::string               conditionName;

            std::string name() const;
            std::string toString() const;
        };

        /**
         * AssertOp - Represents an assert.
         *
         * Must have nodes connected via the following outgoing edges:
         *
         * - True  Sequence:
         * - False <terminates kernel>:
         *
         * Currently generates code that behaves like:
         *
         * if(not condition)
         *   <terminates kernel>
         * <Sequence>
         *
         * Where <terminates kernel> is a instruction sequence that causes a trap or exception in kernel code.
         *
        */
        struct AssertOp
        {
            std::string assertName;

            Expression::ExpressionPtr condition;

            std::string name() const;
            std::string toString() const;
        };

        /**
         * UnrollOp - a kernel unroll.
         */
        struct UnrollOp
        {
            Expression::ExpressionPtr size;

            std::string name() const;
            std::string toString() const;
        };

        /*
         * Computes the value of `expression` and stores it into the associated register.
         *
         * If the register already exists, it must be of type 'regType'.  If not, `regType`
         * specifies which type of register will be allocated.
         */
        struct Assign
        {
            Register::Type            regType = Register::Type::Count;
            Expression::ExpressionPtr expression;

            size_t valueCount = 1;

            // If variableType is a packed type then
            // (valueCount / variableType.packing) registers will be allocated.
            std::optional<VariableType> variableType = std::nullopt;

            // If the destination coordinate is Stride then
            // set the register expression attributes
            std::optional<RegisterExpressionAttributes> strideExpressionAttributes = std::nullopt;

            std::string name() const;
            std::string toString() const;
        };

        /**
         * @brief Represents a memory barrier
         *
         */
        RR_EMPTY_STRUCT_WITH_NAME(Barrier);

        /**
         * @brief Deallocates a register tag.
         */
        struct Deallocate
        {
            std::vector<std::string> arguments;

            std::string name() const;
            std::string toString() const;
        };

        /**
         * LoadLinear - Load linear dimension.
         */
        struct LoadLinear
        {
            LoadLinear();
            explicit LoadLinear(rocRoller::VariableType const varType);

            rocRoller::VariableType varType;

            std::string name() const;
        };

        /**
         * LoadTiled.  Loads a tile (typically a MacroTile or
         * WaveTile).
         *
         * Storage location (LDS, VGPR, etc) is specified by the
         * `MemoryType` member of the MacroTile node.
         *
         * When loading a WaveTile, the storage layout (for MFMA
         * instructions) is specified by the `LayoutType` member of
         * the the WaveTile node.
         */
        struct LoadTiled
        {
            LoadTiled();
            explicit LoadTiled(VariableType const varType);

            VariableType varType;

            std::string name() const;
        };

        /**
         * LoadVGPR - replaces LoadLinear.
         */
        struct LoadVGPR
        {
            LoadVGPR();
            explicit LoadVGPR(VariableType const varType, bool const scalar = false);

            VariableType varType;
            bool         scalar;

            std::string name() const;
        };

        /**
         * LoadSGPR - load scalar value from memory.
         */
        struct LoadSGPR
        {
            LoadSGPR();
            LoadSGPR(VariableType const varType, BufferInstructionOptions const bio);

            VariableType             varType;
            BufferInstructionOptions bufOpts;

            std::string name() const;
        };

        /**
         * LoadLDSTile - loads a tile from LDS
         */
        struct LoadLDSTile
        {
            LoadLDSTile();
            explicit LoadLDSTile(VariableType const varType, bool const isTransposedTile = false);

            VariableType varType;
            bool         isTransposedTile;

            std::string name() const;
            std::string toString() const;
        };

        struct LoadTileDirect2LDS
        {
            LoadTileDirect2LDS();
            explicit LoadTileDirect2LDS(VariableType const varType);

            VariableType varType;

            std::string name() const;
        };

        struct LoadTiledTDMToLDS
        {
            LoadTiledTDMToLDS();
            explicit LoadTiledTDMToLDS(VariableType const varType);

            VariableType varType;

            std::string name() const;
        };

        /**
         * Multiply - Multiply two MacroTiles
         */
        struct Multiply
        {
            Multiply();
            Multiply(Operations::ScaleMode scaleA, Operations::ScaleMode scaleB);

            Operations::ScaleMode scaleA;
            Operations::ScaleMode scaleB;

            std::string name() const;
        };

        /**
         * NOP - Do nothing.
         */
        RR_EMPTY_STRUCT_WITH_NAME(NOP);

        /**
         * Block - similar to NOP but Block locks the scheduler
         */
        RR_EMPTY_STRUCT_WITH_NAME(Block);

        /**
         * StoreLinear - Store linear dimension.
         */
        RR_EMPTY_STRUCT_WITH_NAME(StoreLinear);

        /**
         * StoreTiled.  Stores a tile.
         *
         * Storage location and affinity is specified by the MacroTile
         * node.
         */
        struct StoreTiled
        {
            StoreTiled();
            explicit StoreTiled(VariableType const dtype);

            VariableType             varType = DataType::Count;
            BufferInstructionOptions bufOpts;

            std::string name() const;
        };

        /**
         * StoreVGPR - replaces StoreLinear.
         */
        RR_EMPTY_STRUCT_WITH_NAME(StoreVGPR);

        /**
         * StoreSGPR - stores a scalar value to memory.
         */
        struct StoreSGPR
        {
            StoreSGPR();
            StoreSGPR(VariableType const varType, BufferInstructionOptions const bio);

            VariableType             varType;
            BufferInstructionOptions bufOpts;

            std::string name() const;
        };

        /**
         * StoreLDSTile - store a tile into LDS
         */
        struct StoreLDSTile
        {
            StoreLDSTile();
            explicit StoreLDSTile(VariableType const varType);

            VariableType varType;

            std::string name() const;
        };

        /**
         * TensorContraction - Tensor contraction operation.
         */
        struct TensorContraction
        {
            TensorContraction();
            TensorContraction(std::vector<int> const& aContractedDimensions,
                              std::vector<int> const& bContractedDimensions,
                              VariableType const      accType = DataType::Float);

            std::vector<int>      aDims, bDims; // contracted dimensions
            Operations::ScaleMode scaleModeA = Operations::ScaleMode::None;
            Operations::ScaleMode scaleModeB = Operations::ScaleMode::None;
            std::vector<size_t>   scaleStridesA;
            std::vector<size_t>   scaleStridesB;
            std::vector<size_t>   scalePreShuffledTileA;
            std::vector<size_t>   scalePreShuffledTileB;
            VariableType          accType = DataType::Float;

            std::string name() const;
        };

        /**
         * WaitZero - Emit a Wait Count of zero on all wait queues.
         *
         * This is important in preventing certain race conditions.
         * It forces the wait queues to be emptied before proceeding
         * to the next graph nodes (connected by Sequence edges).
         *
         * Example:
         * Store tile -> WaitZero -> Store sync flags
         */
        RR_EMPTY_STRUCT_WITH_NAME(WaitZero);

        /**
         * Exchange - permute the lanes data within a wave.
         */
        struct Exchange
        {
            Exchange();
            explicit Exchange(VariableType const varType);

            VariableType varType;

            std::string name() const;
        };

        /**
         * SeedPRNG - Set the initial seed value of a random number generator
         */
        struct SeedPRNG
        {
            SeedPRNG();
            explicit SeedPRNG(bool addTID);
            std::string toString() const;

            std::string name() const;

            // Add workitem ID to the seed if this flag is true
            bool addTID = false;
        };

        template <CConcreteOperation Op>
        std::string name(const Op& x);

        /*
         * Helpers
         */
        std::string name(const Operation& x);

        std::string toString(const Operation& x);

        /**
         * @brief Return the datatype associated with the Operation.
         */
        DataType getDataType(const Operation& x);

        VariableType getVariableType(Operation const& op);
        void         setVariableType(Operation& op, VariableType varType);
    }
}

#include <rocRoller/KernelGraph/ControlGraph/Operation_impl.hpp>