/* ************************************************************************
 * Copyright (C) 2025-2026 Advanced Micro Devices, Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 *
 * ************************************************************************ */
#pragma once

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "stinkytofu/Export.hpp"
#include "stinkytofu/core/Function.hpp"
#include "stinkytofu/core/IRBase.hpp"
#include "stinkytofu/pipeline/PassBuilder.hpp"

/*
 * @brief Define the options for the ModuleOptions struct
 * @note This macro is used to define the options for the ModuleOptions struct
 * @note SwPrefetchScratchSgpr: -1 disables SwPrefetchInsertionPass; >=0 runs and uses that scratch.
 *        StinkyAsmModule sets EnableSwPrefetchInsertion = (SwPrefetchScratchSgpr != -1) in its
 * constructor.
 */
#define MODULE_OPTIONS_LIST(X)              \
    X(DebugLevel, int)                      \
    X(OptLevel, int)                        \
    X(TileA0, int)                          \
    X(TileB0, int)                          \
    X(TileM0, int)                          \
    X(NumGRA, uint32_t)                     \
    X(NumGRB, uint32_t)                     \
    X(NumGRM, uint32_t)                     \
    X(wavefrontSize, int)                   \
    X(SubGroup0, int)                       \
    X(SubGroup1, int)                       \
    X(WaveGroup0, int)                      \
    X(WaveGroup1, int)                      \
    X(VectorWidthA, int)                    \
    X(VectorWidthB, int)                    \
    X(GlobalReadVectorWidthA, int)          \
    X(GlobalReadVectorWidthB, int)          \
    X(DirectToLdsA, bool)                   \
    X(DirectToLdsB, bool)                   \
    X(UseSgprForGRO, int)                   \
    X(PrintBeforePass, std::string)         \
    X(PrintAfterPass, std::string)          \
    X(DebugPass, std::string)               \
    X(VerifyEach, bool)                     \
    X(EnableRemarks, bool)                  \
    X(EnableWaitCntInsertion, bool)         \
    X(EnableLoopCarriedTokenDeps, bool)     \
    X(EnableESM2, bool)                     \
    X(VgprMsbMode, int)                     \
    X(EnableSwPrefetchInsertion, bool)      \
    X(SwPrefetchScratchSgpr, int)           \
    X(ClusterBarrier, bool)                 \
    X(PrefetchGlobalRead, int)              \
    X(PrefetchLocalRead, int)

namespace stinkytofu {
/**
 * @brief Assembly IR Module container
 *
 * StinkyAsmModule holds low-level, architecture-specific assembly instructions
 * (stinkytofu::IRBase*) that can be emitted as assembly text.
 *
 * This class provides a container for assembly instructions generated by
 * lowering passes or directly created through the StinkyTofu IR builders.
 *
 * Note: This class is planned for deprecation in favor of using Function/BasicBlock
 * directly, but is currently needed for compatibility with existing Python bindings
 * and rocisa conversion utilities.
 *
 * Architecture:
 *   LogicalModule (high-level IR) -> Lowering Passes -> StinkyAsmModule (assembly IR)
 *
 * Example usage:
 * @code
 *   auto module = std::make_shared<StinkyAsmModule>("myKernel");
 *
 *   // Add assembly instructions
 *   std::vector<IRBase*> insts = {...};
 *   module->add(insts);
 *
 *   // Generate assembly string
 *   std::string asm = module->emitAssembly();
 * @endcode
 */
class STINKYTOFU_EXPORT StinkyAsmModule {
   public:
    /**
     * @brief Options for the StinkyAsmModule
     * @note This struct is used to store the information for the StinkyAsmModule
     */
    struct ModuleOptions {
#define GEN_MEMBER_OPTION(name, type) type name{};
        MODULE_OPTIONS_LIST(GEN_MEMBER_OPTION)
#undef GEN_MEMBER_OPTION
    };

    /**
     * @brief Construct a new StinkyAsmModule
     * @param name Module/kernel name
     * @param arch Target GPU architecture [major, minor, stepping]
     * @param moduleOptions Module options
     */
    StinkyAsmModule(const std::string& name, const std::array<int, 3>& arch,
                    const ModuleOptions& moduleOptions);

    /**
     * @brief Destructor - cleans up owned instructions
     */
    ~StinkyAsmModule();

    // Disable copy (we manage instruction ownership)
    StinkyAsmModule(const StinkyAsmModule&) = delete;
    StinkyAsmModule& operator=(const StinkyAsmModule&) = delete;

    // Enable move
    StinkyAsmModule(StinkyAsmModule&&) noexcept;
    StinkyAsmModule& operator=(StinkyAsmModule&&) noexcept;

