#include <mln/mtl/render_pass.hpp>

#include <mln/mtl/command_encoder.hpp>
#include <mln/mtl/renderable_resource.hpp>
#include <mln/mtl/context.hpp>
#include <mln/util/logging.hpp>

#include <Metal/Metal.hpp>
#include <stdexcept>

namespace mln {
namespace mtl {

detail::ClearedAttachments detail::applyRenderPassClears(MTL::RenderPassDescriptor& rpd,
                                                          const std::optional<Color>& clearColor,
                                                          const std::optional<float>& clearDepth,
                                                          const std::optional<int32_t>& clearStencil) {
    ClearedAttachments result;
    if (clearColor) {
        if (auto* target = rpd.colorAttachments()->object(0); target && target->texture()) {
            const auto& c = *clearColor;
            target->setLoadAction(MTL::LoadActionClear);
            target->setClearColor(MTL::ClearColor::Make(c.r, c.g, c.b, c.a));
            target->setStoreAction(MTL::StoreActionStore);
            result.color = true;
        }
    }
    if (clearDepth) {
        if (auto* target = rpd.depthAttachment(); target && target->texture()) {
            target->setLoadAction(MTL::LoadActionClear);
            target->setClearDepth(*clearDepth);
            target->setStoreAction(MTL::StoreActionStore);
            result.depth = true;
        }
    }
    if (clearStencil) {
        if (auto* target = rpd.stencilAttachment(); target && target->texture()) {
            target->setLoadAction(MTL::LoadActionClear);
            target->setClearStencil(static_cast<uint32_t>(*clearStencil));
            target->setStoreAction(MTL::StoreActionStore);
            result.stencil = true;
        }
    }
    return result;
}

RenderPass::RenderPass(CommandEncoder& commandEncoder_, const char* name, const gfx::RenderPassDescriptor& descriptor)
    : descriptor(descriptor),
      commandEncoder(commandEncoder_) {
    auto& resource = descriptor.renderable.getResource<RenderableResource>();

    resource.bind();

    if (const auto& buffer = resource.getCommandBuffer()) {
        commandBuffer = buffer;
        if (const auto& original = resource.getRenderPassDescriptor()) {
            auto rpd = NS::TransferPtr(original->copy());
            if (!rpd) throw std::runtime_error("Metal render pass descriptor copy failed");
            detail::applyRenderPassClears(*rpd.get(), descriptor.clearColor, descriptor.clearDepth, descriptor.clearStencil);
            encoder = NS::RetainPtr(buffer->renderCommandEncoder(rpd.get()));
            if (!encoder) throw std::runtime_error("Metal render command encoder unavailable");

            const auto* colorTarget = rpd->colorAttachments()->object(0);
            if (!colorTarget || !colorTarget->texture()) throw std::runtime_error("Metal color attachment unavailable");
            const auto& texture = colorTarget->texture();
            width = texture->width();
            height = texture->height();
        }
    }

    if (!encoder) throw std::runtime_error("Metal render pass unavailable");

    // Push the groups already accumulated by the encoder
    commandEncoder.visitDebugGroups([this](const auto& group) {
        debugGroups.emplace_back(gfx::DebugGroup<gfx::RenderPass>{*this, group.c_str()});
    });

    // Push the group for the name provided
    debugGroups.emplace_back(gfx::DebugGroup<gfx::RenderPass>{*this, name});

    // Let the encoder pass along any groups pushed to it after this
    commandEncoder.trackRenderPass(this);

    commandEncoder.context.performCleanup();
}

RenderPass::~RenderPass() {
    commandEncoder.forgetRenderPass(this);
    endEncoding();
}

void RenderPass::endEncoding() {
    debugGroups.clear();

    if (encoder) {
        encoder->endEncoding();
        encoder.reset();
    }

    resetState();
}

void RenderPass::resetState() {
    currentPipelineState.reset();
    currentDepthStencilState.reset();
    currentStencilReferenceValue = 0;
    for (int i = 0; i < maxBinds; ++i) {
        vertexBinds[i].reset();
        fragmentBinds[i].reset();
        vertexTextureBindings[i].reset();
        vertexSamplerStates[i].reset();
        fragmentTextureBindings[i].reset();
        fragmentSamplerStates[i].reset();
    }

    currentCullMode = MTL::CullModeNone;
    currentWinding = MTL::WindingClockwise;
    currentScissorRect = {.x = 0, .y = 0, .width = 0, .height = 0};
}

namespace {
constexpr auto missing = "<none>";
NS::String* toNSString(const char* str) {
    return NS::String::string(str ? str : missing, NS::UTF8StringEncoding);
}
} // namespace

void RenderPass::pushDebugGroup(const char* theName) {
    assert(encoder);
    if (encoder) {
        encoder->pushDebugGroup(toNSString(theName));
    }
}

void RenderPass::popDebugGroup() {
    assert(encoder);
    if (encoder) {
        encoder->popDebugGroup();
    }
}

void RenderPass::addDebugSignpost(const char* theName) {
    assert(encoder);
    if (encoder) {
        encoder->insertDebugSignpost(toNSString(theName));
    }
}

void RenderPass::bindVertex(const BufferResource& buf, std::size_t offset, std::size_t index, std::size_t size) {
    const auto actualSize = size ? size : buf.getSizeInBytes() - offset;
    assert(actualSize <= buf.getSizeInBytes());
    assert(0 <= index && index < maxBinds);
    if (0 <= index && index < maxBinds) {
        if (auto& bind = vertexBinds[index]) {
            // Is this the same buffer already bound to this index?
            if (bind->buf == &buf && !buf.needReBind(bind->version)) {
                // Yes, but is the offset different?
                if (bind->offset != offset) {
                    // Yes, update just the offset
                    buf.updateVertexBindOffset(encoder, offset, index, actualSize);
                    bind->offset = offset;
                }
                return;
            }
        }
        vertexBinds[index] = BindInfo{.buf = &buf, .size = actualSize, .offset = offset};
    }
    buf.bindVertex(encoder, offset, index, actualSize);
}

void RenderPass::unbindVertex(std::size_t index) {
    vertexBinds[index] = std::nullopt;
}

void RenderPass::bindFragment(const BufferResource& buf, std::size_t offset, std::size_t index, std::size_t size) {
    const auto actualSize = size ? size : buf.getSizeInBytes() - offset;
    assert(actualSize <= buf.getSizeInBytes());
    assert(0 <= index && index < maxBinds);
    if (0 <= index && index < maxBinds) {
        if (auto& bind = fragmentBinds[index]) {
            // Is this the same buffer already bound to this index?
            if (bind->buf == &buf && !buf.needReBind(bind->version)) {
                // Yes, but is the offset different?
                if (bind->offset != offset) {
                    // Yes, update just the offset
                    buf.updateFragmentBindOffset(encoder, offset, index, actualSize);
                    bind->offset = offset;
                }
                return;
            }
        }
        fragmentBinds[index] = BindInfo{.buf = &buf, .size = actualSize, .offset = offset};
    }
    buf.bindFragment(encoder, offset, index, actualSize);
}

void RenderPass::unbindFragment(std::size_t index) {
    fragmentBinds[index] = std::nullopt;
}

void RenderPass::setDepthStencilState(const MTLDepthStencilStatePtr& state) {
    if (state != currentDepthStencilState) {
        currentDepthStencilState = state;
        encoder->setDepthStencilState(currentDepthStencilState.get());
    }
}

void RenderPass::setStencilReference(int32_t referenceValue) {
    if (referenceValue != currentStencilReferenceValue) {
        currentStencilReferenceValue = referenceValue;
        encoder->setStencilReferenceValue(currentStencilReferenceValue);
    }
}

void RenderPass::setVertexTexture(const MTLTexturePtr& texture, int32_t location) {
    assert(0 <= location && location < maxBinds);
    if (0 <= location && location < maxBinds && vertexTextureBindings[location] != texture) {
        vertexTextureBindings[location] = texture;
        encoder->setVertexTexture(texture.get(), location);
    }
}

void RenderPass::setVertexSamplerState(const MTLSamplerStatePtr& state, int32_t location) {
    assert(0 <= location && location < maxBinds);
    if (0 <= location && location < maxBinds && vertexSamplerStates[location] != state) {
        vertexSamplerStates[location] = state;
        encoder->setVertexSamplerState(state.get(), location);
    }
}

void RenderPass::setFragmentTexture(const MTLTexturePtr& texture, int32_t location) {
    assert(0 <= location && location < maxBinds);
    if (0 <= location && location < maxBinds) {
        if (fragmentTextureBindings[location] != texture) {
            fragmentTextureBindings[location] = texture;
            encoder->setFragmentTexture(texture.get(), location);
        }
    }
}

void RenderPass::setFragmentSamplerState(const MTLSamplerStatePtr& state, int32_t location) {
    assert(0 <= location && location < maxBinds);
    if (0 <= location && location < maxBinds) {
        if (fragmentSamplerStates[location] != state) {
            fragmentSamplerStates[location] = state;
            encoder->setFragmentSamplerState(state.get(), location);
        }
    }
}

/// Set the render pipeline state
void RenderPass::setRenderPipelineState(const MTLRenderPipelineStatePtr& pipelineState) {
    if (pipelineState != currentPipelineState) {
        currentPipelineState = pipelineState;
        encoder->setRenderPipelineState(currentPipelineState.get());
    }
}

void RenderPass::setCullMode(const MTL::CullMode mode) {
    if (mode != currentCullMode) {
        encoder->setCullMode(mode);
        currentCullMode = mode;
    }
}

void RenderPass::setFrontFacingWinding(const MTL::Winding winding) {
    if (winding != currentWinding) {
        encoder->setFrontFacingWinding(winding);
        currentWinding = winding;
    }
}

void RenderPass::setScissorRect(MTL::ScissorRect rect) {
    if (rect.x != currentScissorRect.x || rect.y != currentScissorRect.y || rect.width != currentScissorRect.width ||
        rect.height != currentScissorRect.height) {
        if (rect.width + rect.x > width) {
            rect.width = width - rect.x;
        }
        if (rect.height + rect.y > height) {
            rect.height = height - rect.y;
        }
        encoder->setScissorRect(rect);
        currentScissorRect = rect;
    }
}

} // namespace mtl
} // namespace mln