// Copyright 2020 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/metal/CommandRecordingContext.h"

#include "common/Assert.h"

namespace dawn_native::metal {

    CommandRecordingContext::CommandRecordingContext() = default;

    CommandRecordingContext::~CommandRecordingContext() {
        // Commands must be acquired.
        ASSERT(mCommands == nullptr);
    }

    id<MTLCommandBuffer> CommandRecordingContext::GetCommands() {
        return mCommands.Get();
    }

    void CommandRecordingContext::MarkUsed() {
        mUsed = true;
    }
    bool CommandRecordingContext::WasUsed() const {
        return mUsed;
    }

    MaybeError CommandRecordingContext::PrepareNextCommandBuffer(id<MTLCommandQueue> queue) {
        ASSERT(mCommands == nil);
        ASSERT(!mUsed);

        // The MTLCommandBuffer will be autoreleased by default.
        // The autorelease pool may drain before the command buffer is submitted. Retain so it stays
        // alive.
        mCommands = AcquireNSPRef([[queue commandBuffer] retain]);
        if (mCommands == nil) {
            return DAWN_INTERNAL_ERROR("Failed to allocate an MTLCommandBuffer");
        }

        return {};
    }

    NSPRef<id<MTLCommandBuffer>> CommandRecordingContext::AcquireCommands() {
        // A blit encoder can be left open from WriteBuffer, make sure we close it.
        if (mCommands != nullptr) {
            EndBlit();
        }

        ASSERT(!mInEncoder);
        mUsed = false;
        return std::move(mCommands);
    }

    id<MTLBlitCommandEncoder> CommandRecordingContext::EnsureBlit() {
        ASSERT(mCommands != nullptr);

        if (mBlit == nullptr) {
            ASSERT(!mInEncoder);
            mInEncoder = true;

            // The encoder is created autoreleased. Retain it to avoid the autoreleasepool from
            // draining from under us.
            mBlit.Acquire([[*mCommands blitCommandEncoder] retain]);
        }
        return mBlit.Get();
    }

    void CommandRecordingContext::EndBlit() {
        ASSERT(mCommands != nullptr);

        if (mBlit != nullptr) {
            [*mBlit endEncoding];
            mBlit = nullptr;
            mInEncoder = false;
        }
    }

    id<MTLComputeCommandEncoder> CommandRecordingContext::BeginCompute() {
        ASSERT(mCommands != nullptr);
        ASSERT(mCompute == nullptr);
        ASSERT(!mInEncoder);

        mInEncoder = true;
        // The encoder is created autoreleased. Retain it to avoid the autoreleasepool from
        // draining from under us.
        mCompute.Acquire([[*mCommands computeCommandEncoder] retain]);
        return mCompute.Get();
    }

    void CommandRecordingContext::EndCompute() {
        ASSERT(mCommands != nullptr);
        ASSERT(mCompute != nullptr);

        [*mCompute endEncoding];
        mCompute = nullptr;
        mInEncoder = false;
    }

    id<MTLRenderCommandEncoder> CommandRecordingContext::BeginRender(
        MTLRenderPassDescriptor* descriptor) {
        ASSERT(mCommands != nullptr);
        ASSERT(mRender == nullptr);
        ASSERT(!mInEncoder);

        mInEncoder = true;
        // The encoder is created autoreleased. Retain it to avoid the autoreleasepool from
        // draining from under us.
        mRender.Acquire([[*mCommands renderCommandEncoderWithDescriptor:descriptor] retain]);
        return mRender.Get();
    }

    void CommandRecordingContext::EndRender() {
        ASSERT(mCommands != nullptr);
        ASSERT(mRender != nullptr);

        [*mRender endEncoding];
        mRender = nullptr;
        mInEncoder = false;
    }

}  // namespace dawn_native::metal