// Copyright (C) 2023 The Android Open Source Project
//
// 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 "ResourceTracker.h"
#include "VkEncoder.h"
#include "gfxstream_vk_private.h"

VkResult gfxstream_vk_CreateCommandPool(VkDevice device, const VkCommandPoolCreateInfo* pCreateInfo,
                                        const VkAllocationCallbacks* pAllocator,
                                        VkCommandPool* pCommandPool) {
    AEMU_SCOPED_TRACE("vkCreateCommandPool");
    VK_FROM_HANDLE(gfxstream_vk_device, gfxstream_device, device);
    VkResult result = (VkResult)0;
    struct gfxstream_vk_command_pool* gfxstream_pCommandPool =
        (gfxstream_vk_command_pool*)vk_zalloc2(&gfxstream_device->vk.alloc, pAllocator,
                                               sizeof(gfxstream_vk_command_pool), 8,
                                               VK_SYSTEM_ALLOCATION_SCOPE_OBJECT);
    result = gfxstream_pCommandPool ? VK_SUCCESS : VK_ERROR_OUT_OF_HOST_MEMORY;
    if (VK_SUCCESS == result) {
        result = vk_command_pool_init(&gfxstream_device->vk, &gfxstream_pCommandPool->vk,
                                      pCreateInfo, pAllocator);
    }
    if (VK_SUCCESS == result) {
        auto vkEnc = gfxstream::vk::ResourceTracker::getThreadLocalEncoder();
        result = vkEnc->vkCreateCommandPool(gfxstream_device->internal_object, pCreateInfo,
                                            pAllocator, &gfxstream_pCommandPool->internal_object,
                                            true /* do lock */);
    }
    *pCommandPool = gfxstream_vk_command_pool_to_handle(gfxstream_pCommandPool);
    return result;
}

void gfxstream_vk_DestroyCommandPool(VkDevice device, VkCommandPool commandPool,
                                     const VkAllocationCallbacks* pAllocator) {
    AEMU_SCOPED_TRACE("vkDestroyCommandPool");
    if (VK_NULL_HANDLE == commandPool) {
        return;
    }
    VK_FROM_HANDLE(gfxstream_vk_device, gfxstream_device, device);
    VK_FROM_HANDLE(gfxstream_vk_command_pool, gfxstream_commandPool, commandPool);
    {
        auto vkEnc = gfxstream::vk::ResourceTracker::getThreadLocalEncoder();
        vkEnc->vkDestroyCommandPool(gfxstream_device->internal_object,
                                    gfxstream_commandPool->internal_object, pAllocator,
                                    true /* do lock */);
    }
    vk_command_pool_finish(&gfxstream_commandPool->vk);
    vk_free(&gfxstream_commandPool->vk.alloc, gfxstream_commandPool);
}

VkResult gfxstream_vk_ResetCommandPool(VkDevice device, VkCommandPool commandPool,
                                       VkCommandPoolResetFlags flags) {
    AEMU_SCOPED_TRACE("vkResetCommandPool");
    VK_FROM_HANDLE(gfxstream_vk_device, gfxstream_device, device);
    VK_FROM_HANDLE(gfxstream_vk_command_pool, gfxstream_commandPool, commandPool);
    VkResult vkResetCommandPool_VkResult_return = (VkResult)0;
    {
        auto vkEnc = gfxstream::vk::ResourceTracker::getThreadLocalEncoder();
        vkResetCommandPool_VkResult_return = vkEnc->vkResetCommandPool(
            gfxstream_device->internal_object, gfxstream_commandPool->internal_object, flags,
            true /* do lock */);
        if (vkResetCommandPool_VkResult_return == VK_SUCCESS) {
            gfxstream::vk::ResourceTracker::get()->resetCommandPoolStagingInfo(
                gfxstream_commandPool->internal_object);
        }
    }
    return vkResetCommandPool_VkResult_return;
}

static VkResult vk_command_buffer_createOp(struct vk_command_pool*, struct vk_command_buffer**);
static void vk_command_buffer_resetOp(struct vk_command_buffer*, VkCommandBufferResetFlags);
static void vk_command_buffer_destroyOp(struct vk_command_buffer*);

static vk_command_buffer_ops gfxstream_vk_commandBufferOps = {
    .create = vk_command_buffer_createOp,
    .reset = vk_command_buffer_resetOp,
    .destroy = vk_command_buffer_destroyOp};

VkResult vk_command_buffer_createOp(struct vk_command_pool* commandPool,
                                    struct vk_command_buffer** pCommandBuffer) {
    VkResult result = VK_SUCCESS;
    struct gfxstream_vk_command_buffer* gfxstream_commandBuffer =
        (struct gfxstream_vk_command_buffer*)vk_zalloc(&commandPool->alloc,
                                                       sizeof(struct gfxstream_vk_command_buffer),
                                                       8, VK_SYSTEM_ALLOCATION_SCOPE_OBJECT);
    if (gfxstream_commandBuffer) {
        result =
            vk_command_buffer_init(commandPool, &gfxstream_commandBuffer->vk,
                                   &gfxstream_vk_commandBufferOps, VK_COMMAND_BUFFER_LEVEL_PRIMARY);
        if (VK_SUCCESS == result) {
            *pCommandBuffer = &gfxstream_commandBuffer->vk;
        }
    } else {
        result = VK_ERROR_OUT_OF_HOST_MEMORY;
    }
    return result;
}

