// SPDX-FileCopyrightText: Copyright 2026 Eden Emulator Project
// SPDX-License-Identifier: GPL-3.0-or-later

// SPDX-FileCopyrightText: Copyright 2018 yuzu Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

#pragma once

#include <array>
#include <atomic>
#include <cstddef>
#include <memory>
#include <optional>
#include <span>
#include <vector>
#include "common/bit_field.h"
#include "common/common_funcs.h"
#include "common/common_types.h"
#include "video_core/engines/engine_interface.h"
#include "video_core/engines/engine_upload.h"
#include "video_core/macro.h"
#include "video_core/native_qmd.h"
#include "video_core/textures/texture.h"

namespace Core {
class System;
}

namespace VideoCommon {
class GpuSplit;
}

namespace Tegra {
class MemoryManager;
}

namespace VideoCore {
class RasterizerInterface;
}

namespace Tegra::Engines {

/// Draws left to trace after an upscale fallback (debug aid).
extern std::atomic<u32> g_nx2_trace_draws_after_upscale;

/**
 * This Engine is known as GK104_Compute. Documentation can be found in:
 * https://github.com/envytools/envytools/blob/master/rnndb/graph/gk104_compute.xml
 * https://cgit.freedesktop.org/mesa/mesa/tree/src/gallium/drivers/nouveau/nvc0/nve4_compute.xml.h
 */

#define KEPLER_COMPUTE_REG_INDEX(field_name)                                                       \
    (offsetof(Tegra::Engines::KeplerCompute::Regs, field_name) / sizeof(u32))

#define LAUNCH_REG_INDEX(field_name)                                                               \
    (offsetof(Tegra::Engines::KeplerCompute::LaunchParams, field_name) / sizeof(u32))

class KeplerCompute final : public EngineInterface {
public:
    explicit KeplerCompute(MemoryManager& memory_manager);
    ~KeplerCompute();

    /// Binds a rasterizer to this engine.
    void BindRasterizer(VideoCore::RasterizerInterface* rasterizer);

    /// The channel's MME data RAM, shared with its 3D engine.
    void SetNativeDataRam(std::array<u32, 0xc00>* data_ram) { native_data_ram = data_ram; }
    void SetNativeMacroMode(MacroEngine* shared_macro, std::span<u32> positions, NativeMacroContext* context = nullptr) {
        if (native_context != context) native_compute_uploaded = false;
        native_context = context;
        native_mode = shared_macro != nullptr;
        upload_state.SetNativePartialTransfers(native_mode);
        native_macro = shared_macro;
        native_macro_positions = positions;
        if (native_mode) execution_mask.set();
        else {
            execution_mask.reset();
            execution_mask[KEPLER_COMPUTE_REG_INDEX(exec_upload)] = true;
            execution_mask[KEPLER_COMPUTE_REG_INDEX(data_upload)] = true;
            execution_mask[KEPLER_COMPUTE_REG_INDEX(launch)] = true;
        }
    }

    static constexpr std::size_t NumConstBuffers = 8;

    struct Regs {
        static constexpr std::size_t NUM_REGS = 0xCF8;

        union {
            struct {
                INSERT_PADDING_WORDS_NOINIT(0x60);

                Upload::Registers upload;

                struct {
                    union {
                        BitField<0, 1, u32> linear;
                    };
                } exec_upload;

                u32 data_upload;

                INSERT_PADDING_WORDS_NOINIT(0x3F);

                struct {
                    u32 address;
                    GPUVAddr Address() const {
                        return GPUVAddr{address} << 8;
                    }
                } launch_desc_loc;

                INSERT_PADDING_WORDS_NOINIT(0x1);

                u32 launch;

                INSERT_PADDING_WORDS_NOINIT(0x4A7);

                struct {
                    u32 address_high;
                    u32 address_low;
                    u32 limit;
                    GPUVAddr Address() const {
                        return (GPUVAddr{address_high} << 32) | GPUVAddr{address_low};
                    }
                } tsc;

                INSERT_PADDING_WORDS_NOINIT(0x3);

                struct {
                    u32 address_high;
                    u32 address_low;
                    u32 limit;
                    GPUVAddr Address() const {
                        return (GPUVAddr{address_high} << 32) | GPUVAddr{address_low};
                    }
                } tic;

                INSERT_PADDING_WORDS_NOINIT(0x22);

