/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#ifndef MOZILLA_GFX_RemoteTextureMap_H
#define MOZILLA_GFX_RemoteTextureMap_H

#include <deque>
#include <functional>
#include <list>
#include <map>
#include <memory>
#include <queue>
#include <unordered_set>
#include <utility>

#include "mozilla/gfx/Point.h"  // for IntSize
#include "mozilla/gfx/Types.h"  // for SurfaceFormat
#include "mozilla/ipc/Shmem.h"
#include "mozilla/layers/CompositorTypes.h"  // for TextureFlags, etc
#include "mozilla/layers/LayersSurfaces.h"   // for SurfaceDescriptor
#include "mozilla/layers/TextureHost.h"
#include "mozilla/Monitor.h"
#include "mozilla/StaticPtr.h"
#include "mozilla/UniquePtr.h"
#include "mozilla/webrender/WebRenderTypes.h"

class nsISerialEventTarget;

namespace mozilla {

namespace ipc {
class IProtocol;
}

namespace gl {
class SharedSurface;
}

namespace webgpu {
class SharedTexture;
}

namespace layers {

class CompositableHost;
class RemoteTextureHostWrapper;
class TextureData;
class TextureHost;

struct RemoteTextureInfo {
  RemoteTextureInfo(const RemoteTextureId aTextureId,
                    const RemoteTextureOwnerId aOwnerId,
                    const base::ProcessId aForPid)
      : mTextureId(aTextureId),
        mOwnerId(aOwnerId),
        mForPid(aForPid),
        mWaitForRemoteTextureOwner(false) {}

  RemoteTextureInfo(const RemoteTextureId aTextureId,
                    const RemoteTextureOwnerId aOwnerId,
                    const base::ProcessId aForPid,
                    const bool aWaitForRemoteTextureOwner)
      : mTextureId(aTextureId),
        mOwnerId(aOwnerId),
        mForPid(aForPid),
        mWaitForRemoteTextureOwner(aWaitForRemoteTextureOwner) {}

  const RemoteTextureId mTextureId;
  const RemoteTextureOwnerId mOwnerId;
  const base::ProcessId mForPid;
  const bool mWaitForRemoteTextureOwner;
};

struct RemoteTextureInfoList {
  std::queue<RemoteTextureInfo> mList;
};

class SharedResourceWrapper {
 public:
  enum class Tag { SharedSurface, SharedTexture };
  const Tag mTag;

  static UniquePtr<SharedResourceWrapper> SharedSurface(
      const std::shared_ptr<gl::SharedSurface>& aSharedSurface) {
    return MakeUnique<SharedResourceWrapper>(Tag::SharedSurface,
                                             aSharedSurface);
  }

  static UniquePtr<SharedResourceWrapper> SharedTexture(
      const std::shared_ptr<webgpu::SharedTexture>& aSharedTexture) {
    return MakeUnique<SharedResourceWrapper>(Tag::SharedTexture,
                                             aSharedTexture);
  }

  SharedResourceWrapper(
      const Tag aTag, const std::shared_ptr<gl::SharedSurface>& aSharedSurface)
      : mTag(aTag), mSharedSurface(aSharedSurface) {
    MOZ_ASSERT(mTag == Tag::SharedSurface);
  }
  SharedResourceWrapper(
      const Tag aTag,
      const std::shared_ptr<webgpu::SharedTexture>& aSharedTexture)
      : mTag(aTag), mSharedTexture(aSharedTexture) {
    MOZ_ASSERT(mTag == Tag::SharedTexture);
  }

  const std::shared_ptr<gl::SharedSurface> mSharedSurface;
  const std::shared_ptr<webgpu::SharedTexture> mSharedTexture;

  std::shared_ptr<gl::SharedSurface> SharedSurface() {
    if (mTag == Tag::SharedSurface) {
      return mSharedSurface;
    }
    MOZ_ASSERT_UNREACHABLE("unexpected to be called");
    return nullptr;
  }

  std::shared_ptr<webgpu::SharedTexture> SharedTexture() {
    if (mTag == Tag::SharedTexture) {
      return mSharedTexture;
    }
    MOZ_ASSERT_UNREACHABLE("unexpected to be called");
    return nullptr;
  }
};

class RemoteTextureRecycleBin final {
 public:
  NS_INLINE_DECL_THREADSAFE_REFCOUNTING(RemoteTextureRecycleBin)

  explicit RemoteTextureRecycleBin(bool aIsShared);

