// 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 "VulkanHandleMapping.h"
#include "VulkanHandles.h"
#include "aemu/base/containers/EntityManager.h"
#include "aemu/base/HealthMonitor.h"
#include "aemu/base/files/Stream.h"
#include "common/goldfish_vk_marshaling.h"
#include "utils/GfxApiLogger.h"

namespace gfxstream {
namespace vk {

// A class that captures all important data structures for
// reconstructing a Vulkan system state via trimmed API record and replay.
class VkReconstruction {
   public:
    VkReconstruction();

    void save(android::base::Stream* stream);
    void load(android::base::Stream* stream, emugl::GfxApiLogger& gfxLogger,
              emugl::HealthMonitor<>* healthMonitor);

    struct ApiInfo {
        // Fast
        uint32_t opCode;
        std::vector<uint8_t> trace;
        size_t traceBytes = 0;
        // Book-keeping for which handles were created by this API
        std::vector<uint64_t> createdHandles;
    };

    using ApiTrace = android::base::EntityManager<32, 16, 16, ApiInfo>;
    using ApiHandle = ApiTrace::EntityHandle;

    enum HandleState { BEGIN = 0, CREATED = 0, BOUND_MEMORY = 1, HANDLE_STATE_COUNT };

    typedef std::pair<uint64_t, HandleState> HandleWithState;
    struct HandleWithStateHash {
        inline size_t operator()(const HandleWithState& v) const {
            std::hash<uint64_t> int_hasher;
            return int_hasher(v.first) ^ int_hasher(v.second);
        }
    };

    struct HandleReconstruction {
        std::vector<ApiHandle> apiRefs;
        std::unordered_set<HandleWithState, HandleWithStateHash> childHandles;
        std::vector<HandleWithState> parentHandles;
    };

    struct HandleWithStateReconstruction {
        std::vector<HandleReconstruction> states =
            std::vector<HandleReconstruction>(HANDLE_STATE_COUNT);
        bool delayed_destroy = false;
        bool destroying = false;
    };

    using HandleWithStateReconstructions =
        android::base::UnpackedComponentManager<32, 16, 16, HandleWithStateReconstruction>;

    struct HandleModification {
        std::vector<ApiHandle> apiRefs;
        uint32_t order = 0;
    };

    using HandleModifications =
        android::base::UnpackedComponentManager<32, 16, 16, HandleModification>;

    ApiHandle createApiInfo();
    void destroyApiInfo(ApiHandle h);

    void removeHandleFromApiInfo(ApiHandle h, uint64_t toRemove);

    ApiInfo* getApiInfo(ApiHandle h);

    void setApiTrace(ApiInfo* apiInfo, uint32_t opcode, 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 forEachHandleAddApi(const uint64_t* toProcess, uint32_t count, uint64_t apiHandle,
                             HandleState state = CREATED);
    void forEachHandleDeleteApi(const uint64_t* toProcess, uint32_t count);

    void addHandleDependency(const uint64_t* handles, uint32_t count, uint64_t parentHandle,
                             HandleState childState = CREATED, HandleState parentState = CREATED);

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

    void forEachHandleAddModifyApi(const uint64_t* toProcess, uint32_t count, uint64_t apiHandle);

    void forEachHandleClearModifyApi(const uint64_t* toProcess, uint32_t count);

    void setModifiedHandlesForApi(uint64_t apiHandle, const uint64_t* modified, uint32_t count);

    // 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);

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

    ApiTrace mApiTrace;

    HandleWithStateReconstructions mHandleReconstructions;
    HandleModifications mHandleModifications;

    std::vector<uint64_t> mExtraHandlesForNextApi;

    std::vector<uint8_t> mLoadedTrace;
};

}  // namespace vk
}  // namespace gfxstream