                struct {
                    u32 address_high;
                    u32 address_low;
                    GPUVAddr Address() const {
                        return (GPUVAddr{address_high} << 32) | GPUVAddr{address_low};
                    }
                } code_loc;

                INSERT_PADDING_WORDS_NOINIT(0x3FE);

                u32 tex_cb_index;

                INSERT_PADDING_WORDS_NOINIT(0x374);
            };
            std::array<u32, NUM_REGS> reg_array;
        };
    } regs{};

    struct LaunchParams {
        static constexpr std::size_t NUM_LAUNCH_PARAMETERS = 0x40;

        INSERT_PADDING_WORDS(0x8);

        u32 program_start;

        INSERT_PADDING_WORDS(0x2);

        BitField<30, 1, u32> linked_tsc;

        BitField<0, 31, u32> grid_dim_x;
        union {
            BitField<0, 16, u32> grid_dim_y;
            BitField<16, 16, u32> grid_dim_z;
        };

        INSERT_PADDING_WORDS(0x3);

        BitField<0, 18, u32> shared_alloc;

        BitField<16, 16, u32> block_dim_x;
        union {
            BitField<0, 16, u32> block_dim_y;
            BitField<16, 16, u32> block_dim_z;
        };

        union {
            BitField<0, 8, u32> const_buffer_enable_mask;
            BitField<29, 2, u32> cache_layout;
        };

        INSERT_PADDING_WORDS(0x8);

        struct ConstBufferConfig {
            u32 address_low;
            union {
                BitField<0, 8, u32> address_high;
                BitField<15, 17, u32> size;
            };
            GPUVAddr Address() const {
                return (GPUVAddr{address_high.Value()} << 32) | GPUVAddr{address_low};
            }
        };
        std::array<ConstBufferConfig, NumConstBuffers> const_buffer_config;

        union {
            BitField<0, 20, u32> local_pos_alloc;
            BitField<27, 5, u32> barrier_alloc;
        };

        union {
            BitField<0, 20, u32> local_neg_alloc;
            BitField<24, 5, u32> gpr_alloc;
        };

        union {
            BitField<0, 20, u32> local_crs_alloc;
            BitField<24, 5, u32> sass_version;
        };

        INSERT_PADDING_WORDS(0x10);
    } launch_description{};

    struct {
        u32 write_offset = 0;
        u32 copy_size = 0;
        std::vector<u8> inner_buffer;
    } state{};

    static_assert(sizeof(Regs) == Regs::NUM_REGS * sizeof(u32),
                  "KeplerCompute Regs has wrong size");

    static_assert(sizeof(LaunchParams) == LaunchParams::NUM_LAUNCH_PARAMETERS * sizeof(u32),
                  "KeplerCompute LaunchParams has wrong size");

    /// Write the value to the register identified by method.
    void CallMethod(Core::System& system, u32 method, u32 method_argument, bool is_last_call) override;

    /// Write multiple values to the register identified by method.
    void CallMultiMethod(Core::System& system, u32 method, const u32* base_start, u32 amount, u32 methods_pending) override;

    std::optional<GPUVAddr> GetIndirectComputeAddress() const {
        return indirect_compute;
    }

    bool IsNativeMode() const { return native_mode; }

    /// GPU thread split with macros on its front half (gpu_split.h, EDEN_NX2_SPLIT=2), called by
    /// the front instead of CallMethod for native compute engines: runs their macros and macro
    /// uploads against a copy of the registers and sends the register writes to the back.
    bool FrontRunsMacros() const { return native_mode; }
    void FrontCallMethod(Core::System& system, VideoCommon::GpuSplit& split, u32 method,
                         u32 argument, bool is_last_call, const GPUVAddr* segment);
    void FrontCallMultiMethod(Core::System& system, VideoCommon::GpuSplit& split, u32 method,
                              const u32* base_start, u32 amount, u32 methods_pending,
                              const GPUVAddr* segment);

private:
    struct FrontState;
    FrontState& Front(VideoCommon::GpuSplit& split);
    void FrontApply(FrontState& state, u32 method, u32 argument);
    MacroEngine& FrontMacroEngine(FrontState& state);
    std::span<u32> FrontMacroPositions(FrontState& state);
    void FrontRunMacro(Core::System& system, VideoCommon::GpuSplit& split, FrontState& state);
    std::unique_ptr<FrontState> front;
    bool front_macros{}; ///< The front owns the macro code and bindings (the back skips them)