 private:
  friend class RemoteTextureMap;

  ~RemoteTextureRecycleBin();

  struct RecycledTextureHolder {
    gfx::IntSize mSize;
    gfx::SurfaceFormat mFormat = gfx::SurfaceFormat::UNKNOWN;
    SurfaceDescriptor::Type mType = SurfaceDescriptor::Tnull_t;
    UniquePtr<TextureData> mTextureData;
    UniquePtr<SharedResourceWrapper> mResourceWrapper;
  };

  bool mIsShared = false;
  std::list<RecycledTextureHolder> mRecycledTextures;
};

/**
 * RemoteTextureTxnScheduler manages dependencies on transaction numbers for a
 * given top-level protocol ("type"). It expects that transaction numbers are
 * all sequenced and comparable for this top-level protocol. Dependencies may
 * then be issued on a given future transaction number. Clients must notify the
 * scheduler when transactions are completed, so that any dependencies can be
 * cleared as the transaction number advances. Generally, transaction numbers
 * will be generated by a FwdTransactionCounter on a top-level protocol child,
 * and there should be a RemoteTextureTxnScheduler instantiated on the top-level
 * protocol parent that corresponds to the same protocol lifetime as the child.
 * To ease sharing in sub-protocols, RemoteTextureTxnScheduler is ref-counted
 * and may be multiply-registered by the various sub-protocols that derive from
 * a given top-level protocol, without having to directly refer to the original
 * top-level protocol.
 */
class RemoteTextureTxnScheduler final {
 public:
  NS_INLINE_DECL_THREADSAFE_REFCOUNTING(RemoteTextureTxnScheduler)

  static already_AddRefed<RemoteTextureTxnScheduler> Create(
      mozilla::ipc::IProtocol* aProtocol);

  void NotifyTxn(RemoteTextureTxnId aTxnId);

 private:
  friend class RemoteTextureMap;

  RemoteTextureTxnScheduler(base::ProcessId aForPid, RemoteTextureTxnType aType)
      : mForPid(aForPid), mType(aType) {}
  ~RemoteTextureTxnScheduler();

  bool WaitForTxn(const MonitorAutoLock& aProofOfLock,
                  RemoteTextureOwnerId aOwnerId, RemoteTextureTxnId aTxnId);

  struct Wait {
    RemoteTextureOwnerId mOwnerId;
    RemoteTextureTxnId mTxnId;

    friend bool operator<(RemoteTextureTxnId aTxnId, const Wait& aWait) {
      return aTxnId < aWait.mTxnId;
    }
  };

  base::ProcessId mForPid = base::kInvalidProcessId;
  RemoteTextureTxnType mType = 0;
  RemoteTextureTxnId mLastTxnId = 0;
  std::deque<Wait> mWaits;
};

typedef std::unordered_set<RemoteTextureOwnerId, RemoteTextureOwnerId::HashFn>
    RemoteTextureOwnerIdSet;

/**
 * A class provides API for remote texture owners.
 */
class RemoteTextureOwnerClient final {
 public:
  NS_INLINE_DECL_THREADSAFE_REFCOUNTING(RemoteTextureOwnerClient)

  explicit RemoteTextureOwnerClient(const base::ProcessId aForPid);

