import { getInlineEditorByModel } from '@blocksuite/affine-rich-text'; import { getSelectedBlocksCommand } from '@blocksuite/affine-shared/commands'; import { BlockElementCommentManager, type CommentId, CommentProviderIdentifier, findAllCommentedBlocks, findAllCommentedElements, } from '@blocksuite/affine-shared/services'; import type { AffineInlineEditor } from '@blocksuite/affine-shared/types'; import { DisposableGroup } from '@blocksuite/global/disposable'; import { LifeCycleWatcher, type TextRangePoint, TextSelection, } from '@blocksuite/std'; import type { BaseSelection, BlockModel } from '@blocksuite/store'; import { signal } from '@preact/signals-core'; import difference from 'lodash-es/difference'; import { extractCommentIdFromDelta, findAllCommentedTexts, findCommentedTexts, } from './utils'; export class InlineCommentManager extends LifeCycleWatcher { static override key = 'inline-comment-manager'; private readonly _disposables = new DisposableGroup(); private readonly _highlightedCommentId$ = signal(null); private get _provider() { return this.std.getOptional(CommentProviderIdentifier); } override mounted() { const provider = this._provider; if (!provider) return; this._init().catch(console.error); this._disposables.add(provider.onCommentAdded(this._handleAddComment)); this._disposables.add( provider.onCommentDeleted(id => this._handleDeleteAndResolve(id, 'delete') ) ); this._disposables.add( provider.onCommentResolved(id => this._handleDeleteAndResolve(id, 'resolve') ) ); this._disposables.add( provider.onCommentHighlighted(this._handleHighlightComment) ); this._disposables.add( this.std.selection.slots.changed.subscribe(this._handleSelectionChanged) ); } override unmounted() { this._disposables.dispose(); } private async _init() { const provider = this._provider; if (!provider) return; const commentsInProvider = await provider.getComments('all'); const commentsInEditor = this.getCommentsInEditor(); // remove comments that are in editor but not in provider // which means the comment may be removed or resolved in provider side difference(commentsInEditor, commentsInProvider).forEach(comment => { this.std .get(BlockElementCommentManager) .handleDeleteAndResolve(comment, 'delete'); }); } getCommentsInEditor() { const inlineComments = [...findAllCommentedTexts(this.std.store).values()]; const blockComments = findAllCommentedBlocks(this.std.store).flatMap( block => Object.keys(block.props.comments) ); const surfaceComments = findAllCommentedElements(this.std.store).flatMap( element => Object.keys(element.comments) ); const commentsInEditor = [ ...new Set([...inlineComments, ...blockComments, ...surfaceComments]), ]; return commentsInEditor; } private readonly _handleAddComment = ( id: CommentId, selections: BaseSelection[] ) => { const needCommentTexts = selections .map(selection => { if (!selection.is(TextSelection)) return []; const [_, { selectedBlocks }] = this.std.command .chain() .pipe(getSelectedBlocksCommand, { textSelection: selection, }) .run(); if (!selectedBlocks) return []; type MakeRequired = T & { [key in K]: NonNullable; }; return selectedBlocks .map( ({ model }) => [model, getInlineEditorByModel(this.std, model)] as const ) .filter( ( pair ): pair is [MakeRequired, AffineInlineEditor] => !!pair[0].text && !!pair[1] ) .map(([model, inlineEditor]) => { let from: TextRangePoint; let to: TextRangePoint | null; if (model.id === selection.from.blockId) { from = selection.from; to = null; } else if (model.id === selection.to?.blockId) { from = selection.to; to = null; } else { from = { blockId: model.id, index: 0, length: model.text.yText.length, }; to = null; } return [new TextSelection({ from, to }), inlineEditor] as const; }); }) .flat(); if (needCommentTexts.length === 0) return; needCommentTexts.forEach(([selection, inlineEditor]) => { inlineEditor.formatText( selection.from, { [`comment-${id}`]: true, }, { withoutTransact: true, } ); }); }; private readonly _handleDeleteAndResolve = ( id: CommentId, type: 'delete' | 'resolve' ) => { const commentedTexts = findCommentedTexts(this.std.store, id); if (commentedTexts.length === 0) return; this.std.store.withoutTransact(() => { commentedTexts.forEach(selection => { const inlineEditor = getInlineEditorByModel( this.std, selection.from.blockId ); inlineEditor?.formatText( selection.from, { [`comment-${id}`]: type === 'delete' ? null : false, }, { withoutTransact: true, } ); }); }); }; private readonly _handleHighlightComment = (id: CommentId | null) => { this._highlightedCommentId$.value = id; }; private readonly _handleSelectionChanged = (selections: BaseSelection[]) => { const currentHighlightedCommentId = this._highlightedCommentId$.peek(); if (selections.length === 1) { const selection = selections[0]; // InlineCommentManager only handle text selection if (!selection.is(TextSelection)) return; if (!selection.isCollapsed() && currentHighlightedCommentId !== null) { this._provider?.highlightComment(null); return; } const model = this.std.store.getModelById(selection.from.blockId); if (!model) return; const inlineEditor = getInlineEditorByModel(this.std, model); if (!inlineEditor) return; const delta = inlineEditor.getDeltaByRangeIndex(selection.from.index); if (!delta) return; const commentIds = extractCommentIdFromDelta(delta); if (commentIds.length !== 0) return; } if (currentHighlightedCommentId !== null) { this._provider?.highlightComment(null); } }; }