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

#pragma once

#include <fstream>
#include <iostream>
#include <memory>
#include <ranges>
#include <string>
#include <unordered_map>
#include <vector>

#include <rocRoller/AssemblyKernel_fwd.hpp>
#include <rocRoller/CodeGen/ArgumentLoader_fwd.hpp>
#include <rocRoller/CodeGen/BranchGenerator_fwd.hpp>
#include <rocRoller/CodeGen/CopyGenerator_fwd.hpp>
#include <rocRoller/CodeGen/CrashKernelGenerator_fwd.hpp>
#include <rocRoller/CodeGen/Instruction_fwd.hpp>
#include <rocRoller/CodeGen/MemoryInstructions_fwd.hpp>
#include <rocRoller/Context_fwd.hpp>
#include <rocRoller/DataTypes/DataTypes.hpp>
#include <rocRoller/Expression_fwd.hpp>
#include <rocRoller/GPUArchitecture/GPUArchitecture.hpp>
#include <rocRoller/GPUArchitecture/GPUArchitectureTarget_fwd.hpp>
#include <rocRoller/InstructionValues/LDSAllocator_fwd.hpp>
#include <rocRoller/InstructionValues/LabelAllocator_fwd.hpp>
#include <rocRoller/InstructionValues/RegisterAllocator_fwd.hpp>
#include <rocRoller/InstructionValues/Register_fwd.hpp>
#include <rocRoller/KernelGraph/RegisterTagManager_fwd.hpp>
#include <rocRoller/KernelGraph/ScopeManager_fwd.hpp>
#include <rocRoller/KernelOptions.hpp>
#include <rocRoller/Operations/Scratch_fwd.hpp>
#include <rocRoller/ScheduledInstructions_fwd.hpp>
#include <rocRoller/Scheduling/Scheduling_fwd.hpp>
#include <rocRoller/Utilities/Random_fwd.hpp>

class ContextFixture;

namespace rocRoller
{
    class Context : public std::enable_shared_from_this<Context>
    {
    public:
        Context();
        ~Context();

        static ContextPtr ForDefaultHipDevice(std::string const&   kernelName,
                                              KernelOptions const& kernelOpts = {});

        static ContextPtr ForHipDevice(int                  deviceIdx,
                                       std::string const&   kernelName,
                                       KernelOptions const& kernelOpts = {});

        static ContextPtr ForTarget(GPUArchitectureTarget const& arch,
                                    std::string const&           kernelName,
                                    KernelOptions const&         kernelOpts = {});

        static ContextPtr ForTarget(GPUArchitecture const& arch,
                                    std::string const&     kernelName,
                                    KernelOptions const&   kernelOpts = {});

        Scheduling::InstructionStatus peek(Instruction const& inst);

        void schedule(Instruction& inst);
        template <std::ranges::input_range T>
        void schedule(T const& insts);
        template <std::ranges::input_range T>
        void schedule(T&& insts);

        std::shared_ptr<Register::Allocator> allocator(Register::Type registerType);

        Register::ValuePtr getM0();
        Register::ValuePtr getVCC();
        Register::ValuePtr getVCC_LO();
        Register::ValuePtr getVCC_HI();
        Register::ValuePtr getSCC();
        Register::ValuePtr getEXECZ();
        Register::ValuePtr getEXEC();
        Register::ValuePtr getTTMP6();
        Register::ValuePtr getTTMP7();
        Register::ValuePtr getTTMP9();
        Register::ValuePtr getSpecial(Register::Type t);

        std::shared_ptr<Scheduling::IObserver> observer() const;

        AssemblyKernelPtr                      kernel() const;
        std::shared_ptr<ArgumentLoader>        argLoader() const;
        std::shared_ptr<ScheduledInstructions> instructions() const;
        std::shared_ptr<MemoryInstructions>    mem() const;
        std::shared_ptr<CopyGenerator>         copier() const;
        std::shared_ptr<BranchGenerator>       brancher() const;
        std::shared_ptr<CrashKernelGenerator>  crasher() const;
        KernelOptions const&                   kernelOptions();

        std::shared_ptr<RandomGenerator> random() const;
        void                             setRandomSeed(int seed);

        std::string assemblyFileName() const;

        LabelAllocatorPtr             labelAllocator() const;
        std::shared_ptr<LDSAllocator> ldsAllocator() const;
        RegTagManPtr                  registerTagManager() const;
        GPUArchitecture const&        targetArchitecture() const;
        int                           hipDeviceIndex() const;

        void setKernel(AssemblyKernelPtr);

        /**
         * @brief Allocate scratch space for the specified scratch policy.
         *
         * @param policy The scratch policy to allocate for
         * @param size Number of bytes requested
         * @return Expression::ExpressionPtr The offset before this allocation
         */
        Expression::ExpressionPtr allocateScratch(Operations::ScratchPolicy policy,
                                                  Expression::ExpressionPtr size);

        /**
         * @brief Returns an expression representing how much scratch space is required (in bytes)
         *        for the specified scratch policy.
         *
         * @param policy The scratch policy to query
         * @return Expression::ExpressionPtr
         */
        Expression::ExpressionPtr getScratchAmount(Operations::ScratchPolicy policy) const;

        /**
         * @brief Get register scope manager.
         */
        std::shared_ptr<KernelGraph::ScopeManager> getScopeManager() const;

        /**
         * @brief Set register scope manager.
         */
        void setScopeManager(std::shared_ptr<KernelGraph::ScopeManager>);

        friend class ::ContextFixture;

    private:
        static ContextPtr Create(int                    deviceIndex,
                                 GPUArchitecture const& arch,
                                 std::string const&     kernelName,
                                 KernelOptions const&   kernelOpts);

        void scheduleCopy(Instruction const& inst);

        // If we are generating code for a real Hip device, refers to its
        // device index.
        int             m_hipDeviceIdx = -1;
        GPUArchitecture m_targetArch;
        RegTagManPtr    m_registerTagMan;
        std::array<std::shared_ptr<Register::Allocator>, static_cast<size_t>(Register::Type::Count)>
            m_allocators;

        std::shared_ptr<Scheduling::IObserver> m_observer;
        AssemblyKernelPtr                      m_kernel;
        std::shared_ptr<ArgumentLoader>        m_argLoader;
        std::shared_ptr<ScheduledInstructions> m_instructions;
        std::shared_ptr<MemoryInstructions>    m_mem;
        LabelAllocatorPtr                      m_labelAllocator;
        std::shared_ptr<LDSAllocator>          m_ldsAllocator;
        std::array<Expression::ExpressionPtr, static_cast<size_t>(Operations::ScratchPolicy::Count)>
                                                   m_scratchSizes;
        std::shared_ptr<CopyGenerator>             m_copier;
        std::shared_ptr<BranchGenerator>           m_brancher;
        std::shared_ptr<CrashKernelGenerator>      m_crasher;
        std::shared_ptr<RandomGenerator>           m_random;
        std::shared_ptr<KernelGraph::ScopeManager> m_scope;

        std::string   m_assemblyFileName;
        KernelOptions m_kernelOptions;
    };

    std::ostream& operator<<(std::ostream&, ContextPtr const&);
}

#include <rocRoller/Context_impl.hpp>