From 6f4610355c5ede9a768ffce8a5b44fcd33e5e6d4 Mon Sep 17 00:00:00 2001
From: Simon Knott <info@simonknott.de>
Date: Wed, 21 Aug 2024 12:32:16 +0200
Subject: [PATCH] first attempt at screenshot in trace viewer

---
 packages/trace-viewer/src/snapshotServer.ts  |  3 +-
 packages/trace-viewer/src/ui/snapshotTab.tsx | 36 ++++++++++++++++----
 packages/trace-viewer/src/ui/uiModeView.tsx  |  3 ++
 packages/trace-viewer/src/ui/workbench.tsx   |  4 ++-
 4 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/packages/trace-viewer/src/snapshotServer.ts b/packages/trace-viewer/src/snapshotServer.ts
index b1dd371cb3..d41dcbdbbd 100644
--- a/packages/trace-viewer/src/snapshotServer.ts
+++ b/packages/trace-viewer/src/snapshotServer.ts
@@ -44,7 +44,8 @@ export class SnapshotServer {
     const snapshot = this._snapshot(pathname.substring('/snapshotInfo'.length), searchParams);
     return this._respondWithJson(snapshot ? {
       viewport: snapshot.viewport(),
-      url: snapshot.snapshot().frameUrl
+      url: snapshot.snapshot().frameUrl,
+      timestamp: snapshot.snapshot().timestamp,
     } : {
       error: 'No snapshot found'
     });
diff --git a/packages/trace-viewer/src/ui/snapshotTab.tsx b/packages/trace-viewer/src/ui/snapshotTab.tsx
index 4faa668677..6e16823302 100644
--- a/packages/trace-viewer/src/ui/snapshotTab.tsx
+++ b/packages/trace-viewer/src/ui/snapshotTab.tsx
@@ -17,10 +17,10 @@
 import './snapshotTab.css';
 import * as React from 'react';
 import type { ActionTraceEvent } from '@trace/trace';
-import { context, prevInList } from './modelUtil';
+import { context, type MultiTraceModel, prevInList } from './modelUtil';
 import { Toolbar } from '@web/components/toolbar';
 import { ToolbarButton } from '@web/components/toolbarButton';
-import { clsx, useMeasure } from '@web/uiUtils';
+import { clsx, useMeasure, useSetting } from '@web/uiUtils';
 import { InjectedScript } from '@injected/injectedScript';
 import { Recorder } from '@injected/recorder/recorder';
 import ConsoleAPI from '@injected/consoleApi';
@@ -30,8 +30,18 @@ import { locatorOrSelectorAsSelector } from '@isomorphic/locatorParser';
 import { TabbedPaneTab } from '@web/components/tabbedPane';
 import { BrowserFrame } from './browserFrame';
 
+function findClosest<T extends { timestamp: number }>(items: T[], target: number) {
+  return items.find((item, index) => {
+    if (index === items.length - 1)
+      return true;
+    const next = items[index + 1];
+    return Math.abs(item.timestamp - target) < Math.abs(next.timestamp - target);
+  });
+}
+
 export const SnapshotTab: React.FunctionComponent<{
   action: ActionTraceEvent | undefined,
+  model?: MultiTraceModel,
   sdkLanguage: Language,
   testIdAttributeName: string,
   isInspecting: boolean,
@@ -39,9 +49,10 @@ export const SnapshotTab: React.FunctionComponent<{
   highlightedLocator: string,
   setHighlightedLocator: (locator: string) => void,
   openPage?: (url: string, target?: string) => Window | any,
-}> = ({ action, sdkLanguage, testIdAttributeName, isInspecting, setIsInspecting, highlightedLocator, setHighlightedLocator, openPage }) => {
+}> = ({ action, model, sdkLanguage, testIdAttributeName, isInspecting, setIsInspecting, highlightedLocator, setHighlightedLocator, openPage }) => {
   const [measure, ref] = useMeasure<HTMLDivElement>();
   const [snapshotTab, setSnapshotTab] = React.useState<'action'|'before'|'after'>('action');
+  const [showScreenshotInsteadOfSnapshot] = useSetting('screenshot-instead-of-snapshot', false, 'Show screenshot instead of snapshot');
 
   type Snapshot = { action: ActionTraceEvent, snapshotName: string, point?: { x: number, y: number } };
   const { snapshots } = React.useMemo(() => {
@@ -90,7 +101,7 @@ export const SnapshotTab: React.FunctionComponent<{
 
   const iframeRef0 = React.useRef<HTMLIFrameElement>(null);
   const iframeRef1 = React.useRef<HTMLIFrameElement>(null);
-  const [snapshotInfo, setSnapshotInfo] = React.useState({ viewport: kDefaultViewport, url: '' });
+  const [snapshotInfo, setSnapshotInfo] = React.useState<{ viewport: typeof kDefaultViewport, url: string, timestamp?: number }>({ viewport: kDefaultViewport, url: '', timestamp: undefined });
   const loadingRef = React.useRef({ iteration: 0, visibleIframe: 0 });
 
   React.useEffect(() => {
@@ -99,13 +110,14 @@ export const SnapshotTab: React.FunctionComponent<{
       const newVisibleIframe = 1 - loadingRef.current.visibleIframe;
       loadingRef.current.iteration = thisIteration;
 
-      const newSnapshotInfo = { url: '', viewport: kDefaultViewport };
+      const newSnapshotInfo = { url: '', viewport: kDefaultViewport, timestamp: undefined };
       if (snapshotInfoUrl) {
         const response = await fetch(snapshotInfoUrl);
         const info = await response.json();
         if (!info.error) {
           newSnapshotInfo.url = info.url;
           newSnapshotInfo.viewport = info.viewport;
+          newSnapshotInfo.timestamp = info.timestamp;
         }
       }
 
@@ -154,6 +166,14 @@ export const SnapshotTab: React.FunctionComponent<{
     y: (measure.height - snapshotContainerSize.height) / 2,
   };
 
+  const page = model?.pages[0]; // TODO: figure out what to do about multiple pages.
+  const screencastFrame = React.useMemo(
+      () => snapshotInfo.timestamp && page?.screencastFrames
+        ? findClosest(page.screencastFrames, snapshotInfo.timestamp)
+        : undefined
+      , [page?.screencastFrames, snapshotInfo.timestamp]
+  );
+
   return <div
     className='snapshot-tab'
     tabIndex={0}
@@ -181,7 +201,7 @@ export const SnapshotTab: React.FunctionComponent<{
       iframe={iframeRef1.current}
       iteration={loadingRef.current.iteration} />
     <Toolbar>
-      <ToolbarButton className='pick-locator' title='Pick locator' icon='target' toggled={isInspecting} onClick={() => setIsInspecting(!isInspecting)} />
+      <ToolbarButton className='pick-locator' title='Pick locator' icon='target' toggled={isInspecting} onClick={() => setIsInspecting(!isInspecting)} disabled={showScreenshotInsteadOfSnapshot /* TODO: proper tooltip */} />
       {['action', 'before', 'after'].map(tab => {
         return <TabbedPaneTab
           key={tab}
@@ -192,6 +212,7 @@ export const SnapshotTab: React.FunctionComponent<{
         ></TabbedPaneTab>;
       })}
       <div style={{ flex: 'auto' }}></div>
+      {/* TODO: disable + proper tooltip. also for locator tab */}
       <ToolbarButton icon='link-external' title='Open snapshot in a new tab' disabled={!popoutUrl} onClick={() => {
         if (!openPage)
           openPage = window.open;
@@ -209,7 +230,8 @@ export const SnapshotTab: React.FunctionComponent<{
         transform: `translate(${translate.x}px, ${translate.y}px) scale(${scale})`,
       }}>
         <BrowserFrame url={snapshotInfo.url} />
-        <div className='snapshot-switcher'>
+        {(showScreenshotInsteadOfSnapshot && screencastFrame) && <img src={`sha1/${screencastFrame.sha1}`} width={screencastFrame.width} height={screencastFrame.height} />}
+        <div className='snapshot-switcher' style={showScreenshotInsteadOfSnapshot ? { display: 'none' } : undefined}>
           <iframe ref={iframeRef0} name='snapshot' title='DOM Snapshot' className={clsx(loadingRef.current.visibleIframe === 0 && 'snapshot-visible')}></iframe>
           <iframe ref={iframeRef1} name='snapshot' title='DOM Snapshot' className={clsx(loadingRef.current.visibleIframe === 1 && 'snapshot-visible')}></iframe>
         </div>
diff --git a/packages/trace-viewer/src/ui/uiModeView.tsx b/packages/trace-viewer/src/ui/uiModeView.tsx
index b12617c300..afbc356352 100644
--- a/packages/trace-viewer/src/ui/uiModeView.tsx
+++ b/packages/trace-viewer/src/ui/uiModeView.tsx
@@ -129,6 +129,8 @@ export const UIModeView: React.FC<{}> = ({
   }, [runUpdateSnapshots, setRunUpdateSnapshots]);
 
   const [, , showRouteActionsSetting] = useSetting('show-route-actions', true, 'Show route actions');
+  const [, , showScreenshotSetting] = useSetting('screenshot-instead-of-snapshot', false, 'Show screenshot instead of snapshot');
+
 
   const darkModeSetting = useDarkModeSetting();
 
@@ -537,6 +539,7 @@ export const UIModeView: React.FC<{}> = ({
         {settingsVisible && <SettingsView settings={[
           darkModeSetting,
           showRouteActionsSetting,
+          showScreenshotSetting,
         ]} />}
       </div>
       }
diff --git a/packages/trace-viewer/src/ui/workbench.tsx b/packages/trace-viewer/src/ui/workbench.tsx
index 01ae6142bd..0f72792af9 100644
--- a/packages/trace-viewer/src/ui/workbench.tsx
+++ b/packages/trace-viewer/src/ui/workbench.tsx
@@ -73,6 +73,7 @@ export const Workbench: React.FunctionComponent<{
   const [selectedTime, setSelectedTime] = React.useState<Boundaries | undefined>();
   const [sidebarLocation, setSidebarLocation] = useSetting<'bottom' | 'right'>('propertiesSidebarLocation', 'bottom');
   const [showRouteActions, , showRouteActionsSetting] = useSetting('show-route-actions', true, 'Show route actions');
+  const [, , showScreenshotSetting] = useSetting('screenshot-instead-of-snapshot', false, 'Show screenshot instead of snapshot');
 
   const filteredActions = React.useMemo(() => {
     return (model?.actions || []).filter(action => showRouteActions || !isRouteAction(action));
@@ -291,7 +292,7 @@ export const Workbench: React.FunctionComponent<{
   const settingsTab: TabbedPaneTabModel = {
     id: 'settings',
     title: 'Settings',
-    component: <SettingsView settings={[showRouteActionsSetting]}/>,
+    component: <SettingsView settings={[showRouteActionsSetting, showScreenshotSetting]}/>,
   };
 
   return <div className='vbox workbench' {...(inert ? { inert: 'true' } : {})}>
@@ -317,6 +318,7 @@ export const Workbench: React.FunctionComponent<{
         settingName='actionListSidebar'
         main={<SnapshotTab
           action={activeAction}
+          model={model}
           sdkLanguage={sdkLanguage}
           testIdAttributeName={model?.testIdAttributeName || 'data-testid'}
           isInspecting={isInspecting}