// Copyright 2017 The Dawn Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "dawn/native/d3d12/BindGroupLayoutD3D12.h"

#include <utility>

#include "dawn/common/BitSetIterator.h"
#include "dawn/native/d3d12/DeviceD3D12.h"
#include "dawn/native/d3d12/SamplerHeapCacheD3D12.h"
#include "dawn/native/d3d12/StagingDescriptorAllocatorD3D12.h"

namespace dawn::native::d3d12 {
    namespace {
        D3D12_DESCRIPTOR_RANGE_TYPE WGPUBindingInfoToDescriptorRangeType(
            const BindingInfo& bindingInfo) {
            switch (bindingInfo.bindingType) {
                case BindingInfoType::Buffer:
                    switch (bindingInfo.buffer.type) {
                        case wgpu::BufferBindingType::Uniform:
                            return D3D12_DESCRIPTOR_RANGE_TYPE_CBV;
                        case wgpu::BufferBindingType::Storage:
                        case kInternalStorageBufferBinding:
                            return D3D12_DESCRIPTOR_RANGE_TYPE_UAV;
                        case wgpu::BufferBindingType::ReadOnlyStorage:
                            return D3D12_DESCRIPTOR_RANGE_TYPE_SRV;
                        case wgpu::BufferBindingType::Undefined:
                            UNREACHABLE();
                    }

                case BindingInfoType::Sampler:
                    return D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER;

                case BindingInfoType::Texture:
                case BindingInfoType::ExternalTexture:
                    return D3D12_DESCRIPTOR_RANGE_TYPE_SRV;

                case BindingInfoType::StorageTexture:
                    switch (bindingInfo.storageTexture.access) {
                        case wgpu::StorageTextureAccess::WriteOnly:
                            return D3D12_DESCRIPTOR_RANGE_TYPE_UAV;
                        case wgpu::StorageTextureAccess::Undefined:
                            UNREACHABLE();
                    }
            }
        }
    }  // anonymous namespace

    // static
    Ref<BindGroupLayout> BindGroupLayout::Create(
        Device* device,
        const BindGroupLayoutDescriptor* descriptor,
        PipelineCompatibilityToken pipelineCompatibilityToken) {
        return AcquireRef(new BindGroupLayout(device, descriptor, pipelineCompatibilityToken));
    }

    BindGroupLayout::BindGroupLayout(Device* device,
                                     const BindGroupLayoutDescriptor* descriptor,
                                     PipelineCompatibilityToken pipelineCompatibilityToken)
        : BindGroupLayoutBase(device, descriptor, pipelineCompatibilityToken),
          mDescriptorHeapOffsets(GetBindingCount()),
          mShaderRegisters(GetBindingCount()),
          mCbvUavSrvDescriptorCount(0),
          mSamplerDescriptorCount(0),
          mBindGroupAllocator(MakeFrontendBindGroupAllocator<BindGroup>(4096)) {
        for (BindingIndex bindingIndex{0}; bindingIndex < GetBindingCount(); ++bindingIndex) {
            const BindingInfo& bindingInfo = GetBindingInfo(bindingIndex);

            D3D12_DESCRIPTOR_RANGE_TYPE descriptorRangeType =
                WGPUBindingInfoToDescriptorRangeType(bindingInfo);
            mShaderRegisters[bindingIndex] = uint32_t(bindingInfo.binding);

            // For dynamic resources, Dawn uses root descriptor in D3D12 backend. So there is no
            // need to allocate the descriptor from descriptor heap or create descriptor ranges.
            if (bindingIndex < GetDynamicBufferCount()) {
                continue;
            }
            ASSERT(!bindingInfo.buffer.hasDynamicOffset);

            mDescriptorHeapOffsets[bindingIndex] =
                descriptorRangeType == D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER
                    ? mSamplerDescriptorCount++
                    : mCbvUavSrvDescriptorCount++;

            D3D12_DESCRIPTOR_RANGE range;
            range.RangeType = descriptorRangeType;
            range.NumDescriptors = 1;
            range.BaseShaderRegister = GetShaderRegister(bindingIndex);
            range.RegisterSpace = kRegisterSpacePlaceholder;
            range.OffsetInDescriptorsFromTableStart = D3D12_DESCRIPTOR_RANGE_OFFSET_APPEND;

            std::vector<D3D12_DESCRIPTOR_RANGE>& descriptorRanges =
                descriptorRangeType == D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER
                    ? mSamplerDescriptorRanges
                    : mCbvUavSrvDescriptorRanges;

            // Try to join this range with the previous one, if the current range is a continuation
            // of the previous. This is possible because the binding infos in the base type are
            // sorted.
            if (descriptorRanges.size() >= 2) {
                D3D12_DESCRIPTOR_RANGE& previous = descriptorRanges.back();
                if (previous.RangeType == range.RangeType &&
                    previous.BaseShaderRegister + previous.NumDescriptors ==
                        range.BaseShaderRegister) {
                    previous.NumDescriptors += range.NumDescriptors;
                    continue;
                }
            }

            descriptorRanges.push_back(range);
        }

        mViewAllocator = device->GetViewStagingDescriptorAllocator(GetCbvUavSrvDescriptorCount());
        mSamplerAllocator =
            device->GetSamplerStagingDescriptorAllocator(GetSamplerDescriptorCount());
    }