void vk_command_buffer_resetOp(struct vk_command_buffer* commandBuffer,
                               VkCommandBufferResetFlags flags) {
    (void)flags;
    vk_command_buffer_reset(commandBuffer);
}

void vk_command_buffer_destroyOp(struct vk_command_buffer* commandBuffer) {
    vk_command_buffer_finish(commandBuffer);
    vk_free(&commandBuffer->pool->alloc, commandBuffer);
}

VkResult gfxstream_vk_AllocateCommandBuffers(VkDevice device,
                                             const VkCommandBufferAllocateInfo* pAllocateInfo,
                                             VkCommandBuffer* pCommandBuffers) {
    AEMU_SCOPED_TRACE("vkAllocateCommandBuffers");
    VK_FROM_HANDLE(gfxstream_vk_device, gfxstream_device, device);
    VK_FROM_HANDLE(gfxstream_vk_command_pool, gfxstream_commandPool, pAllocateInfo->commandPool);
    VkResult result = (VkResult)0;
    std::vector<gfxstream_vk_command_buffer*> gfxstream_commandBuffers(
        pAllocateInfo->commandBufferCount);
    for (uint32_t i = 0; i < pAllocateInfo->commandBufferCount; i++) {
        result = vk_command_buffer_createOp(&gfxstream_commandPool->vk,
                                            (vk_command_buffer**)&gfxstream_commandBuffers[i]);
        if (VK_SUCCESS == result) {
            gfxstream_commandBuffers[i]->vk.level = pAllocateInfo->level;
        } else {
            break;
        }
    }
    if (VK_SUCCESS == result) {
        // Create gfxstream-internal commandBuffer array
        std::vector<VkCommandBuffer> internal_objects(pAllocateInfo->commandBufferCount);
        auto vkEnc = gfxstream::vk::ResourceTracker::getThreadLocalEncoder();
        auto resources = gfxstream::vk::ResourceTracker::get();
        VkCommandBufferAllocateInfo internal_allocateInfo;
        internal_allocateInfo = *pAllocateInfo;
        internal_allocateInfo.commandPool = gfxstream_commandPool->internal_object;
        result = resources->on_vkAllocateCommandBuffers(
            vkEnc, VK_SUCCESS, gfxstream_device->internal_object, &internal_allocateInfo,
            internal_objects.data());
        if (result == VK_SUCCESS) {
            gfxstream::vk::ResourceTracker::get()->addToCommandPool(
                gfxstream_commandPool->internal_object, pAllocateInfo->commandBufferCount,
                internal_objects.data());
            for (uint32_t i = 0; i < (uint32_t)internal_objects.size(); i++) {
                gfxstream_commandBuffers[i]->internal_object = internal_objects[i];
                // TODO: Also vk_command_buffer_init() on every mesa command buffer?
                pCommandBuffers[i] =
                    gfxstream_vk_command_buffer_to_handle(gfxstream_commandBuffers[i]);
            }
        }
    }
    return result;
}

void gfxstream_vk_FreeCommandBuffers(VkDevice device, VkCommandPool commandPool,
                                     uint32_t commandBufferCount,
                                     const VkCommandBuffer* pCommandBuffers) {
    AEMU_SCOPED_TRACE("vkFreeCommandBuffers");
    VK_FROM_HANDLE(gfxstream_vk_device, gfxstream_device, device);
    VK_FROM_HANDLE(gfxstream_vk_command_pool, gfxstream_commandPool, commandPool);
    {
        // Set up internal commandBuffer array for gfxstream-internal call
        std::vector<VkCommandBuffer> internal_objects(commandBufferCount);
        for (uint32_t i = 0; i < commandBufferCount; i++) {
            VK_FROM_HANDLE(gfxstream_vk_command_buffer, gfxstream_commandBuffer,
                           pCommandBuffers[i]);
            internal_objects[i] = gfxstream_commandBuffer->internal_object;
        }
        auto vkEnc = gfxstream::vk::ResourceTracker::getThreadLocalEncoder();
        vkEnc->vkFreeCommandBuffers(gfxstream_device->internal_object,
                                    gfxstream_commandPool->internal_object, commandBufferCount,
                                    internal_objects.data(), true /* do lock */);
    }
    for (uint32_t i = 0; i < commandBufferCount; i++) {
        VK_FROM_HANDLE(gfxstream_vk_command_buffer, gfxstream_commandBuffer, pCommandBuffers[i]);
        vk_command_buffer_destroyOp(&gfxstream_commandBuffer->vk);
    }
}