#include "XCEngine/RHI/OpenGL/OpenGLPipelineState.h"
#include "XCEngine/RHI/OpenGL/OpenGLShader.h"
#include "XCEngine/RHI/OpenGL/OpenGLEnums.h"
#include <cstring>
#include <glad/glad.h>

namespace XCEngine {
namespace RHI {

OpenGLPipelineState::OpenGLPipelineState()
    : m_program(0), m_programAttached(false) {
}

OpenGLPipelineState::~OpenGLPipelineState() {
}

void OpenGLPipelineState::SetInputLayout(const InputLayoutDesc& layout) {
    m_inputLayoutDesc = layout;
}

void OpenGLPipelineState::SetRasterizerState(const RasterizerDesc& state) {
    m_rasterizerDesc = state;
    m_glRasterizerState.cullFaceEnable = static_cast<CullMode>(state.cullMode) != CullMode::None;
    m_glRasterizerState.cullFace = static_cast<CullMode>(state.cullMode);
    m_glRasterizerState.frontFace = static_cast<FrontFace>(state.frontFace);
    m_glRasterizerState.polygonMode = static_cast<FillMode>(state.fillMode);
    m_glRasterizerState.depthClipEnable = state.depthClipEnable;
    m_glRasterizerState.scissorTestEnable = state.scissorTestEnable;
    m_glRasterizerState.multisampleEnable = state.multisampleEnable;
    m_glRasterizerState.polygonOffsetFactor = state.slopeScaledDepthBias;
    m_glRasterizerState.polygonOffsetUnits = static_cast<float>(state.depthBias);
}

void OpenGLPipelineState::SetBlendState(const BlendDesc& state) {
    m_blendDesc = state;
    m_glBlendState.blendEnable = state.blendEnable;
    m_glBlendState.srcBlend = static_cast<BlendFactor>(state.srcBlend);
    m_glBlendState.dstBlend = static_cast<BlendFactor>(state.dstBlend);
    m_glBlendState.srcBlendAlpha = static_cast<BlendFactor>(state.srcBlendAlpha);
    m_glBlendState.dstBlendAlpha = static_cast<BlendFactor>(state.dstBlendAlpha);
    m_glBlendState.blendOp = static_cast<BlendOp>(state.blendOp);
    m_glBlendState.blendOpAlpha = static_cast<BlendOp>(state.blendOpAlpha);
    m_glBlendState.colorWriteMask = state.colorWriteMask;
    std::memcpy(m_glBlendState.blendFactor, state.blendFactor, sizeof(state.blendFactor));
}

void OpenGLPipelineState::SetDepthStencilState(const DepthStencilStateDesc& state) {
    m_depthStencilDesc = state;
    m_glDepthStencilState.depthTestEnable = state.depthTestEnable;
    m_glDepthStencilState.depthWriteEnable = state.depthWriteEnable;
    m_glDepthStencilState.depthFunc = static_cast<ComparisonFunc>(state.depthFunc);
    m_glDepthStencilState.stencilEnable = state.stencilEnable;
    m_glDepthStencilState.stencilReadMask = state.stencilReadMask;
    m_glDepthStencilState.stencilWriteMask = state.stencilWriteMask;
    m_glDepthStencilState.frontStencilFunc = static_cast<ComparisonFunc>(state.front.func);
    m_glDepthStencilState.frontStencilFailOp = static_cast<StencilOp>(state.front.failOp);
    m_glDepthStencilState.frontStencilDepthFailOp = static_cast<StencilOp>(state.front.depthFailOp);
    m_glDepthStencilState.frontStencilDepthPassOp = static_cast<StencilOp>(state.front.passOp);
    m_glDepthStencilState.backStencilFunc = static_cast<ComparisonFunc>(state.back.func);
    m_glDepthStencilState.backStencilFailOp = static_cast<StencilOp>(state.back.failOp);
    m_glDepthStencilState.backStencilDepthFailOp = static_cast<StencilOp>(state.back.depthFailOp);
    m_glDepthStencilState.backStencilDepthPassOp = static_cast<StencilOp>(state.back.passOp);
}

void OpenGLPipelineState::SetTopology(uint32_t topologyType) {
    m_topologyType = topologyType;
}

void OpenGLPipelineState::SetRenderTargetFormats(uint32_t count, const uint32_t* formats, uint32_t depthFormat) {
}

void OpenGLPipelineState::SetSampleCount(uint32_t count) {
}

void OpenGLPipelineState::SetComputeShader(RHIShader* shader) {
    m_computeShader = shader;
    if (shader) {
        OpenGLShader* glShader = static_cast<OpenGLShader*>(shader);
        m_computeProgram = glShader->GetID();
    } else {
        m_computeProgram = 0;
    }
}

PipelineStateHash OpenGLPipelineState::GetHash() const {
    PipelineStateHash hash = {};
    return hash;
}

void OpenGLPipelineState::Shutdown() {
    m_graphicsShader.reset();
    m_program = 0;
    m_computeProgram = 0;
    m_computeShader = nullptr;
    m_programAttached = false;
}

void OpenGLPipelineState::Bind() {
    if (HasComputeShader()) {
        glUseProgram(m_computeProgram);
    } else if (m_programAttached && m_program != 0) {
        glUseProgram(m_program);
    }
    Apply();
}

void OpenGLPipelineState::Unbind() {
    glUseProgram(0);
}

void OpenGLPipelineState::Apply() {
    ApplyDepthStencil();
    ApplyBlend();
    ApplyRasterizer();
}

void OpenGLPipelineState::ApplyDepthStencil() {
    if (m_glDepthStencilState.depthTestEnable) {
        glEnable(GL_DEPTH_TEST);
    } else {
        glDisable(GL_DEPTH_TEST);
    }
    glDepthMask(m_glDepthStencilState.depthWriteEnable ? GL_TRUE : GL_FALSE);
    glDepthFunc(ToOpenGL(m_glDepthStencilState.depthFunc));

    if (m_glDepthStencilState.stencilEnable) {
        glEnable(GL_STENCIL_TEST);
        glStencilMaskSeparate(GL_FRONT, m_glDepthStencilState.stencilWriteMask);
        glStencilMaskSeparate(GL_BACK, m_glDepthStencilState.stencilWriteMask);
        glStencilFuncSeparate(
            GL_FRONT,
            ToOpenGL(m_glDepthStencilState.frontStencilFunc),
            m_glDepthStencilState.stencilRef,
            m_glDepthStencilState.stencilReadMask
        );
        glStencilOpSeparate(
            GL_FRONT,
            ToOpenGL(m_glDepthStencilState.frontStencilFailOp),
            ToOpenGL(m_glDepthStencilState.frontStencilDepthFailOp),
            ToOpenGL(m_glDepthStencilState.frontStencilDepthPassOp)
        );
        glStencilFuncSeparate(
            GL_BACK,
            ToOpenGL(m_glDepthStencilState.backStencilFunc),
            m_glDepthStencilState.stencilRef,
            m_glDepthStencilState.stencilReadMask
        );
        glStencilOpSeparate(
            GL_BACK,
            ToOpenGL(m_glDepthStencilState.backStencilFailOp),
            ToOpenGL(m_glDepthStencilState.backStencilDepthFailOp),
            ToOpenGL(m_glDepthStencilState.backStencilDepthPassOp)
        );
    } else {
        glDisable(GL_STENCIL_TEST);
    }
}

void OpenGLPipelineState::ApplyBlend() {
    if (m_glBlendState.blendEnable) {
        glEnable(GL_BLEND);
        glBlendFuncSeparate(
            ToOpenGL(m_glBlendState.srcBlend),
            ToOpenGL(m_glBlendState.dstBlend),
            ToOpenGL(m_glBlendState.srcBlendAlpha),
            ToOpenGL(m_glBlendState.dstBlendAlpha)
        );
        glBlendEquationSeparate(
            ToOpenGL(m_glBlendState.blendOp),
            ToOpenGL(m_glBlendState.blendOpAlpha)
        );
        glColorMask(
            (m_glBlendState.colorWriteMask & 1) != 0,
            (m_glBlendState.colorWriteMask & 2) != 0,
            (m_glBlendState.colorWriteMask & 4) != 0,
            (m_glBlendState.colorWriteMask & 8) != 0
        );
    } else {
        glDisable(GL_BLEND);
    }
}

void OpenGLPipelineState::ApplyRasterizer() {
    if (m_glRasterizerState.cullFaceEnable) {
        glEnable(GL_CULL_FACE);
        glCullFace(ToOpenGL(m_glRasterizerState.cullFace));
        glFrontFace(ToOpenGL(m_glRasterizerState.frontFace));
    } else {
        glDisable(GL_CULL_FACE);
    }

    glPolygonMode(GL_FRONT_AND_BACK, ToOpenGL(m_glRasterizerState.polygonMode));

    if (m_glRasterizerState.multisampleEnable) {
        glEnable(GL_MULTISAMPLE);
    } else {
        glDisable(GL_MULTISAMPLE);
    }

    const bool polygonOffsetEnable =
        m_glRasterizerState.polygonOffsetFactor != 0.0f ||
        m_glRasterizerState.polygonOffsetUnits != 0.0f;
    if (polygonOffsetEnable) {
        glEnable(GL_POLYGON_OFFSET_FILL);
        glPolygonOffset(
            m_glRasterizerState.polygonOffsetFactor,
            m_glRasterizerState.polygonOffsetUnits);
    } else {
        glDisable(GL_POLYGON_OFFSET_FILL);
    }
}

void OpenGLPipelineState::ApplyViewport() {
    glViewport((GLint)m_viewportState.x, (GLint)m_viewportState.y,
               (GLsizei)m_viewportState.width, (GLsizei)m_viewportState.height);
    glDepthRange(m_viewportState.minDepth, m_viewportState.maxDepth);
}

void OpenGLPipelineState::ApplyScissor() {
    if (m_scissorState.enable) {
        glEnable(GL_SCISSOR_TEST);
        glScissor(m_scissorState.x, m_scissorState.y, m_scissorState.width, m_scissorState.height);
    } else {
        glDisable(GL_SCISSOR_TEST);
    }
}

void OpenGLPipelineState::SetProgram(unsigned int program) {
    m_program = program;
    m_programAttached = (program != 0);
}

void OpenGLPipelineState::SetDepthStencilState(const OpenGLDepthStencilState& state) {
    m_glDepthStencilState = state;
}

void OpenGLPipelineState::SetBlendState(const OpenGLBlendState& state) {
    m_glBlendState = state;
}

void OpenGLPipelineState::SetRasterizerState(const OpenGLRasterizerState& state) {
    m_glRasterizerState = state;
}

void OpenGLPipelineState::SetViewport(const ViewportState& state) {
    m_viewportState = state;
}

void OpenGLPipelineState::SetScissor(const ScissorState& state) {
    m_scissorState = state;
}

void OpenGLPipelineState::SetClearColor(float r, float g, float b, float a) {
    m_clearColor[0] = r;
    m_clearColor[1] = g;
    m_clearColor[2] = b;
    m_clearColor[3] = a;
    glClearColor(r, g, b, a);
}

void OpenGLPipelineState::Clear(unsigned int buffers) {
    glClear(ToOpenGLClearBuffer(buffers));
}

void OpenGLPipelineState::AttachShader(unsigned int program) {
    m_program = program;
    m_programAttached = true;
    glUseProgram(program);
}

void OpenGLPipelineState::DetachShader() {
    m_program = 0;
    m_programAttached = false;
    glUseProgram(0);
}

void OpenGLPipelineState::SetOwnedGraphicsShader(std::unique_ptr<OpenGLShader> shader) {
    m_graphicsShader = std::move(shader);
    if (m_graphicsShader) {
        SetProgram(m_graphicsShader->GetID());
    } else {
        DetachShader();
    }
}

const OpenGLDepthStencilState& OpenGLPipelineState::GetOpenGLDepthStencilState() const {
    return m_glDepthStencilState;
}

const OpenGLBlendState& OpenGLPipelineState::GetOpenGLBlendState() const {
    return m_glBlendState;
}

const OpenGLRasterizerState& OpenGLPipelineState::GetOpenGLRasterizerState() const {
    return m_glRasterizerState;
}

} // namespace RHI
} // namespace XCEngine