/*
 * Copyright 2021 Google Inc.
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#ifndef skgpu_graphite_DrawBufferManager_DEFINED
#define skgpu_graphite_DrawBufferManager_DEFINED

#include "include/core/SkRefCnt.h"
#include "src/gpu/BufferWriter.h"
#include "src/gpu/graphite/DrawTypes.h"
#include "src/gpu/graphite/ResourceTypes.h"

#include <unordered_map>
#include <vector>

namespace skgpu::graphite {

class Buffer;
class Recording;
class ResourceProvider;

class DrawBufferManager {
public:
    DrawBufferManager(ResourceProvider*, size_t uniformStartAlignment, size_t ssboStartAlignment);
    ~DrawBufferManager();

    std::tuple<VertexWriter, BindBufferInfo> getVertexWriter(size_t requiredBytes);
    std::tuple<IndexWriter, BindBufferInfo> getIndexWriter(size_t requiredBytes);
    std::tuple<UniformWriter, BindBufferInfo> getUniformWriter(size_t requiredBytes);
    std::tuple<UniformWriter, BindBufferInfo> getSsboWriter(size_t requiredBytes);

    // Returns the last 'unusedBytes' from the last call to getVertexWriter(). Assumes that
    // 'unusedBytes' is less than the 'requiredBytes' to the original allocation.
    void returnVertexBytes(size_t unusedBytes);

    size_t alignUniformBlockSize(size_t dataSize) {
        return SkAlignTo(dataSize, fUniformStartAlignment);
    }

    // Get the shared static buffer filled with contents computed by the InitializeBufferFn.
    // Both InitializeBufferFn and BufferSizeFn should be static functions since their addresses are
    // used in the static buffer's internal unique keys.
    // TODO(skbug.com/13059): getStaticBuffer is only temporary until RenderSteps can create buffers
    // during initialization of the Context. Only use with RenderSteps that will be updated
    // to the new initialization pattern in the near future.
    using BufferSizeFn = size_t(*)();
    using InitializeBufferFn = void(*)(skgpu::VertexWriter, size_t bufferSize);
    BindBufferInfo getStaticBuffer(BufferType type,
                                   InitializeBufferFn,
                                   BufferSizeFn);

    // Finalizes all buffers and transfers ownership of them to a Recording.
    void transferToRecording(Recording*);

private:
    ResourceProvider* fResourceProvider;

    sk_sp<Buffer> fCurrentVertexBuffer;
    size_t fVertexOffset = 0;

    sk_sp<Buffer> fCurrentIndexBuffer;
    size_t fIndexOffset = 0;

    sk_sp<Buffer> fCurrentUniformBuffer;
    size_t fUniformOffset = 0;

    sk_sp<Buffer> fCurrentStorageBuffer;
    size_t fSsboOffset = 0;

    const size_t fUniformStartAlignment;
    const size_t fSsboStartAlignment;

    std::vector<sk_sp<Buffer>> fUsedBuffers;
    // TODO(skbug.com/13059): This is likely not the final location for static buffers, but makes it
    // convenient to maintain ownership and call trackResources() on the CommandBuffer.
    std::unordered_map<uintptr_t, sk_sp<Buffer>> fStaticBuffers;
};

} // namespace skgpu::graphite

#endif // skgpu_graphite_DrawBufferManager_DEFINED