  bool IsRegistered(const RemoteTextureOwnerId aOwnerId);
  void RegisterTextureOwner(const RemoteTextureOwnerId aOwnerId,
                            bool aSharedRecycling = false);
  void UnregisterTextureOwner(const RemoteTextureOwnerId aOwnerId);
  void UnregisterAllTextureOwners();
  bool WaitForTxn(const RemoteTextureOwnerId aOwnerId,
                  RemoteTextureTxnType aTxnType, RemoteTextureTxnId aTxnId);
  void ClearRecycledTextures();
  void NotifyContextLost(const RemoteTextureOwnerIdSet* aOwnerIds = nullptr);
  void NotifyContextRestored(
      const RemoteTextureOwnerIdSet* aOwnerIds = nullptr);
  void PushTexture(const RemoteTextureId aTextureId,
                   const RemoteTextureOwnerId aOwnerId,
                   UniquePtr<TextureData>&& aTextureData);
  void PushTexture(const RemoteTextureId aTextureId,
                   const RemoteTextureOwnerId aOwnerId,
                   const std::shared_ptr<gl::SharedSurface>& aSharedSurface,
                   const gfx::IntSize& aSize, gfx::SurfaceFormat aFormat,
                   const SurfaceDescriptor& aDesc);
  void PushTexture(const RemoteTextureId aTextureId,
                   const RemoteTextureOwnerId aOwnerId,
                   const std::shared_ptr<webgpu::SharedTexture>& aSharedTexture,
                   const gfx::IntSize& aSize, gfx::SurfaceFormat aFormat,
                   const SurfaceDescriptor& aDesc);
  void PushDummyTexture(const RemoteTextureId aTextureId,
                        const RemoteTextureOwnerId aOwnerId);
  void GetLatestBufferSnapshot(const RemoteTextureOwnerId aOwnerId,
                               const mozilla::ipc::Shmem& aDestShmem,
                               const gfx::IntSize& aDestSize,
                               size_t aDestStride);
  UniquePtr<TextureData> GetRecycledTextureData(
      const gfx::IntSize& aSize, gfx::SurfaceFormat aFormat,
      TextureType aTextureType,
      RemoteTextureOwnerId aOwnerId = RemoteTextureOwnerId());
  UniquePtr<TextureData> CreateOrRecycleBufferTextureData(
      const gfx::IntSize& aSize, gfx::SurfaceFormat aFormat,
      RemoteTextureOwnerId aOwnerId = RemoteTextureOwnerId());
  std::shared_ptr<gl::SharedSurface> GetRecycledSharedSurface(
      const gfx::IntSize& aSize, gfx::SurfaceFormat aFormat,
      SurfaceDescriptor::Type aType,
      RemoteTextureOwnerId aOwnerId = RemoteTextureOwnerId());
  std::shared_ptr<webgpu::SharedTexture> GetRecycledSharedTexture(
      const gfx::IntSize& aSize, gfx::SurfaceFormat aFormat,
      SurfaceDescriptor::Type aType,
      RemoteTextureOwnerId aOwnerId = RemoteTextureOwnerId());

  const base::ProcessId mForPid;

 protected:
  ~RemoteTextureOwnerClient();

  RemoteTextureOwnerIdSet mOwnerIds;
  RefPtr<RemoteTextureRecycleBin> mSharedRecycleBin;
};

/**
 * A class to map RemoteTextureId to remote texture(TextureHost).
 * Remote textures are provided by texture owner.
 */
class RemoteTextureMap {
 public:
  static void Init();
  static void Shutdown();
  static RemoteTextureMap* Get() { return sInstance; }

  RemoteTextureMap();
  ~RemoteTextureMap();

  // Push remote texture data and gl::SharedSurface from texture owner.
  // The texture data is used for creating TextureHost.
  // gl::SharedSurface is pushed only when the surface needs to be kept alive
  // during TextureHost usage. The texture data and the surface might be
  // recycled when TextureHost is destroyed.
  void PushTexture(const RemoteTextureId aTextureId,
                   const RemoteTextureOwnerId aOwnerId,
                   const base::ProcessId aForPid,
                   UniquePtr<TextureData>&& aTextureData,
                   RefPtr<TextureHost>& aTextureHost,
                   UniquePtr<SharedResourceWrapper>&& aResourceWrapper);

  // Remove waiting texture that will not be used.
  bool RemoveTexture(const RemoteTextureId aTextureId,
                     const RemoteTextureOwnerId aOwnerId,
                     const base::ProcessId aForPid);

  void GetLatestBufferSnapshot(const RemoteTextureOwnerId aOwnerId,
                               const base::ProcessId aForPid,
                               const mozilla::ipc::Shmem& aDestShmem,
                               const gfx::IntSize& aDestSize,
                               size_t aDestStride);

  void RegisterTextureOwner(
      const RemoteTextureOwnerId aOwnerId, const base::ProcessId aForPid,
      const RefPtr<RemoteTextureRecycleBin>& aRecycleBin = nullptr);
  void UnregisterTextureOwner(const RemoteTextureOwnerId aOwnerId,
                              const base::ProcessId aForPid);
  void UnregisterTextureOwners(const RemoteTextureOwnerIdSet& aOwnerIds,
                               const base::ProcessId aForPid);

  void ClearRecycledTextures(
      const RemoteTextureOwnerIdSet& aOwnerIds, const base::ProcessId aForPid,
      const RefPtr<RemoteTextureRecycleBin>& aRecycleBin = nullptr);
  void NotifyContextLost(const RemoteTextureOwnerIdSet& aOwnerIds,
                         const base::ProcessId aForPid);
  void NotifyContextRestored(const RemoteTextureOwnerIdSet& aOwnerIds,
                             const base::ProcessId aForPid);