    ResultOrError<Ref<BindGroup>> BindGroupLayout::AllocateBindGroup(
        Device* device,
        const BindGroupDescriptor* descriptor) {
        uint32_t viewSizeIncrement = 0;
        CPUDescriptorHeapAllocation viewAllocation;
        if (GetCbvUavSrvDescriptorCount() > 0) {
            DAWN_TRY_ASSIGN(viewAllocation, mViewAllocator->AllocateCPUDescriptors());
            viewSizeIncrement = mViewAllocator->GetSizeIncrement();
        }

        Ref<BindGroup> bindGroup = AcquireRef<BindGroup>(
            mBindGroupAllocator.Allocate(device, descriptor, viewSizeIncrement, viewAllocation));

        if (GetSamplerDescriptorCount() > 0) {
            Ref<SamplerHeapCacheEntry> samplerHeapCacheEntry;
            DAWN_TRY_ASSIGN(samplerHeapCacheEntry, device->GetSamplerHeapCache()->GetOrCreate(
                                                       bindGroup.Get(), mSamplerAllocator));
            bindGroup->SetSamplerAllocationEntry(std::move(samplerHeapCacheEntry));
        }

        return bindGroup;
    }

    void BindGroupLayout::DeallocateBindGroup(BindGroup* bindGroup,
                                              CPUDescriptorHeapAllocation* viewAllocation) {
        if (viewAllocation->IsValid()) {
            mViewAllocator->Deallocate(viewAllocation);
        }

        mBindGroupAllocator.Deallocate(bindGroup);
    }

    ityp::span<BindingIndex, const uint32_t> BindGroupLayout::GetDescriptorHeapOffsets() const {
        return {mDescriptorHeapOffsets.data(), mDescriptorHeapOffsets.size()};
    }

    uint32_t BindGroupLayout::GetShaderRegister(BindingIndex bindingIndex) const {
        return mShaderRegisters[bindingIndex];
    }

    uint32_t BindGroupLayout::GetCbvUavSrvDescriptorCount() const {
        return mCbvUavSrvDescriptorCount;
    }

    uint32_t BindGroupLayout::GetSamplerDescriptorCount() const {
        return mSamplerDescriptorCount;
    }

    const std::vector<D3D12_DESCRIPTOR_RANGE>& BindGroupLayout::GetCbvUavSrvDescriptorRanges()
        const {
        return mCbvUavSrvDescriptorRanges;
    }

    const std::vector<D3D12_DESCRIPTOR_RANGE>& BindGroupLayout::GetSamplerDescriptorRanges() const {
        return mSamplerDescriptorRanges;
    }

}  // namespace dawn::native::d3d12