    bool DispatchNativeCompute(const NativeQmdScheduler::Descriptor& descriptor);
    void ProcessNativeReportSemaphore(u32 argument);
    void DescribeNativeResources(const NativeQmdScheduler::Descriptor& qmd);
    void RememberDlssMotionVectors();
    bool UpscaleInsteadOfNativeProgram();

    /// Motion vectors and depth of Fire Emblem: Fortune's Weave's DLSS evaluation, from its first
    /// pass (run earlier in the frame than the network pass the stand-in replaces), with the input
    /// colour image and camera jitter they go with.
    std::optional<Texture::TICEntry> dlss_motion_vectors;
    std::optional<Texture::TICEntry> dlss_depth;
    GPUVAddr dlss_motion_color{};
    std::array<f32, 2> dlss_jitter{};
    bool ProcessNativeMethod(Core::System& system, u32 method, u32 argument, bool is_last_call);
    NativeQmdScheduler native_qmd;
    std::array<u32, 0xc00>* native_data_ram{};
    bool native_mme_dumped{};
    std::optional<u32> native_executing_macro_logged;
    bool native_mode{};
    NativeMacroContext* native_context{};
    bool native_compute_uploaded{};
    MacroEngine* native_macro{};
    std::array<u32, 0x1000> native_registers{};
    /// Native methods that only store their value (no macro, QMD, upload, semaphore or code
    /// side effect): CallMethod writes those directly.
    std::bitset<0xe00> native_plain_methods;
    std::span<u32> native_macro_positions;
    std::vector<u32> native_parameters;
    std::vector<u32> native_memo_parameters; ///< The current call's parameters, for the memo
    std::optional<u32> native_executing_macro;
    u32 native_bind_pointer{};

    void ProcessLaunch();

    void ConsumeSinkImpl(Core::System& system) override;

    /// Retrieves information about a specific TIC entry from the TIC buffer.
    Texture::TICEntry GetTICEntry(u32 tic_index) const;

    /// Retrieves information about a specific TSC entry from the TSC buffer.
    Texture::TSCEntry GetTSCEntry(u32 tsc_index) const;

    MemoryManager& memory_manager;
    VideoCore::RasterizerInterface* rasterizer = nullptr;
    Upload::State upload_state;
    GPUVAddr upload_address;
    bool upload_dirty{};

    struct UploadInfo {
        GPUVAddr upload_address;
        GPUVAddr exec_address;
        u32 copy_size;
        bool was_dirty;
    };
    std::vector<UploadInfo> uploads;
    std::optional<GPUVAddr> indirect_compute{};
};

#define ASSERT_REG_POSITION(field_name, position)                                                  \
    static_assert(offsetof(KeplerCompute::Regs, field_name) == position * 4,                       \
                  "Field " #field_name " has invalid position")

#define ASSERT_LAUNCH_PARAM_POSITION(field_name, position)                                         \
    static_assert(offsetof(KeplerCompute::LaunchParams, field_name) == position * 4,               \
                  "Field " #field_name " has invalid position")

ASSERT_REG_POSITION(upload, 0x60);
ASSERT_REG_POSITION(exec_upload, 0x6C);
ASSERT_REG_POSITION(data_upload, 0x6D);
ASSERT_REG_POSITION(launch, 0xAF);
ASSERT_REG_POSITION(tsc, 0x557);
ASSERT_REG_POSITION(tic, 0x55D);
ASSERT_REG_POSITION(code_loc, 0x582);
ASSERT_REG_POSITION(tex_cb_index, 0x982);
ASSERT_LAUNCH_PARAM_POSITION(program_start, 0x8);
ASSERT_LAUNCH_PARAM_POSITION(grid_dim_x, 0xC);
ASSERT_LAUNCH_PARAM_POSITION(shared_alloc, 0x11);
ASSERT_LAUNCH_PARAM_POSITION(block_dim_x, 0x12);
ASSERT_LAUNCH_PARAM_POSITION(const_buffer_enable_mask, 0x14);
ASSERT_LAUNCH_PARAM_POSITION(const_buffer_config, 0x1D);

#undef ASSERT_REG_POSITION

} // namespace Tegra::Engines