From a3ca41fd6ae03f3f2e1312049a0867dfba25f0af Mon Sep 17 00:00:00 2001
From: doouding <github@douding.io>
Date: Thu, 13 Jun 2024 07:05:14 +0000
Subject: [PATCH] fix: surface-ref peek view (#7208)

### Change
Add ref to `SurfaceRefPeekView`. It provide `fitViewportToTarget` method to fit the `surface-ref` content.

Related to [AFF-1200](https://linear.app/affine-design/issue/AFF-1200/center-peek-frame-rendering-issue).
---
 .../modules/peek-view/view/doc-peek-view.tsx  | 67 +++++++++++++------
 .../peek-view/view/modal-container.tsx        |  8 +--
 .../peek-view/view/peek-view-manager.tsx      | 43 +++++++++---
 3 files changed, 83 insertions(+), 35 deletions(-)

diff --git a/packages/frontend/core/src/modules/peek-view/view/doc-peek-view.tsx b/packages/frontend/core/src/modules/peek-view/view/doc-peek-view.tsx
index 6450299c35..13727f4e50 100644
--- a/packages/frontend/core/src/modules/peek-view/view/doc-peek-view.tsx
+++ b/packages/frontend/core/src/modules/peek-view/view/doc-peek-view.tsx
@@ -9,7 +9,13 @@ import { DisposableGroup } from '@blocksuite/global/utils';
 import { type AffineEditorContainer, AIProvider } from '@blocksuite/presets';
 import type { DocMode } from '@toeverything/infra';
 import { DocsService, FrameworkScope, useService } from '@toeverything/infra';
-import { forwardRef, useEffect, useState } from 'react';
+import {
+  forwardRef,
+  useCallback,
+  useEffect,
+  useImperativeHandle,
+  useState,
+} from 'react';
 
 import { WorkbenchService } from '../../workbench';
 import { PeekViewService } from '../services/peek-view';
@@ -113,6 +119,7 @@ const DocPreview = forwardRef<
     </AffineErrorBoundary>
   );
 });