  bool WaitForRemoteTextureOwner(RemoteTextureHostWrapper* aTextureHostWrapper);

  // Get remote texture's TextureHost for RemoteTextureHostWrapper.
  void GetRemoteTexture(RemoteTextureHostWrapper* aTextureHostWrapper);

  bool WaitForTxn(const RemoteTextureOwnerId aOwnerId,
                  const base::ProcessId aForPid, RemoteTextureTxnType aTxnType,
                  RemoteTextureTxnId aTxnId);

  void ReleaseRemoteTextureHost(RemoteTextureHostWrapper* aTextureHostWrapper);

  RefPtr<TextureHost> GetOrCreateRemoteTextureHostWrapper(
      const RemoteTextureId aTextureId, const RemoteTextureOwnerId aOwnerId,
      const base::ProcessId aForPid, const gfx::IntSize& aSize,
      const TextureFlags aFlags);

  void UnregisterRemoteTextureHostWrapper(const RemoteTextureId aTextureId,
                                          const RemoteTextureOwnerId aOwnerId,
                                          const base::ProcessId aForPid);

  bool CheckRemoteTextureReady(
      const RemoteTextureInfo& aInfo,
      std::function<void(const RemoteTextureInfo&)>&& aCallback);

  bool WaitRemoteTextureReady(const RemoteTextureInfo& aInfo);

  void SuppressRemoteTextureReadyCheck(const RemoteTextureInfo& aInfo);

  UniquePtr<TextureData> GetRecycledTextureData(
      const RemoteTextureOwnerId aOwnerId, const base::ProcessId aForPid,
      const RefPtr<RemoteTextureRecycleBin>& aRecycleBin,
      const gfx::IntSize& aSize, gfx::SurfaceFormat aFormat,
      TextureType aTextureType);

  UniquePtr<SharedResourceWrapper> GetRecycledSharedTexture(
      const RemoteTextureOwnerId aOwnerId, const base::ProcessId aForPid,
      const RefPtr<RemoteTextureRecycleBin>& aRecycleBin,
      const gfx::IntSize& aSize, gfx::SurfaceFormat aFormat,
      SurfaceDescriptor::Type aType);

  static RefPtr<TextureHost> CreateRemoteTexture(TextureData* aTextureData,
                                                 TextureFlags aTextureFlags);

  already_AddRefed<RemoteTextureTxnScheduler> RegisterTxnScheduler(
      base::ProcessId aForPid, RemoteTextureTxnType aType);

 protected:
  friend class RemoteTextureTxnScheduler;

  // Holds data related to remote texture
  struct TextureDataHolder {
    TextureDataHolder(const RemoteTextureId aTextureId,
                      RefPtr<TextureHost> aTextureHost,
                      UniquePtr<TextureData>&& aTextureData,
                      UniquePtr<SharedResourceWrapper>&& aResourceWrapper);

    const RemoteTextureId mTextureId;
    // TextureHost of remote texture
    // Compositable ref of the mTextureHost should be updated within mMonitor.
    // The compositable ref is used to check if TextureHost(remote texture) is
    // still in use by WebRender.
    RefPtr<TextureHost> mTextureHost;
    // Holds BufferTextureData of TextureHost
    UniquePtr<TextureData> mTextureData;
    // Holds gl::SharedSurface of TextureHost
    UniquePtr<SharedResourceWrapper> mResourceWrapper;
  };

  struct RenderingReadyCallbackHolder {
    RenderingReadyCallbackHolder(
        const RemoteTextureId aTextureId,
        std::function<void(const RemoteTextureInfo&)>&& aCallback);

    const RemoteTextureId mTextureId;
    // callback of RemoteTexture ready
    std::function<void(const RemoteTextureInfo&)> mCallback;
  };

  struct WaitingTextureOwner {
    std::deque<UniquePtr<RenderingReadyCallbackHolder>>
        mRenderingReadyCallbackHolders;
  };

  struct TextureOwner {
    bool mIsContextLost = false;
    // Whether to wait for a transaction to complete before unregistering.
    bool mWaitForTxn = false;
    // The thread on which to finally unregister when ready.
    RefPtr<nsISerialEventTarget> mDeferUnregister;

