From a769bc0c0ebd9e8fee0cfad09534dc3b05777871 Mon Sep 17 00:00:00 2001
From: Ayush Agrawal <54364088+AyushAgrawal-A2@users.noreply.github.com>
Date: Tue, 19 Dec 2023 12:41:37 +0530
Subject: [PATCH] fix: element handling with `getPageByElement` in doc &
 edgeless (#5771)

---
 .../rich-text/inline/nodes/reference-node.ts  | 10 +++----
 packages/blocks/src/_common/utils/query.ts    | 27 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/packages/blocks/src/_common/components/rich-text/inline/nodes/reference-node.ts b/packages/blocks/src/_common/components/rich-text/inline/nodes/reference-node.ts
index d9d5e1035071..ba443f9a0bf9 100644
--- a/packages/blocks/src/_common/components/rich-text/inline/nodes/reference-node.ts
+++ b/packages/blocks/src/_common/components/rich-text/inline/nodes/reference-node.ts
@@ -10,12 +10,12 @@ import type { Page, PageMeta } from '@blocksuite/store';
 import { css, html } from 'lit';
 import { customElement, property, state } from 'lit/decorators.js';
 
-import type { DocPageBlockComponent } from '../../../../../page-block/doc/doc-page-block.js';
+import type { PageBlockComponent } from '../../../../../index.js';
 import { FontLinkedPageIcon, FontPageIcon } from '../../../../icons/index.js';
 import {
   getClosestBlockElementByElement,
-  getDocPageByElement,
   getModelByElement,
+  getPageByElement,
 } from '../../../../utils/index.js';
 import { DEFAULT_PAGE_NAME, REFERENCE_NODE } from '../../consts.js';
 import type { AffineTextAttributes } from '../types.js';
@@ -127,9 +127,9 @@ export class AffineReference extends WithDisposable(ShadowlessElement) {
     const targetPageId = refMeta.id;
     const root = model.page.root;
     assertExists(root);
-    const docPageElement = getDocPageByElement(this) as DocPageBlockComponent;
-    assertExists(docPageElement);
-    docPageElement.slots.pageLinkClicked.emit({ pageId: targetPageId });
+    const pageElement = getPageByElement(this) as PageBlockComponent;
+    assertExists(pageElement);
+    pageElement.slots.pageLinkClicked.emit({ pageId: targetPageId });
   }
 
   override render() {
diff --git a/packages/blocks/src/_common/utils/query.ts b/packages/blocks/src/_common/utils/query.ts
index 4039666c7d2a..af6cd21e7bf6 100644
--- a/packages/blocks/src/_common/utils/query.ts
+++ b/packages/blocks/src/_common/utils/query.ts
@@ -5,6 +5,7 @@ import type { BaseBlockModel, Page } from '@blocksuite/store';
 
 import type { Loader } from '../../_common/components/loader.js';
 import type { RichText } from '../../_common/components/rich-text/rich-text.js';
+import type { PageBlockComponent } from '../../index.js';
 import type { DocPageBlockComponent } from '../../page-block/doc/doc-page-block.js';
 import type { EdgelessCanvasTextEditor } from '../../page-block/edgeless/components/text/types.js';
 import type { EdgelessPageBlockComponent } from '../../page-block/edgeless/edgeless-page-block.js';
@@ -166,6 +167,30 @@ export function buildPath(model: BaseBlockModel | null): string[] {
   return path;
 }
 
+export function getPageByElement(element: Element): PageBlockComponent | null {
+  const docPageElement = getDocPageByElement(element);
+  if (docPageElement) return docPageElement;
+
+  const edgelessPageElement = getEdgelessPageByElement(element);
+  if (edgelessPageElement) return edgelessPageElement;
+
+  return null;
+}
+
+export function getPageByEditorHost(
+  editorHost: EditorHost
+): PageBlockComponent | null {
+  if (isInsideDocEditor(editorHost)) {
+    return getDocPageByEditorHost(editorHost);
+  }
+
+  if (isInsideEdgelessEditor(editorHost)) {
+    return getEdgelessPageByEditorHost(editorHost);
+  }
+
+  return null;
+}
+
 /** If it's not in the page mode, it will return `null` directly
  * Use `getDocPageByElement` or `getDocPageByEditorHost` instead.
  * @deprecated
@@ -250,7 +275,7 @@ export function getLitRoot() {
  * ```
  */
 export function getViewportElement(editorHost: EditorHost) {
-  if (isInsideDocEditor(editorHost)) return null;
+  if (!isInsideDocEditor(editorHost)) return null;
   const page = editorHost.page;
   assertExists(page.root);
   const pageComponent = editorHost.view.viewFromPath('block', [page.root.id]);