1import { useCallback, useEffect, useState, useRef } from 'react';
2import {
3  TDUserStatus,
4  TDAsset,
5  TDBinding,
6  TDShape,
7  TDUser,
8  TldrawApp,
9} from '@tldraw/tldraw';
10import { useThrottleCallback } from '@react-hook/throttle';
11import * as yorkie from '@yorkie-js/sdk';
12import randomColor from 'randomcolor';
13import { uniqueNamesGenerator, names } from 'unique-names-generator';
14import _ from 'lodash';
15
16import type { YorkieDocType, YorkiePresenceType } from './types';
17
18/**
19 * Custom hook for managing multiplayer state in Tldraw using Yorkie
20 * Handles real-time collaboration, presence updates, and undo/redo functionality
21 */
22export function useMultiplayerState(roomId: string) {
23  const [app, setApp] = useState<TldrawApp>();
24  const [loading, setLoading] = useState(true);
25
26  const clientRef = useRef<yorkie.Client>();
27  const docRef = useRef<yorkie.Document<YorkieDocType, YorkiePresenceType>>();
28
29  /**
30   * Handles changes from remote users and updates the local Tldraw app
31   * Parses document content and replaces page content in the app
32   */
33  const handleChanges = useCallback(() => {
34    if (!docRef.current || !app) return;
35
36    const root = docRef.current.getRoot();
37
38    try {
39      const shapesJson = root.shapes.toJSON?.();
40      const bindingsJson = root.bindings.toJSON?.();
41      const assetsJson = root.assets.toJSON?.();
42
43      if (!shapesJson || !bindingsJson || !assetsJson) {
44        console.warn('Document content is not yet initialized');
45        return;
46      }
47
48      const shapeRecord: Record<string, TDShape> = JSON.parse(shapesJson);
49      const bindingRecord: Record<string, TDBinding> = JSON.parse(bindingsJson);
50      const assetRecord: Record<string, TDAsset> = JSON.parse(assetsJson);
51
52      app.replacePageContent(shapeRecord, bindingRecord, assetRecord);
53    } catch (error) {
54      console.error('Error parsing document content:', error);
55    }
56  }, [app]);
57
58  /**
59   * Utility function to get the list of properties that have changed
60   * between source and target objects using deep equality comparison
61   */
62  const getUpdatedPropertyList = useCallback(
63    <T extends object>(source: T, target: T) => {
64      return (Object.keys(source) as Array<keyof T>).filter(
65        (key) => !_.isEqual(source[key], target[key]),
66      );
67    },
68    [],
69  );
70
71  /**
72   * Callback for when Tldraw app is mounted
73   * Initializes the room, sets loading state, and creates initial user
74   */
75  const onMount = useCallback(
76    (app: TldrawApp) => {
77      app.loadRoom(roomId);
78      app.setIsLoading(true);
79      app.pause();
80      setApp(app);
81
82      const randomName = uniqueNamesGenerator({
83        dictionaries: [names],
84      });
85
86      app.updateUsers([
87        {
88          id: app!.currentUser!.id,
89          point: [0, 0],
90          color: randomColor(),
91          status: TDUserStatus.Connected,
92          activeShapes: [],
93          selectedIds: [],
94          metadata: { name: randomName },
95        },
96      ]);
97    },
98    [roomId],
99  );
100
101  /** Handle undo operation using Yorkie's history API */
102  const onUndo = useCallback(() => {
103    if (docRef.current?.history.canUndo()) {
104      docRef.current.history.undo();
105      handleChanges();
106    }
107  }, [handleChanges]);
108
109  /** Handle redo operation using Yorkie's history API */
110  const onRedo = useCallback(() => {
111    if (docRef.current?.history.canRedo()) {
112      docRef.current.history.redo();
113      handleChanges();
114    }
115  }, [handleChanges]);
116
117  /**
118   * Throttled callback for handling page content changes
119   * Updates shapes, bindings, and assets in the Yorkie document
120   * Throttled to 60ms to prevent excessive updates
121   */
122  const onChangePage = useThrottleCallback(
123    (
124      app: TldrawApp,
125      shapes: Record<string, TDShape | undefined>,
126      bindings: Record<string, TDBinding | undefined>,
127    ) => {
128      if (!app || !clientRef.current || !docRef.current) return;
129
130      // Update shapes in the document
131      for (const [id, shape] of Object.entries(shapes)) {
132        docRef.current.update((root: YorkieDocType) => {
133          if (!shape) {
134            delete root.shapes[id];
135          } else if (!root.shapes[id]) {
136            root.shapes[id] = shape as yorkie.JSONObject<TDShape>;
137          } else {
138            const updatedPropertyList = getUpdatedPropertyList(
139              shape,
140              root.shapes[id].toJS(),
141            );
142
143            for (const key of updatedPropertyList) {
144              (root.shapes[id] as any)[key] = shape[key];
145            }
146          }
147        });
148      }
149
150      // Update bindings in the document
151      for (const [id, binding] of Object.entries(bindings)) {
152        docRef.current.update((root: YorkieDocType) => {
153          if (!binding) {
154            delete root.bindings[id];
155          } else if (!root.bindings[id]) {
156            root.bindings[id] = binding as yorkie.JSONObject<TDBinding>;
157          } else {
158            const updatedPropertyList = getUpdatedPropertyList(
159              binding,
160              root.bindings[id].toJS(),
161            );
162
163            for (const key of updatedPropertyList) {
164              const newValue = binding[key];
165              if (newValue !== undefined) {
166                (root.bindings[id] as any)[key] = newValue;
167              }
168            }
169          }
170        });
171      }
172
173      /**
174       * Update assets in the document
175       * Assets are stored globally in app.assets and referenced by inner page assets
176       * Document keys for assets use asset.id (string), not array indices
177       */
178      for (const [, asset] of Object.entries(app.assets)) {
179        docRef.current.update((root: YorkieDocType) => {
180          if (!asset.id) {
181            // Skip assets without valid IDs
182            return;
183          } else if (!root.assets[asset.id]) {
184            root.assets[asset.id] = asset as yorkie.JSONObject<TDAsset>;
185          } else {
186            const updatedPropertyList = getUpdatedPropertyList(
187              asset,
188              root.assets[asset.id].toJS(),
189            );
190
191            for (const key of updatedPropertyList) {
192              (root.assets[asset.id] as any)[key] = asset[key];
193            }
194          }
195        });
196      }
197    },
198    60,
199    false,
200  );
201
202  /**
203   * Throttled callback for handling user presence updates
204   * Updates cursor position, selection, and other user state
205   * Throttled to 60ms to prevent excessive presence updates
206   */
207  const onChangePresence = useThrottleCallback(
208    (app: TldrawApp, user: TDUser) => {
209      if (!app || !clientRef.current?.isActive() || !docRef.current) return;
210
211      docRef.current.update((root, presence) => {
212        presence.set({ tdUser: user });
213      });
214    },
215    60,
216    false,
217  );
218
219  useEffect(() => {
220    if (!app) return;
221
222    const unsubs: Array<Function> = [];
223    let stillAlive = true;
224
225    /**
226     * Set up document change subscription for remote updates
227     * Triggers handleChanges when remote changes are received
228     */
229    const setupDocumentSubscription = (
230      doc: yorkie.Document<YorkieDocType, YorkiePresenceType>,
231    ) => {
232      doc.subscribe((event) => {
233        if (event.type === 'remote-change') {
234          handleChanges();
235        }
236      });
237    };
238
239    /**
240     * Finalize setup after document is attached
241     * Syncs data, handles initial zoom, and sets loading state
242     */
243    const finalizeSetup = async (client: yorkie.Client) => {
244      await client.sync();
245
246      if (stillAlive) {
247        handleChanges();
248        if (app) {
249          app.zoomToFit();
250          if (app.zoom > 1) {
251            app.resetZoom();
252          }
253          app.setIsLoading(false);
254        }
255
256        setLoading(false);
257      }
258    };
259
260    /**
261     * Main setup function that orchestrates the entire initialization process
262     * Creates document and client, sets up subscriptions, and handles errors
263     */
264    const setupDocument = async () => {
265      try {
266        // 01. Activate Yorkie client.
267        const client = new yorkie.Client({
268          rpcAddr: import.meta.env.VITE_YORKIE_API_ADDR,
269          apiKey: import.meta.env.VITE_YORKIE_API_KEY,
270          syncLoopDuration: 0,
271          reconnectStreamDelay: 1000,
272        });
273        await client.activate();
274        clientRef.current = client;
275
276        // 02. Create document and subscribe to events.
277        const doc = new yorkie.Document<YorkieDocType, YorkiePresenceType>(
278          roomId,
279          { enableDevtools: true },
280        );
281        docRef.current = doc;
282        unsubs.push(
283          doc.subscribe('my-presence', (event) => {
284            if (event.type === yorkie.DocEventType.Initialized) {
285              const allPeers = doc
286                .getPresences()
287                .map((peer) => peer.presence.tdUser);
288              app?.updateUsers(allPeers);
289            }
290          }),
291        );
292
293        unsubs.push(
294          doc.subscribe('others', (event) => {
295            if (event.type === yorkie.DocEventType.Unwatched) {
296              app?.removeUser(event.value.presence.tdUser.id);
297            }
298
299            const allPeers = doc
300              .getPresences()
301              .map((peer) => peer.presence.tdUser);
302            app?.updateUsers(allPeers);
303          }),
304        );
305
306        unsubs.push(
307          doc.subscribe((event) => {
308            if (event.type === 'remote-change') {
309              handleChanges();
310            }
311          }),
312        );
313
314        await client.attach(doc, {
315          initialRoot: {
316            shapes: {},
317            bindings: {},
318            assets: {},
319          } as YorkieDocType,
320          initialPresence: app?.currentUser
321            ? { tdUser: app.currentUser }
322            : undefined,
323        });
324
325        await finalizeSetup(client);
326      } catch (error) {
327        console.error('Error setting up document:', error);
328        setLoading(false);
329      }
330    };
331
332    setupDocument();
333
334    // Cleanup function to properly dispose of resources
335    return () => {
336      stillAlive = false;
337      for (const unsub of unsubs) {
338        unsub();
339      }
340
341      const cleanup = async () => {
342        try {
343          if (docRef.current && clientRef.current) {
344            await clientRef.current.detach(docRef.current);
345          }
346          if (clientRef.current) {
347            await clientRef.current.deactivate();
348          }
349        } catch (error) {
350          console.error('Error during cleanup:', error);
351        } finally {
352          docRef.current = undefined;
353          clientRef.current = undefined;
354        }
355      };
356      cleanup();
357    };
358  }, [app, handleChanges]);
359
360  return {
361    onMount,
362    onChangePage,
363    loading,
364    onChangePresence,
365    onUndo,
366    onRedo,
367  };
368}