    // Holds TextureDataHolders that wait to be used for building wr display
    // list.
    std::deque<UniquePtr<TextureDataHolder>> mWaitingTextureDataHolders;
    // Holds TextureDataHolders that are used for building wr display list.
    std::deque<UniquePtr<TextureDataHolder>> mUsingTextureDataHolders;
    std::deque<UniquePtr<TextureDataHolder>> mReleasingTextureDataHolders;
    // Holds RemoteTexture ready callbacks.
    std::deque<UniquePtr<RenderingReadyCallbackHolder>>
        mRenderingReadyCallbackHolders;

    RemoteTextureId mLatestPushedTextureId = {0};
    RemoteTextureId mLatestUsingTextureId = {0};
    CompositableTextureHostRef mLatestTextureHost;
    CompositableTextureHostRef mLatestRenderedTextureHost;
    // Holds compositable refs to TextureHosts of RenderTextureHosts that are
    // waiting to be released in non-RenderThread.
    std::deque<CompositableTextureHostRef> mReleasingRenderedTextureHosts;
    RefPtr<RemoteTextureRecycleBin> mRecycleBin;
  };

  // Holds data related to remote texture wrapper
  struct RemoteTextureHostWrapperHolder {
    explicit RemoteTextureHostWrapperHolder(
        RefPtr<TextureHost> aRemoteTextureHostWrapper);

    const RefPtr<TextureHost> mRemoteTextureHostWrapper;
    // Hold compositable ref of remote texture of the RemoteTextureId.
    CompositableTextureHostRef mRemoteTextureHost;
    bool mReadyCheckSuppressed = false;
  };

  void UpdateTexture(const MonitorAutoLock& aProofOfLock,
                     RemoteTextureMap::TextureOwner* aOwner,
                     const RemoteTextureId aTextureId);

  UniquePtr<TextureOwner> UnregisterTextureOwner(
      MonitorAutoLock& aProofOfLock, const RemoteTextureOwnerId aOwnerId,
      const base::ProcessId aForPid,
      std::vector<RefPtr<TextureHost>>& aReleasingTextures,
      std::vector<std::function<void(const RemoteTextureInfo&)>>&
          aRenderingReadyCallbacks);

  void GetRenderingReadyCallbacks(
      const MonitorAutoLock& aProofOfLock,
      RemoteTextureMap::TextureOwner* aOwner, const RemoteTextureId aTextureId,
      std::vector<std::function<void(const RemoteTextureInfo&)>>& aFunctions);

  void GetAllRenderingReadyCallbacks(
      const MonitorAutoLock& aProofOfLock,
      RemoteTextureMap::TextureOwner* aOwner,
      std::vector<std::function<void(const RemoteTextureInfo&)>>& aFunctions);

  void KeepTextureDataAliveForTextureHostIfNecessary(
      const MonitorAutoLock& aProofOfLock,
      RemoteTextureMap::TextureOwner* aOwner,
      std::deque<UniquePtr<TextureDataHolder>>& aHolders);

  bool RecycleTexture(const RefPtr<RemoteTextureRecycleBin>& aRecycleBin,
                      TextureDataHolder& aHolder, bool aExpireOldTextures);

  RemoteTextureMap::TextureOwner* GetTextureOwner(
      const MonitorAutoLock& aProofOfLock, const RemoteTextureOwnerId aOwnerId,
      const base::ProcessId aForPid);

  void NotifyTxn(const MonitorAutoLock& aProofOfLock,
                 const RemoteTextureOwnerId aOwnerId,
                 const base::ProcessId aForPid);

  void UnregisterTxnScheduler(base::ProcessId aForPid,
                              RemoteTextureTxnType aType);

  Monitor mMonitor MOZ_UNANNOTATED;

  std::map<std::pair<base::ProcessId, RemoteTextureOwnerId>,
           UniquePtr<WaitingTextureOwner>>
      mWaitingTextureOwners;

  std::map<std::pair<base::ProcessId, RemoteTextureOwnerId>,
           UniquePtr<TextureOwner>>
      mTextureOwners;

  std::map<std::pair<base::ProcessId, RemoteTextureId>,
           UniquePtr<RemoteTextureHostWrapperHolder>>
      mRemoteTextureHostWrapperHolders;

  std::map<std::pair<base::ProcessId, RemoteTextureTxnType>,
           RemoteTextureTxnScheduler*>
      mTxnSchedulers;

  static StaticAutoPtr<RemoteTextureMap> sInstance;
};

}  // namespace layers
}  // namespace mozilla

#endif  // MOZILLA_GFX_RemoteTextureMap_H