    /**
     * @brief Get the module name
     * @return Module name string
     */
    std::string getName() const;

    /**
     * @brief Set the name used for output files (e.g. aggregated_instruction_cost.txt).
     * When set, Backend writes <outputName>_aggregated_instruction_cost.txt so it matches
     * the full kernel name (e.g. .o basename). When empty, getName() is used.
     * @param name Full kernel name for output file basename
     */
    void setOutputName(const std::string& name);

    /**
     * @brief Get the output file basename (cost file, etc.). Empty means use getName().
     * @return Output name string, or empty to use module name
     */
    std::string getOutputName() const;

    /**
     * @brief Set the directory for output files (e.g. cost file).
     * When set, Backend writes to <outputDir>/<kernel_full_name>/aggregated_instruction_cost.txt
     * (e.g. comparison_output/1024_vgpr_gfx1250/<full_name>/). When empty, files go to cwd.
     * @param dir Path such as "comparison_output/1024_vgpr_gfx1250"
     */
    void setOutputDir(const std::string& dir);

    /**
     * @brief Get the output directory. Empty means use current working directory.
     */
    std::string getOutputDir() const;

    /**
     * @brief Get the target architecture
     * @return Architecture array [major, minor, stepping]
     */
    std::array<int, 3> getArch() const;

    /**
     * @brief Emit assembly code for all instructions
     * @return Assembly code as string
     */
    std::string emitAssembly() const;

    /**
     * @brief Run optimization pipeline on the module
     */
    void runOptimizationPipeline();

    /**
     * @brief Read uint64 metadata from the underlying Function by key.
     * @param key Metadata key
     * @return Metadata value if key exists
     */
    std::optional<uint64_t> getMetaDataU64(const std::string& key) const;

    /**
     * @brief Get the underlying Function
     *
     * This provides access to the internal Function representation.
     * The Function is owned by the StinkyAsmModule.
     *
     * @return Reference to the Function
     */
    Function& getFunction();

    const Function& getFunction() const;

    /**
     * @brief Add a group name to the module
     * @param name Group name
     */
    void addGroup(const std::string& name);

    /**
     * @brief Check if the module has the given group name
     * @param name Group name
     * @return True if the module has the given group name, false otherwise
     */
    bool hasGroup(const std::string& name) const;

    /**
     * @brief Find the range of instructions in the given group
     * @param groupName Group name
     * @return Optional range of instructions (begin, end) for the given group
     */
    std::optional<std::pair<IntrusiveListIterator<IRBase>, IntrusiveListIterator<IRBase>>>
    findGroupRange(const std::string& groupName) const;

    /**
     * @brief Update the group ranges
     * @param groups Group names if exists to update
     * @param instsCountBefore Number of instructions before updating
     */
    void updateInstructionGroups(const std::vector<const std::string*>& groups,
                                 size_t instsCountBefore);

    /**
     * @brief Set the stored range for a group
     * @param groupName Group name
     * @param first First instruction in the range
     * @param last Last instruction in the range
     */
    void setGroupRange(const std::string& groupName, IntrusiveListIterator<IRBase> first,
                       IntrusiveListIterator<IRBase> last);

    /**
     * @brief Get the ModuleOptions
     * @return ModuleOptions
     */
    const ModuleOptions& getModuleOptions() const;

    /**
     * @brief Set the ModuleOptions
     * @param moduleOptions ModuleOptions
     */
    void setModuleOptions(const ModuleOptions& moduleOptions);

    /**
     * @brief Set total instruction size in bytes (encoding size) for the module.
     * Used to emit .amdhsa_inst_pref_size (totalBytes/128). Typically set by the
     * backend after running the optimization pipeline.
     */
    void setTotalInstructionBytes(int64_t totalBytes);

    /**
     * @brief Get total instruction size in bytes, or -1 if not set.
     */
    int64_t getTotalInstructionBytes() const;

    // ---- Plugin data (opaque key-value store for pass plugins) ----

    void setPluginDataI64(const std::string& key, int64_t value);
    int64_t getPluginDataI64(const std::string& key, int64_t defaultVal = 0) const;

    void setPluginDataStr(const std::string& key, const std::string& value);
    std::string getPluginDataStr(const std::string& key, const std::string& defaultVal = "") const;

    // ---- Pass plugin support ----

    PassBuilder& getPassBuilder();
    const PassBuilder& getPassBuilder() const;

   private:
    struct Impl;
    std::unique_ptr<Impl> pImpl;
    ModuleOptions moduleOptions;
};

}  // namespace stinkytofu
