// Copyright (C) 2019 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.

#pragma once

#include <atomic>
#include <map>
#include <memory>
#include <mutex>
#include <set>

#include "DependencyGraph.h"
#include "VkSnapshotHandles.h"
#include "VulkanHandleMapping.h"
#include "VulkanHandles.h"
#include "common/goldfish_vk_marshaling.h"
#include "gfxstream/HealthMonitor.h"
#include "gfxstream/containers/EntityManager.h"
#include "gfxstream/host/GfxApiLogger.h"
#include "render-utils/stream.h"

namespace gfxstream {
namespace vk {

class VkReconstruction {
   public:
    VkReconstruction();

    void clear();

    void saveReplayBuffers(gfxstream::Stream* stream);
    static void loadReplayBuffers(gfxstream::Stream* stream,
                                  std::vector<uint64_t>* outHandleBuffer,
                                  std::vector<uint8_t>* outDecoderBuffer);

    enum HandleState { CREATED = 0 };

    VkSnapshotApiCallHandle createApiCallInfo();
    void destroyApiCallInfoIfUnused(VkSnapshotApiCallHandle apiCallHandle);

    void removeHandleFromApiInfo(VkSnapshotApiCallHandle apiCallHandle, uint64_t toRemove);

    void setApiTrace(VkSnapshotApiCallHandle apiCallHandle, const uint8_t* traceBegin, size_t traceBytes);

    void dump();

    void addHandles(const uint64_t* toAdd, uint32_t count);
    void removeHandles(const uint64_t* toRemove, uint32_t count, bool recursive = true);

    void removeGrandChildren(const uint64_t handle);
    void removeDescendantsOfHandle(const uint64_t handle);

    void forEachHandleAddApi(const uint64_t* toProcess, uint32_t count,
                             VkSnapshotApiCallHandle apiCallHandle, HandleState state = CREATED);

    void addApiCallDependencyOnVkObject(VkSnapshotApiCallHandle apiCallHandle, VkObjectHandle object);

    void addHandleDependenciesForApiCallDependencies(VkSnapshotApiCallHandle apiCallHandle, VkObjectHandle child);

    void addHandleDependency(const VkObjectHandle* childHandles, uint32_t childHandlesCount,
                             VkObjectHandle parentHandle,
                             HandleState childState = CREATED, HandleState parentState = CREATED);

    void setCreatedHandlesForApi(VkSnapshotApiCallHandle handle , const uint64_t* created,
                                 uint32_t count);

    void forEachHandleClearModifyApi(const uint64_t* toProcess, uint32_t count) {}
    void forEachHandleAddModifyApi(const uint64_t* toProcess, uint32_t count, uint64_t api) {}

    // Used by on_vkCreateDescriptorPool.
    //
    // Snapshot keeps track of all the boxed handles created by each function. By default
    // the generated code assumes no extra internal boxed handles are generated by
    // VkDecoderGlobalState. But this is not the case for on_vkCreateDescriptorPool.
    // Thus we add an extra API to VkReconstruction, which gives it the list of all the
    // extra boxed handles.
    //
    // Implementation-wise it is a bit tricky. The regular workflow looks like:
    //
    // on_vkCreateDescriptorPool(... pDescriptorPool)
    // ...
    // mReconstruction.setCreatedHandlesForApi(OP_vkCreateDescriptorPool, pDescriptorPool);
    //
    // It is not easy to directly tell mReconstruction that OP_vkCreateDescriptorPool created
    // extra handles. Instead, we add an API to VkReconstruction to cache the extra handles.
    // Next time setCreatedHandlesForApi is called, it will check the cached handles and
    // add them to OP_vkCreateDescriptorPool.
    void createExtraHandlesForNextApi(const uint64_t* created, uint32_t count);

    void addOrderedBoxedHandlesCreatedByCall(VkSnapshotApiCallHandle handle,
                                             const VkObjectHandle* boxedHandles,
                                             uint32_t boxedHandlesCount);

   private:
    struct VkSnapshotApiCallInfo {
        VkSnapshotApiCallHandle handle = kInvalidSnapshotApiCallHandle;

        // Raw packet from VkDecoder.
        std::vector<uint8_t> packet;

        // Book-keeping for which handles were created by this API
        std::vector<uint64_t> createdHandles;
        std::vector<uint64_t> depends;

        // Extra boxed handles created for this API call that are not identifiable
        // solely from the API parameters itself. For example, the extra boxed `VkQueue`s
        // that are created during `vkCreateDevice()` can not be identified from the
        // parameters to `vkCreateDevice()`.
        //
        // TODO: remove this and require that all of the `new_boxed_*()` take a
        // `VkSnapshotApiCallInfo` as an argument so the creation order of the boxed
        // handles in `createdHandles` is guaranteed to match the replay order. For now,
        // this relies on careful manual ordering.
        std::vector<uint64_t> extraCreatedHandles;
    };

    std::vector<uint64_t> getOrderedUniqueModifyApis() const;

    using VkSnapshotApiCallManager = gfxstream::base::EntityManager<32, 16, 16, VkSnapshotApiCallInfo>;
    VkSnapshotApiCallManager mApiCallManager;

    std::vector<uint8_t> mLoadedTrace;

    DependencyGraph mGraph;
};

}  // namespace vk
}  // namespace gfxstream