From 008ec441783c618e0e9177eec66f4b0d26c86a14 Mon Sep 17 00:00:00 2001
From: Srujan Gurram <srujangs8@gmail.com>
Date: Sun, 26 Nov 2023 14:59:20 +0530
Subject: [PATCH] fixed image capture

---
 package.json                                  |   1 +
 .../Sidebar/chat/ImageCaptureButton.tsx       |   6 +-
 src/lib/getScreenshotImage.ts                 | 161 +++++++++++++-----
 yarn.lock                                     |  34 ++++
 4 files changed, 152 insertions(+), 50 deletions(-)

diff --git a/package.json b/package.json
index d7388f3c..6a50e7f6 100644
--- a/package.json
+++ b/package.json
@@ -32,6 +32,7 @@
     "axios": "^1.3.5",
     "dnd-kit-sortable-tree": "^0.1.58",
     "endent": "^2.1.0",
+    "html2canvas": "^1.4.1",
     "jotai": "^2.4.3",
     "langchain": "^0.0.197-rc.1",
     "object-hash": "^3.0.0",
diff --git a/src/components/Sidebar/chat/ImageCaptureButton.tsx b/src/components/Sidebar/chat/ImageCaptureButton.tsx
index 1bb5c66f..d8c4442e 100644
--- a/src/components/Sidebar/chat/ImageCaptureButton.tsx
+++ b/src/components/Sidebar/chat/ImageCaptureButton.tsx
@@ -8,7 +8,7 @@ const ImageCaptureButton = ({
   addMessageDraftFile,
 }: ImageCaptureButtonProps) => {
   const handleScreenshotClick = async () => {
-    const image = await new Promise((resolve) => {
+    const imageBlob: Blob = await new Promise((resolve) => {
       window.parent.postMessage({ action: 'get-screenshot-image' }, '*')
       window.addEventListener('message', function (event) {
         const { action, payload } = event.data
@@ -17,8 +17,8 @@ const ImageCaptureButton = ({
         }
       })
     })
-    console.log('image', image)
-    addMessageDraftFile(image as Blob)
+
+    addMessageDraftFile(imageBlob as Blob)
   }
   return (
     <button
diff --git a/src/lib/getScreenshotImage.ts b/src/lib/getScreenshotImage.ts
index c32cac83..88ac55b2 100644
--- a/src/lib/getScreenshotImage.ts
+++ b/src/lib/getScreenshotImage.ts
@@ -1,3 +1,5 @@
+import html2canvas from 'html2canvas'
+
 /**
  * We use this function to
  * 1. Create a snipping tool view for the user to select the area of the screen
@@ -5,9 +7,9 @@
  * 3. Crop the image with the user's selection
  * 4. Return the cropped image as a blob
  */
-export const getScreenshotImage = async () => {
+export const getScreenshotImage = async (): Promise<Blob> => {
   // Create a snipping tool view for the user to select the area of the screen
-  const snipeRegion = document.createElement('div')
+  const snipeRegion: HTMLDivElement = document.createElement('div')
   snipeRegion.style.position = 'fixed'
   snipeRegion.style.top = '0'
   snipeRegion.style.left = '0'
@@ -16,68 +18,133 @@ export const getScreenshotImage = async () => {
   snipeRegion.style.zIndex = '2147483646' // Maximum z-index
   snipeRegion.style.backgroundColor = 'rgba(0, 0, 0, 0.5)'
   snipeRegion.style.cursor = 'crosshair'
-  document.body.appendChild(snipeRegion)
 
-  const snipeSelection = document.createElement('div')
+  const snipeSelection: HTMLDivElement = document.createElement('div')
   snipeSelection.style.position = 'fixed'
-  snipeSelection.style.top = '0'
-  snipeSelection.style.left = '0'
-  snipeSelection.style.width = '0'
-  snipeSelection.style.height = '0'
-  snipeSelection.style.border = '1px solid #ffffff0a'
-  snipeSelection.style.backgroundColor = 'rgba(256, 256, 256, 0.1)'
+  snipeSelection.style.border = '1px solid #fff'
+  snipeSelection.style.backgroundColor = 'rgba(255, 255, 255, 0.1)'
   snipeSelection.style.zIndex = '2147483647' // Maximum z-index
+
+  document.body.appendChild(snipeRegion)
   document.body.appendChild(snipeSelection)
 
-  // Grab the screen image with canvas
-  const canvas = document.createElement('canvas')
+  // Create a canvas element
+  const canvas: HTMLCanvasElement = document.createElement('canvas')
   canvas.width = window.innerWidth
   canvas.height = window.innerHeight
-  const ctx = canvas.getContext('2d')
+  const ctx: CanvasRenderingContext2D | null = canvas.getContext('2d')
   if (!ctx) {
     throw new Error('Could not get canvas context')
   }
 
-  // Crop the image with the user's selection
-  await new Promise((resolve) => {
-    snipeRegion.addEventListener('mousedown', (e) => {
-      const startX = e.clientX
-      const startY = e.clientY
-
-      snipeSelection.style.top = `${startY}px`
-      snipeSelection.style.left = `${startX}px`
-
-      // Update the snipping tool view when the user moves the mouse
-      document.addEventListener('mousemove', (e) => {
-        snipeSelection.style.width = `${Math.abs(e.clientX - startX)}px`
-        snipeSelection.style.height = `${Math.abs(e.clientY - startY)}px`
-      })
-
-      // Crop the image when the user releases the mouse
-      snipeRegion.addEventListener('mouseup', (e) => {
-        const endX = e.clientX
-        const endY = e.clientY
-        const width = endX - startX
-        const height = endY - startY
-        const imageData = ctx.getImageData(startX, startY, width, height)
-        ctx.putImageData(imageData, 0, 0)
-        resolve(null)
-      })
-    })
+  // Initially declare the variables with a type and set to undefined
+  let startX: number | undefined
+  let startY: number | undefined
+  let endX: number | undefined
+  let endY: number | undefined
+
+  // Wait for the user to make a selection
+  await new Promise<void>((resolve) => {
+    const onMouseMove = (e: MouseEvent) => {
+      if (startX === undefined || startY === undefined) return
+
+      const currentX: number = e.clientX
+      const currentY: number = e.clientY
+      const width: number = Math.abs(currentX - startX)
+      const height: number = Math.abs(currentY - startY)
+      snipeSelection.style.width = `${width}px`
+      snipeSelection.style.height = `${height}px`
+      snipeSelection.style.left = `${Math.min(startX, currentX)}px`
+      snipeSelection.style.top = `${Math.min(startY, currentY)}px`
+    }
+
+    const onMouseUp = (e: MouseEvent) => {
+      endX = e.clientX
+      endY = e.clientY
+      document.removeEventListener('mousemove', onMouseMove)
+      snipeRegion.removeEventListener('mouseup', onMouseUp)
+      document.body.removeChild(snipeRegion)
+      document.body.removeChild(snipeSelection)
+      resolve()
+    }
+
+    snipeRegion.addEventListener(
+      'mousedown',
+      (e: MouseEvent) => {
+        startX = e.clientX
+        startY = e.clientY
+        snipeSelection.style.left = `${startX}px`
+        snipeSelection.style.top = `${startY}px`
+
+        document.addEventListener('mousemove', onMouseMove)
+        document.addEventListener('mouseup', onMouseUp, { once: true })
+      },
+      { once: true },
+    )
   })
 
-  // Remove the snipping tool view
-  document.body.removeChild(snipeRegion)
-  document.body.removeChild(snipeSelection)
+  // Ensure that the coordinates are defined before using them
+  if (
+    typeof startX === 'undefined' ||
+    typeof startY === 'undefined' ||
+    typeof endX === 'undefined' ||
+    typeof endY === 'undefined'
+  ) {
+    throw new Error('Selection coordinates have not been defined.')
+  }
+
+  // Now we can safely use the variables as they have been assigned during the mouse events
+  const width: number = Math.abs(endX - startX)
+  const height: number = Math.abs(endY - startY)
+  const left: number = Math.min(startX, endX)
+  const top: number = Math.min(startY, endY)
 
-  // Return the cropped image as a blob
-  const blob = await new Promise<Blob | null>((resolve) => {
-    canvas.toBlob((blob) => {
-      resolve(blob)
+  // Use html2canvas to capture the content of the page
+  const screenshotCanvas: HTMLCanvasElement = await html2canvas(document.body, {
+    width: window.innerWidth,
+    height: window.innerHeight,
+    x: window.scrollX,
+    y: window.scrollY,
+    scale: 1,
+  })
+
+  // Create a cropped canvas as before
+  const croppedCanvas: HTMLCanvasElement = document.createElement('canvas')
+  croppedCanvas.width = width
+  croppedCanvas.height = height
+  const croppedCtx: CanvasRenderingContext2D | null =
+    croppedCanvas.getContext('2d')
+  if (!croppedCtx) {
+    throw new Error('Could not get cropped canvas context')
+  }
+
+  // Draw the captured area from the screenshotCanvas onto the cropped canvas
+  croppedCtx.drawImage(
+    screenshotCanvas,
+    left,
+    top,
+    width,
+    height,
+    0,
+    0,
+    width,
+    height,
+  )
+
+  // Convert the cropped canvas to a blob as before
+  const blob: Blob | null = await new Promise((resolve, reject) => {
+    croppedCanvas.toBlob((blob) => {
+      if (blob) {
+        resolve(blob)
+      } else {
+        reject(new Error('Blob conversion failed'))
+      }
     })
   })
+
   if (!blob) {
-    throw new Error('Could not get blob')
+    throw new Error('Blob is null')
   }
+
   return blob
 }
diff --git a/yarn.lock b/yarn.lock
index 4e1d6e5f..5ffe4037 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -1915,6 +1915,11 @@ base-64@^0.1.0:
   resolved "https://registry.yarnpkg.com/base-64/-/base-64-0.1.0.tgz#780a99c84e7d600260361511c4877613bf24f6bb"
   integrity sha512-Y5gU45svrR5tI2Vt/X9GPd3L0HNIKzGu202EjxrXMpuc2V2CiKgemAbUUsqYmZJvPtCXoUKjNZwBJzsNScUbXA==
 
+base64-arraybuffer@^1.0.2:
+  version "1.0.2"
+  resolved "https://registry.yarnpkg.com/base64-arraybuffer/-/base64-arraybuffer-1.0.2.tgz#1c37589a7c4b0746e34bd1feb951da2df01c1bdc"
+  integrity sha512-I3yl4r9QB5ZRY3XuJVEPfc2XhZO6YweFPI+UovAzn+8/hb3oJ6lnysaFcjVpkCPfVWFUDvoZ8kmVDP7WyRtYtQ==
+
 base64-js@^1.3.1, base64-js@^1.5.1:
   version "1.5.1"
   resolved "https://registry.yarnpkg.com/base64-js/-/base64-js-1.5.1.tgz#1b1b440160a5bf7ad40b650f095963481903930a"
@@ -2343,6 +2348,13 @@ crypt@0.0.2:
   resolved "https://registry.yarnpkg.com/crypt/-/crypt-0.0.2.tgz#88d7ff7ec0dfb86f713dc87bbb42d044d3e6c41b"
   integrity sha512-mCxBlsHFYh9C+HVpiEacem8FEBnMXgU9gy4zmNC+SXAZNB/1idgp/aulFJ4FgCi7GPEVbfyng092GqL2k2rmow==
 
+css-line-break@^2.1.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/css-line-break/-/css-line-break-2.1.0.tgz#bfef660dfa6f5397ea54116bb3cb4873edbc4fa0"
+  integrity sha512-FHcKFCZcAha3LwfVBhCQbW2nCNbkZXn7KVUJcsT5/P8YmfsVja0FMPJr0B903j/E69HUphKiV9iQArX8SDYA4w==
+  dependencies:
+    utrie "^1.0.2"
+
 css-select@^5.1.0:
   version "5.1.0"
   resolved "https://registry.yarnpkg.com/css-select/-/css-select-5.1.0.tgz#b8ebd6554c3637ccc76688804ad3f6a6fdaea8a6"
@@ -3340,6 +3352,14 @@ html-void-elements@^2.0.0:
   resolved "https://registry.yarnpkg.com/html-void-elements/-/html-void-elements-2.0.1.tgz#29459b8b05c200b6c5ee98743c41b979d577549f"
   integrity sha512-0quDb7s97CfemeJAnW9wC0hw78MtW7NU3hqtCD75g2vFlDLt36llsYD7uB7SUzojLMP24N5IatXf7ylGXiGG9A==
 
+html2canvas@^1.4.1:
+  version "1.4.1"
+  resolved "https://registry.yarnpkg.com/html2canvas/-/html2canvas-1.4.1.tgz#7cef1888311b5011d507794a066041b14669a543"
+  integrity sha512-fPU6BHNpsyIhr8yyMpTLLxAbkaK8ArIBcmZIRiBLiDhjeqvXolaEmDGmELFuX9I4xDcaKKcJl+TKZLqruBbmWA==
+  dependencies:
+    css-line-break "^2.1.0"
+    text-segmentation "^1.0.3"
+
 htmlparser2@^8.0.1:
   version "8.0.2"
   resolved "https://registry.yarnpkg.com/htmlparser2/-/htmlparser2-8.0.2.tgz#f002151705b383e62433b5cf466f5b716edaec21"
@@ -5984,6 +6004,13 @@ term-size@^2.1.0:
   resolved "https://registry.yarnpkg.com/term-size/-/term-size-2.2.1.tgz#2a6a54840432c2fb6320fea0f415531e90189f54"
   integrity sha512-wK0Ri4fOGjv/XPy8SBHZChl8CM7uMc5VML7SqiQ0zG7+J5Vr+RMQDoHa2CNT6KHUnTGIXH34UDMkPzAUyapBZg==
 
+text-segmentation@^1.0.3:
+  version "1.0.3"
+  resolved "https://registry.yarnpkg.com/text-segmentation/-/text-segmentation-1.0.3.tgz#52a388159efffe746b24a63ba311b6ac9f2d7943"
+  integrity sha512-iOiPUo/BGnZ6+54OsWxZidGCsdU8YbE4PSpdPinp7DeMtUJNJBoJ/ouUSTJjHkh1KntHaltHl/gDs2FC4i5+Nw==
+  dependencies:
+    utrie "^1.0.2"
+
 thenify-all@^1.0.0:
   version "1.6.0"
   resolved "https://registry.yarnpkg.com/thenify-all/-/thenify-all-1.6.0.tgz#1a1918d402d8fc3f98fbf234db0bcc8cc10e9726"
@@ -6225,6 +6252,13 @@ util-deprecate@^1.0.1, util-deprecate@^1.0.2, util-deprecate@~1.0.1:
   resolved "https://registry.yarnpkg.com/util-deprecate/-/util-deprecate-1.0.2.tgz#450d4dc9fa70de732762fbd2d4a28981419a0ccf"
   integrity sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==
 
+utrie@^1.0.2:
+  version "1.0.2"
+  resolved "https://registry.yarnpkg.com/utrie/-/utrie-1.0.2.tgz#d42fe44de9bc0119c25de7f564a6ed1b2c87a645"
+  integrity sha512-1MLa5ouZiOmQzUbjbu9VmjLzn1QLXBhwpUa7kdLUQK+KQ5KA9I1vk5U4YHe/X2Ch7PYnJfWuWT+VbuxbGwljhw==
+  dependencies:
+    base64-arraybuffer "^1.0.2"
+
 uuid@^9.0.0:
   version "9.0.1"
   resolved "https://registry.yarnpkg.com/uuid/-/uuid-9.0.1.tgz#e188d4c8853cc722220392c424cd637f32293f30"