+DocPreview.displayName = 'DocPreview';
 
 export const DocPeekView = ({
   docId,
@@ -126,16 +133,42 @@ export const DocPeekView = ({
   return <DocPreview mode={mode} docId={docId} blockId={blockId} />;
 };
 
-export const SurfaceRefPeekView = ({
-  docId,
-  xywh,
-}: {
-  docId: string;
-  xywh: `[${number},${number},${number},${number}]`;
-}) => {
+export type SurfaceRefPeekViewRef = {
+  fitViewportToTarget: () => void;
+};
+
+export const SurfaceRefPeekView = forwardRef<
+  SurfaceRefPeekViewRef,
+  { docId: string; xywh: `[${number},${number},${number},${number}]` }
+>(function SurfaceRefPeekView({ docId, xywh }, ref) {
   const [editorRef, setEditorRef] = useState<AffineEditorContainer | null>(
     null
   );
+  const fitViewportToTarget = useCallback(() => {
+    if (!editorRef) {
+      return;
+    }
+
+    const viewport = {
+      xywh: xywh,
+      padding: [60, 20, 20, 20] as [number, number, number, number],
+    };
+    const rootService =
+      editorRef.host.std.spec.getService<EdgelessRootService>('affine:page');
+    rootService.viewport.onResize();
+    rootService.viewport.setViewportByBound(
+      Bound.deserialize(viewport.xywh),
+      viewport.padding
+    );
+  }, [editorRef, xywh]);
+
+  useImperativeHandle(
+    ref,
+    () => ({
+      fitViewportToTarget,
+    }),
+    [fitViewportToTarget]
+  );
 
   useEffect(() => {
     let mounted = true;
@@ -143,18 +176,7 @@ export const SurfaceRefPeekView = ({
       editorRef.host?.updateComplete
         .then(() => {
           if (mounted) {
-            const viewport = {
-              xywh: xywh,
-              padding: [60, 20, 20, 20] as [number, number, number, number],
-            };
-            const rootService =
-              editorRef.host.std.spec.getService<EdgelessRootService>(
-                'affine:page'
-              );
-            rootService.viewport.setViewportByBound(
-              Bound.deserialize(viewport.xywh),
-              viewport.padding
-            );
+            fitViewportToTarget();
           }
         })
         .catch(e => {
@@ -164,7 +186,8 @@ export const SurfaceRefPeekView = ({
     return () => {
       mounted = false;
     };
-  }, [editorRef, xywh]);
+  }, [editorRef, fitViewportToTarget]);
 
   return <DocPreview ref={setEditorRef} docId={docId} mode={'edgeless'} />;
-};
+});
+SurfaceRefPeekView.displayName = 'SurfaceRefPeekView';
diff --git a/packages/frontend/core/src/modules/peek-view/view/modal-container.tsx b/packages/frontend/core/src/modules/peek-view/view/modal-container.tsx
index 3354883dcb..a8d21a8ffb 100644
--- a/packages/frontend/core/src/modules/peek-view/view/modal-container.tsx
+++ b/packages/frontend/core/src/modules/peek-view/view/modal-container.tsx
@@ -75,11 +75,6 @@ export const PeekViewModalContainer = ({
 }>) => {
   const [{ status }, toggle] = useTransition({
     timeout: animationTimeout,
-    onStateChange(event) {
-      if (event.current.status === 'exited' && onAnimateEnd) {
-        onAnimateEnd();
-      }
-    },
   });
   const [transformOrigin, setTransformOrigin] = useState<string | null>(null);
   useEffect(() => {
@@ -105,6 +100,9 @@ export const PeekViewModalContainer = ({
               [styles.transformOrigin]: transformOrigin,
               [styles.animationTimeout]: `${animationTimeout}ms`,
             })}
+            onAnimationEnd={() => {
+              onAnimateEnd?.();
+            }}
           />
           <div
             data-peek-view-wrapper
diff --git a/packages/frontend/core/src/modules/peek-view/view/peek-view-manager.tsx b/packages/frontend/core/src/modules/peek-view/view/peek-view-manager.tsx
index 1e7df11b76..0c2747c116 100644
--- a/packages/frontend/core/src/modules/peek-view/view/peek-view-manager.tsx
+++ b/packages/frontend/core/src/modules/peek-view/view/peek-view-manager.tsx
@@ -1,15 +1,26 @@
+import { BlockElement } from '@blocksuite/block-std';
 import { useLiveData, useService } from '@toeverything/infra';
-import { useMemo } from 'react';
+import { useEffect, useMemo, useRef } from 'react';
 
 import type { ActivePeekView } from '../entities/peek-view';
 import { PeekViewService } from '../services/peek-view';
 import { DocPeekViewControls } from './doc-peek-controls';
+import type { SurfaceRefPeekViewRef } from './doc-peek-view';
 import { DocPeekView, SurfaceRefPeekView } from './doc-peek-view';
 import { PeekViewModalContainer } from './modal-container';
 
-function renderPeekView({ info }: ActivePeekView) {
+function renderPeekView(
+  { info }: ActivePeekView,
+  refCallback: (editor: SurfaceRefPeekViewRef) => void
+) {
   if (info.mode === 'edgeless' && info.xywh) {
-    return <SurfaceRefPeekView docId={info.docId} xywh={info.xywh} />;
+    return (
+      <SurfaceRefPeekView
+        ref={refCallback}
+        docId={info.docId}
+        xywh={info.xywh}
+      />
+    );
   }
 
   return (
@@ -31,15 +42,32 @@ export const PeekViewManagerModal = () => {
   const peekViewEntity = useService(PeekViewService).peekView;
   const activePeekView = useLiveData(peekViewEntity.active$);
   const show = useLiveData(peekViewEntity.show$);
+  const peekViewRef = useRef<SurfaceRefPeekViewRef | null>(null);
 
   const preview = useMemo(() => {
-    return activePeekView ? renderPeekView(activePeekView) : null;
+    return activePeekView
+      ? renderPeekView(activePeekView, editor => {
+          peekViewRef.current = editor;
+        })
+      : null;
   }, [activePeekView]);
 
   const controls = useMemo(() => {
     return activePeekView ? renderControls(activePeekView) : null;
   }, [activePeekView]);
 
+  useEffect(() => {
+    const subscription = peekViewEntity.show$.subscribe(() => {
+      if (activePeekView?.target instanceof BlockElement) {
+        activePeekView.target.requestUpdate();
+      }
+    });
+
+    return () => {
+      subscription.unsubscribe();
+    };
+  }, [activePeekView, peekViewEntity]);
+
   return (
     <PeekViewModalContainer
       open={show && !!preview}
@@ -49,15 +77,14 @@ export const PeekViewManagerModal = () => {
           : undefined
       }
       controls={controls}
-      // there is a bug for edgeless mode when showing the peek view during start up animation
-      hideOnEntering={
-        !activePeekView?.info.mode || activePeekView?.info.mode === 'edgeless'
-      }
       onOpenChange={open => {
         if (!open) {
           peekViewEntity.close();
         }
       }}
+      onAnimateEnd={() => {
+        peekViewRef.current?.fitViewportToTarget();
+      }}
     >
       {preview}
     </PeekViewModalContainer>