// Copyright 2015 Citra Emulator Project
// Licensed under GPLv2 or any later version
// Refer to the license.txt file included.

#pragma once

#include <array>
#include <cstddef>
#include <type_traits>
#include <nihstro/shader_bytecode.h>
#include "common/assert.h"
#include "common/common_funcs.h"
#include "common/common_types.h"
#include "common/vector_math.h"
#include "video_core/pica.h"
#include "video_core/pica_types.h"

using nihstro::RegisterType;
using nihstro::SourceRegister;
using nihstro::DestRegister;

namespace Pica {

namespace Shader {

struct AttributeBuffer {
    alignas(16) Math::Vec4<float24> attr[16];
};

struct OutputVertex {
    Math::Vec4<float24> pos;
    Math::Vec4<float24> quat;
    Math::Vec4<float24> color;
    Math::Vec2<float24> tc0;
    Math::Vec2<float24> tc1;
    float24 tc0_w;
    INSERT_PADDING_WORDS(1);
    Math::Vec3<float24> view;
    INSERT_PADDING_WORDS(1);
    Math::Vec2<float24> tc2;

    static OutputVertex FromAttributeBuffer(const Regs& regs, AttributeBuffer& output);
};
#define ASSERT_POS(var, pos)                                                                       \
    static_assert(offsetof(OutputVertex, var) == pos * sizeof(float24), "Semantic at wrong "       \
                                                                        "offset.")
ASSERT_POS(pos, Regs::VSOutputAttributes::POSITION_X);
ASSERT_POS(quat, Regs::VSOutputAttributes::QUATERNION_X);
ASSERT_POS(color, Regs::VSOutputAttributes::COLOR_R);
ASSERT_POS(tc0, Regs::VSOutputAttributes::TEXCOORD0_U);
ASSERT_POS(tc1, Regs::VSOutputAttributes::TEXCOORD1_U);
ASSERT_POS(tc0_w, Regs::VSOutputAttributes::TEXCOORD0_W);
ASSERT_POS(view, Regs::VSOutputAttributes::VIEW_X);
ASSERT_POS(tc2, Regs::VSOutputAttributes::TEXCOORD2_U);
#undef ASSERT_POS
static_assert(std::is_pod<OutputVertex>::value, "Structure is not POD");
static_assert(sizeof(OutputVertex) == 24 * sizeof(float), "OutputVertex has invalid size");

/**
 * This structure contains the state information that needs to be unique for a shader unit. The 3DS
 * has four shader units that process shaders in parallel. At the present, Citra only implements a
 * single shader unit that processes all shaders serially. Putting the state information in a struct
 * here will make it easier for us to parallelize the shader processing later.
 */
struct UnitState {
    struct Registers {
        // The registers are accessed by the shader JIT using SSE instructions, and are therefore
        // required to be 16-byte aligned.
        alignas(16) Math::Vec4<float24> input[16];
        alignas(16) Math::Vec4<float24> temporary[16];
        alignas(16) Math::Vec4<float24> output[16];
    } registers;
    static_assert(std::is_pod<Registers>::value, "Structure is not POD");

    bool conditional_code[2];

    // Two Address registers and one loop counter
    // TODO: How many bits do these actually have?
    s32 address_registers[3];

    static size_t InputOffset(const SourceRegister& reg) {
        switch (reg.GetRegisterType()) {
        case RegisterType::Input:
            return offsetof(UnitState, registers.input) +
                   reg.GetIndex() * sizeof(Math::Vec4<float24>);

        case RegisterType::Temporary:
            return offsetof(UnitState, registers.temporary) +
                   reg.GetIndex() * sizeof(Math::Vec4<float24>);

        default:
            UNREACHABLE();
            return 0;
        }
    }

    static size_t OutputOffset(const DestRegister& reg) {
        switch (reg.GetRegisterType()) {
        case RegisterType::Output:
            return offsetof(UnitState, registers.output) +
                   reg.GetIndex() * sizeof(Math::Vec4<float24>);

        case RegisterType::Temporary:
            return offsetof(UnitState, registers.temporary) +
                   reg.GetIndex() * sizeof(Math::Vec4<float24>);

        default:
            UNREACHABLE();
            return 0;
        }
    }

    /**
     * Loads the unit state with an input vertex.
     *
     * @param config Shader configuration registers corresponding to the unit.
     * @param input Attribute buffer to load into the input registers.
     */
    void LoadInput(const Regs::ShaderConfig& config, const AttributeBuffer& input);

    void WriteOutput(const Regs::ShaderConfig& config, AttributeBuffer& output);
};

struct ShaderSetup {
    struct {
        // The float uniforms are accessed by the shader JIT using SSE instructions, and are
        // therefore required to be 16-byte aligned.
        alignas(16) Math::Vec4<float24> f[96];

        std::array<bool, 16> b;
        std::array<Math::Vec4<u8>, 4> i;
    } uniforms;

    static size_t GetFloatUniformOffset(unsigned index) {
        return offsetof(ShaderSetup, uniforms.f) + index * sizeof(Math::Vec4<float24>);
    }

    static size_t GetBoolUniformOffset(unsigned index) {
        return offsetof(ShaderSetup, uniforms.b) + index * sizeof(bool);
    }

    static size_t GetIntUniformOffset(unsigned index) {
        return offsetof(ShaderSetup, uniforms.i) + index * sizeof(Math::Vec4<u8>);
    }

    std::array<u32, 1024> program_code;
    std::array<u32, 1024> swizzle_data;

    /// Data private to ShaderEngines
    struct EngineData {
        unsigned int entry_point;
        /// Used by the JIT, points to a compiled shader object.
        const void* cached_shader = nullptr;
    } engine_data;
};

class ShaderEngine {
public:
    virtual ~ShaderEngine() = default;

    /**
     * Performs any shader unit setup that only needs to happen once per shader (as opposed to once
     * per vertex, which would happen within the `Run` function).
     */
    virtual void SetupBatch(ShaderSetup& setup, unsigned int entry_point) = 0;

    /**
     * Runs the currently setup shader.
     *
     * @param setup Shader engine state, must be setup with SetupBatch on each shader change.
     * @param state Shader unit state, must be setup with input data before each shader invocation.
     */
    virtual void Run(const ShaderSetup& setup, UnitState& state) const = 0;
};

// TODO(yuriks): Remove and make it non-global state somewhere
ShaderEngine* GetEngine();
void Shutdown();

} // namespace Shader

} // namespace Pica