From 41d6e0f9f9a35c2b80ec938f3d0a85aa0318d82f Mon Sep 17 00:00:00 2001
From: RunDevelopment <mitchi5000.ms@googlemail.com>
Date: Mon, 27 Nov 2023 16:40:37 +0100
Subject: [PATCH] Improved screenshots

---
 src/renderer/components/ReactFlowBox.tsx |  4 +---
 src/renderer/helpers/nodeScreenshot.ts   | 20 +++++++++++++-------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/renderer/components/ReactFlowBox.tsx b/src/renderer/components/ReactFlowBox.tsx
index 625515b83..018ecf25d 100644
--- a/src/renderer/components/ReactFlowBox.tsx
+++ b/src/renderer/components/ReactFlowBox.tsx
@@ -426,9 +426,6 @@ export const ReactFlowBox = memo(({ wrapperRef, nodeTypes, edgeTypes }: ReactFlo
 
     return (
         <Box
-            bg="var(--chain-editor-bg)"
-            borderRadius="lg"
-            borderWidth="0px"
             className={animateChain ? '' : 'no-chain-animation'}
             h="100%"
             ref={wrapperRef}
@@ -452,6 +449,7 @@ export const ReactFlowBox = memo(({ wrapperRef, nodeTypes, edgeTypes }: ReactFlo
                 style={{
                     zIndex: 0,
                     borderRadius: '0.5rem',
+                    backgroundColor: 'var(--chain-editor-bg)',
                 }}
                 onConnect={createConnection}
                 onConnectEnd={onConnectStop}
diff --git a/src/renderer/helpers/nodeScreenshot.ts b/src/renderer/helpers/nodeScreenshot.ts
index b17ad7fbe..5c16eccdf 100644
--- a/src/renderer/helpers/nodeScreenshot.ts
+++ b/src/renderer/helpers/nodeScreenshot.ts
@@ -36,6 +36,9 @@ export const takeScreenshot = async (
 ): Promise<PngDataUrl> => {
     const oldViewport = reactFlow.getViewport();
 
+    const rfElement = (currentFlowWrapper.firstElementChild ?? currentFlowWrapper) as HTMLElement;
+    const oldBorderRadius = rfElement.style.borderRadius;
+
     const reactFlowViewport = currentFlowWrapper.getBoundingClientRect();
     const nodes = reactFlow.getNodes();
     const nodesBoundingBox = getNodesBoundingBox(nodes) ?? { x: 0, y: 0, width: 0, height: 0 };
@@ -52,17 +55,18 @@ export const takeScreenshot = async (
         reactFlowViewport.height / paddedBoundingBox.height
     );
 
-    reactFlow.setViewport({
-        x: paddedBoundingBox.x * -1 * exportZoom,
-        y: paddedBoundingBox.y * -1 * exportZoom,
-        zoom: exportZoom,
-    });
-
     try {
+        reactFlow.setViewport({
+            x: paddedBoundingBox.x * -1 * exportZoom,
+            y: paddedBoundingBox.y * -1 * exportZoom,
+            zoom: exportZoom,
+        });
+        rfElement.style.borderRadius = '0';
+
         // wait for the viewport to be updated
         await delay(10);
 
-        const dataUrl = await toPng(currentFlowWrapper, {
+        const dataUrl = await toPng(rfElement, {
             style: {
                 padding: '0',
                 margin: '0',
@@ -71,6 +75,7 @@ export const takeScreenshot = async (
             pixelRatio: 1 / exportZoom,
             width: paddedBoundingBox.width * exportZoom,
             height: paddedBoundingBox.height * exportZoom,
+            backgroundColor: getComputedStyle(rfElement).backgroundColor,
             filter: (node: unknown) => {
                 if (
                     node instanceof HTMLElement &&
@@ -86,6 +91,7 @@ export const takeScreenshot = async (
         return dataUrl as PngDataUrl;
     } finally {
         reactFlow.setViewport(oldViewport);
+        rfElement.style.borderRadius = oldBorderRadius;
     }
 };