// metal-command-encoder.cpp
#include "metal-command-encoder.h"

#include "metal-buffer.h"
#include "metal-command-buffer.h"
#include "metal-query.h"
#include "metal-render-pass.h"
#include "metal-resource-views.h"
#include "metal-shader-object.h"
#include "metal-shader-program.h"
#include "metal-shader-table.h"
#include "metal-texture.h"

#include "metal-helper-functions.h"

namespace gfx
{

using namespace Slang;

namespace metal
{

void PipelineCommandEncoder::init(CommandBufferImpl* commandBuffer)
{
    m_commandBuffer = commandBuffer;
    m_device = commandBuffer->m_renderer;
    m_metalCommandBuffer = m_commandBuffer->m_commandBuffer;
}

void ResourceCommandEncoder::copyBuffer(
    IBufferResource* dst, Offset dstOffset, IBufferResource* src, Offset srcOffset, Size size)
{
}

void ResourceCommandEncoder::uploadBufferData(
    IBufferResource* buffer, Offset offset, Size size, void* data)
{
}

void ResourceCommandEncoder::textureBarrier(
    GfxCount count, ITextureResource* const* textures, ResourceState src, ResourceState dst)
{
}

// TODO: Change size_t to Count?
void ResourceCommandEncoder::bufferBarrier(
    GfxCount count, IBufferResource* const* buffers, ResourceState src, ResourceState dst)
{
}

void ResourceCommandEncoder::endEncoding()
{
}

void ResourceCommandEncoder::writeTimestamp(IQueryPool* queryPool, GfxIndex index)
{
}

void ResourceCommandEncoder::copyTexture(
    ITextureResource* dst,
    ResourceState dstState,
    SubresourceRange dstSubresource,
    ITextureResource::Offset3D dstOffset,
    ITextureResource* src,
    ResourceState srcState,
    SubresourceRange srcSubresource,
    ITextureResource::Offset3D srcOffset,
    ITextureResource::Extents extent)
{
}

void ResourceCommandEncoder::uploadTextureData(
    ITextureResource* dst,
    SubresourceRange subResourceRange,
    ITextureResource::Offset3D offset,
    ITextureResource::Extents extend,
    ITextureResource::SubresourceData* subResourceData,
    GfxCount subResourceDataCount)
{
}


void ResourceCommandEncoder::clearResourceView(
    IResourceView* view, ClearValue* clearValue, ClearResourceViewFlags::Enum flags)
{
}

void ResourceCommandEncoder::resolveResource(
    ITextureResource* source,
    ResourceState sourceState,
    SubresourceRange sourceRange,
    ITextureResource* dest,
    ResourceState destState,
    SubresourceRange destRange)
{
}

void ResourceCommandEncoder::resolveQuery(
    IQueryPool* queryPool, GfxIndex index, GfxCount count, IBufferResource* buffer, Offset offset)
{
}

void ResourceCommandEncoder::copyTextureToBuffer(
    IBufferResource* dst,
    Offset dstOffset,
    Size dstSize,
    Size dstRowStride,
    ITextureResource* src,
    ResourceState srcState,
    SubresourceRange srcSubresource,
    ITextureResource::Offset3D srcOffset,
    ITextureResource::Extents extent)
{
    assert(srcSubresource.mipLevelCount <= 1);
}

void ResourceCommandEncoder::textureSubresourceBarrier(
    ITextureResource* texture,
    SubresourceRange subresourceRange,
    ResourceState src,
    ResourceState dst)
{
}

void ResourceCommandEncoder::beginDebugEvent(const char* name, float rgbColor[3])
{
}

void ResourceCommandEncoder::endDebugEvent()
{
}

void RenderCommandEncoder::beginPass(IRenderPassLayout* renderPass, IFramebuffer* framebuffer)
{
    FramebufferImpl* fb = static_cast<FramebufferImpl*>(framebuffer);
    if (fb == nullptr)
    {
        return;
    }
    RenderPassLayoutImpl* renderPassLayoutImpl = static_cast<RenderPassLayoutImpl*>(renderPass);

    MTL::RenderPassDescriptor* rpd = renderPassLayoutImpl->m_renderPassDesc->copy();

    if (rpd->depthAttachment() && false)
    {
        TextureResourceViewImpl* depthView = static_cast<TextureResourceViewImpl*>(fb->depthStencilView.get());
        rpd->depthAttachment()->setTexture(depthView->m_texture->m_texture);
    }
    const int colorTargetCount = fb->renderTargetViews.getCount();
    for (int i = 0; i < colorTargetCount; ++i)
    {
        TextureResourceViewImpl* texView = static_cast<TextureResourceViewImpl*>(fb->renderTargetViews[i].get());
        MTL::Texture* tex = nullptr;
        assert(texView->m_texture);
        if (texView->m_texture->m_isCurrentDrawable)
        {
            CA::MetalDrawable* drawable = static_cast<CA::MetalDrawable*>(fb->m_renderer->m_drawable);
            tex = drawable->texture();
        }
        else
        {
            tex = texView->m_texture->m_texture;
        }
        rpd->colorAttachments()->object(i)->setTexture(tex);
        rpd->colorAttachments()->object(i)->setClearColor(MTL::ClearColor(0.2, 0.4, 0.9, 1.0));
    }
    rpd->setRenderTargetWidth(fb->m_width);
    rpd->setRenderTargetHeight(fb->m_height);

    m_encoder = m_metalCommandBuffer->renderCommandEncoder(rpd);
}

void RenderCommandEncoder::endEncoding()
{
    m_encoder->endEncoding();
}

Result RenderCommandEncoder::bindPipeline(
    IPipelineState* pipelineState, IShaderObject** outRootObject)
{
    m_currentPipeline = static_cast<PipelineStateImpl*>(pipelineState);
    // Initialize the root object
    SLANG_RETURN_ON_FAIL(m_commandBuffer->m_rootObject.init(m_commandBuffer->m_renderer,
        m_currentPipeline->getProgram<ShaderProgramImpl>()->m_rootObjectLayout));
    *outRootObject = &m_commandBuffer->m_rootObject;
    //if (pPipelineState->m_renderState == nullptr) return SLANG_ERROR_INVALID_PARAMETER;
    //m_encoder->setRenderPipelineState(pPipelineState->m_renderState);
    return SLANG_OK;
}

Result RenderCommandEncoder::bindPipelineWithRootObject(
    IPipelineState* pipelineState, IShaderObject* rootObject)
{
    return SLANG_E_NOT_IMPLEMENTED;
}

void RenderCommandEncoder::setViewports(GfxCount count, const Viewport* viewports)
{
    static const int kMaxViewports = 8; // TODO: base on device caps
    assert(count <= kMaxViewports);

    m_viewports.setCount(count);
    for (GfxIndex i = 0; i < count; ++i)
    {
        const auto& inViewport = viewports[i];
        auto& metalViewport = m_viewports[i];
        metalViewport.height = inViewport.extentY;
        metalViewport.width = inViewport.extentX;
        metalViewport.originX = inViewport.originX;
        metalViewport.originY = inViewport.originY;
        metalViewport.znear = inViewport.minZ;
        metalViewport.zfar = inViewport.maxZ;
    }
    m_encoder->setViewports(m_viewports.begin(), count);
}

void RenderCommandEncoder::setScissorRects(GfxCount count, const ScissorRect* rects)
{
    static const int kMaxScissorRects = 9; // TODO:
    assert(count < kMaxScissorRects);

    m_scissorRects.setCount(count);
    for (GfxIndex i = 0; i < count; ++i)
    {
        const auto& inRect = rects[i];
        auto& metalRect = m_scissorRects[i];
        metalRect.height = inRect.maxX - inRect.minX;
        metalRect.width = inRect.maxY - inRect.minY;
        metalRect.x = inRect.minX;
        metalRect.y = inRect.minY;
    }
    m_encoder->setScissorRects(m_scissorRects.begin(), count);
}

void RenderCommandEncoder::setPrimitiveTopology(PrimitiveTopology topology)
{
}

void RenderCommandEncoder::setVertexBuffers(
    GfxIndex startSlot,
    GfxCount slotCount,
    IBufferResource* const* buffers,
    const Offset* offsets)
{
    for (GfxIndex i = 0; i < GfxIndex(slotCount); i++)
    {
        GfxIndex slotIndex = startSlot + i;
        BufferResourceImpl* buffer = static_cast<BufferResourceImpl*>(buffers[i]);
        if (buffer)
        {
            MTL::Buffer* vertexBuffers = {buffer->m_buffer};
            m_encoder->setVertexBuffer(buffer->m_buffer, offsets[i], slotIndex);
            // ...
        }
    }
}

void RenderCommandEncoder::setIndexBuffer(
    IBufferResource* buffer, Format indexFormat, Offset offset)
{
}

Result RenderCommandEncoder::prepareDraw()
{
    // Bind render state, including JIT pipeline state object creation
    auto pipeline = static_cast<PipelineStateImpl*>(m_currentPipeline.Ptr());
    if (!pipeline)
    {
        return SLANG_FAIL;
    }
    // TODO: specialization, binding, ...
    SLANG_RETURN_ON_FAIL(pipeline->ensureAPIPipelineStateCreated());
    return SLANG_OK;
}

static Result translatePrimitiveType(gfx::PrimitiveType primType, MTL::PrimitiveType& mtlType)
{
    switch (primType) 
    {
    case PrimitiveType::Triangle:
        mtlType = MTL::PrimitiveTypeTriangle;
        break;
    case PrimitiveType::Line:
        mtlType = MTL::PrimitiveTypeLine;
        break;
    case PrimitiveType::Point:
        mtlType = MTL::PrimitiveTypePoint;
        break;
    case PrimitiveType::Patch:
    default:
        return SLANG_E_INVALID_ARG;
    }
    return SLANG_OK;
}

Result RenderCommandEncoder::draw(GfxCount vertexCount, GfxIndex startVertex)
{
    SLANG_RETURN_ON_FAIL(prepareDraw());

    MTL::PrimitiveType primType;
    Result res = translatePrimitiveType(m_currentPipeline->desc.graphics.primitiveType, primType);
    if (res != SLANG_OK)
        return res;
    m_encoder->drawPrimitives(primType, startVertex, vertexCount);
    return SLANG_OK;
}

Result RenderCommandEncoder::drawIndexed(
    GfxCount indexCount, GfxIndex startIndex, GfxIndex baseVertex)
{
    return SLANG_E_NOT_IMPLEMENTED;
}

void RenderCommandEncoder::setStencilReference(uint32_t referenceValue)
{
}

Result RenderCommandEncoder::drawIndirect(
    GfxCount maxDrawCount,
    IBufferResource* argBuffer,
    Offset argOffset,
    IBufferResource* countBuffer,
    Offset countOffset)
{
    return SLANG_E_NOT_IMPLEMENTED;
}

Result RenderCommandEncoder::drawIndexedIndirect(
    GfxCount maxDrawCount,
    IBufferResource* argBuffer,
    Offset argOffset,
    IBufferResource* countBuffer,
    Offset countOffset)
{
    return SLANG_E_NOT_IMPLEMENTED;
}

Result RenderCommandEncoder::setSamplePositions(
    GfxCount samplesPerPixel, GfxCount pixelCount, const SamplePosition* samplePositions)
{
    return SLANG_E_NOT_AVAILABLE;
}

Result RenderCommandEncoder::drawInstanced(
    GfxCount vertexCount,
    GfxCount instanceCount,
    GfxIndex startVertex,
    GfxIndex startInstanceLocation)
{
    SLANG_RETURN_ON_FAIL(prepareDraw());
    return SLANG_E_NOT_IMPLEMENTED;
}

Result RenderCommandEncoder::drawIndexedInstanced(
    GfxCount indexCount,
    GfxCount instanceCount,
    GfxIndex startIndexLocation,
    GfxIndex baseVertexLocation,
    GfxIndex startInstanceLocation)
{
    SLANG_RETURN_ON_FAIL(prepareDraw());
    return SLANG_E_NOT_IMPLEMENTED;
}

Result RenderCommandEncoder::drawMeshTasks(int x, int y, int z)
{
    SLANG_RETURN_ON_FAIL(prepareDraw());
    return SLANG_E_NOT_IMPLEMENTED;
}

void ComputeCommandEncoder::endEncoding() { }

Result ComputeCommandEncoder::bindPipeline(
    IPipelineState* pipelineState, IShaderObject** outRootObject)
{
    return SLANG_E_NOT_IMPLEMENTED;
}

Result ComputeCommandEncoder::bindPipelineWithRootObject(
    IPipelineState* pipelineState, IShaderObject* rootObject)
{
    return SLANG_E_NOT_IMPLEMENTED;
}

Result ComputeCommandEncoder::dispatchCompute(int x, int y, int z)
{
    auto pipeline = static_cast<PipelineStateImpl*>(m_currentPipeline.Ptr());
    if (!pipeline)
    {
        return SLANG_FAIL;
    }

    // Also create descriptor sets based on the given pipeline layout
    return SLANG_E_NOT_IMPLEMENTED;
}

Result ComputeCommandEncoder::dispatchComputeIndirect(IBufferResource* argBuffer, Offset offset)
{
    SLANG_UNIMPLEMENTED_X("dispatchComputeIndirect");
}

void RayTracingCommandEncoder::_memoryBarrier(
    int count,
    IAccelerationStructure* const* structures,
    AccessFlag srcAccess,
    AccessFlag destAccess)
{
}

void RayTracingCommandEncoder::_queryAccelerationStructureProperties(
    GfxCount accelerationStructureCount,
    IAccelerationStructure* const* accelerationStructures,
    GfxCount queryCount,
    AccelerationStructureQueryDesc* queryDescs)
{
}

void RayTracingCommandEncoder::buildAccelerationStructure(
    const IAccelerationStructure::BuildDesc& desc,
    GfxCount propertyQueryCount,
    AccelerationStructureQueryDesc* queryDescs)
{
}

void RayTracingCommandEncoder::copyAccelerationStructure(
    IAccelerationStructure* dest, IAccelerationStructure* src, AccelerationStructureCopyMode mode)
{
}

void RayTracingCommandEncoder::queryAccelerationStructureProperties(
    GfxCount accelerationStructureCount,
    IAccelerationStructure* const* accelerationStructures,
    GfxCount queryCount,
    AccelerationStructureQueryDesc* queryDescs)
{
    _queryAccelerationStructureProperties(
        accelerationStructureCount, accelerationStructures, queryCount, queryDescs);
}

void RayTracingCommandEncoder::serializeAccelerationStructure(
    DeviceAddress dest, IAccelerationStructure* source)
{
}

void RayTracingCommandEncoder::deserializeAccelerationStructure(
    IAccelerationStructure* dest, DeviceAddress source)
{
}

Result RayTracingCommandEncoder::bindPipeline(IPipelineState* pipeline, IShaderObject** outRootObject)
{
    return SLANG_E_NOT_IMPLEMENTED;
}

Result RayTracingCommandEncoder::bindPipelineWithRootObject(
    IPipelineState* pipelineState, IShaderObject* rootObject)
{
    return SLANG_E_NOT_IMPLEMENTED;
}

Result RayTracingCommandEncoder::dispatchRays(
    GfxIndex raygenShaderIndex,
    IShaderTable* shaderTable,
    GfxCount width,
    GfxCount height,
    GfxCount depth)
{
    return SLANG_E_NOT_IMPLEMENTED;
}

void RayTracingCommandEncoder::endEncoding() { }

} // namespace metal
} // namespace gfx