bolt.diy/app/components/workbench/ScreenshotSelector.tsx

import { memo, useCallback, useState } from 'react';
import { toast } from 'react-toastify';

interface ScreenshotSelectorProps {
  isSelectionMode: boolean;
  setIsSelectionMode: (mode: boolean) => void;
  containerRef: React.RefObject<HTMLElement>;
}

export const ScreenshotSelector = memo(
  ({ isSelectionMode, setIsSelectionMode, containerRef }: ScreenshotSelectorProps) => {
    const [isCapturing, setIsCapturing] = useState(false);
    const [selectionStart, setSelectionStart] = useState<{ x: number; y: number } | null>(null);
    const [selectionEnd, setSelectionEnd] = useState<{ x: number; y: number } | null>(null);

    const handleCopySelection = useCallback(async () => {
      if (!isSelectionMode || !selectionStart || !selectionEnd || !containerRef.current) return;

      setIsCapturing(true);
      try {
        const width = Math.abs(selectionEnd.x - selectionStart.x);
        const height = Math.abs(selectionEnd.y - selectionStart.y);

        // Create a video element to capture the screen
        const video = document.createElement('video');
        video.style.opacity = '0';
        document.body.appendChild(video);

        try {
          // Capture the entire screen
          const stream = await navigator.mediaDevices.getDisplayMedia({
            audio: false,
            video: {
              displaySurface: 'window',
            },
          } as MediaStreamConstraints);

          // Set up video with the stream
          video.srcObject = stream;
          await video.play();

          // Wait for video to be ready
          await new Promise((resolve) => setTimeout(resolve, 300));

          // Create temporary canvas for full screenshot
          const tempCanvas = document.createElement('canvas');
          tempCanvas.width = video.videoWidth;
          tempCanvas.height = video.videoHeight;
          const tempCtx = tempCanvas.getContext('2d');

          if (!tempCtx) {
            throw new Error('Failed to get temporary canvas context');
          }

          // Draw the full video frame
          tempCtx.drawImage(video, 0, 0);

          // Get the container's position in the page
          const containerRect = containerRef.current.getBoundingClientRect();

          // Calculate scale factor between video and screen
          const scaleX = video.videoWidth / window.innerWidth;
          const scaleY = video.videoHeight / window.innerHeight;

          // Calculate the scaled coordinates
          const scaledX = (containerRect.left + Math.min(selectionStart.x, selectionEnd.x)) * scaleX;
          const scaledY = (containerRect.top + Math.min(selectionStart.y, selectionEnd.y)) * scaleY;
          const scaledWidth = width * scaleX;
          const scaledHeight = height * scaleY;

          // Create final canvas for the cropped area
          const canvas = document.createElement('canvas');
          canvas.width = width;
          canvas.height = height;
          const ctx = canvas.getContext('2d');

          if (!ctx) {
            throw new Error('Failed to get canvas context');
          }

          // Draw the cropped area
          ctx.drawImage(tempCanvas, scaledX, scaledY, scaledWidth, scaledHeight, 0, 0, width, height);

          // Convert to blob
          const blob = await new Promise<Blob>((resolve, reject) => {
            canvas.toBlob((blob) => {
              if (blob) resolve(blob);
              else reject(new Error('Failed to create blob'));
            }, 'image/png');
          });

          // Create a FileReader to convert blob to base64
          const reader = new FileReader();
          reader.onload = (e) => {
            const base64Image = e.target?.result as string;

            // Find the textarea element
            const textarea = document.querySelector('textarea');
            if (textarea) {
              // Get the setters from the BaseChat component
              const setUploadedFiles = (window as any).__BOLT_SET_UPLOADED_FILES__;
              const setImageDataList = (window as any).__BOLT_SET_IMAGE_DATA_LIST__;
              const uploadedFiles = (window as any).__BOLT_UPLOADED_FILES__ || [];
              const imageDataList = (window as any).__BOLT_IMAGE_DATA_LIST__ || [];

              if (setUploadedFiles && setImageDataList) {
                // Update the files and image data
                const file = new File([blob], 'screenshot.png', { type: 'image/png' });
                setUploadedFiles([...uploadedFiles, file]);
                setImageDataList([...imageDataList, base64Image]);
                toast.success('Screenshot captured and added to chat');
              } else {
                toast.error('Could not add screenshot to chat');
              }
            }
          };
          reader.readAsDataURL(blob);

          // Stop all tracks
          stream.getTracks().forEach((track) => track.stop());
        } finally {
          // Clean up video element
          document.body.removeChild(video);
        }
      } catch (error) {
        console.error('Failed to capture screenshot:', error);
        toast.error('Failed to capture screenshot');
      } finally {
        setIsCapturing(false);
        setSelectionStart(null);
        setSelectionEnd(null);
        setIsSelectionMode(false);
      }
    }, [isSelectionMode, selectionStart, selectionEnd, containerRef, setIsSelectionMode]);

    const handleSelectionStart = useCallback(
      (e: React.MouseEvent) => {
        e.preventDefault();
        e.stopPropagation();
        if (!isSelectionMode) return;
        const rect = e.currentTarget.getBoundingClientRect();
        const x = e.clientX - rect.left;
        const y = e.clientY - rect.top;
        setSelectionStart({ x, y });
        setSelectionEnd({ x, y });
      },
      [isSelectionMode],
    );

    const handleSelectionMove = useCallback(
      (e: React.MouseEvent) => {
        e.preventDefault();
        e.stopPropagation();
        if (!isSelectionMode || !selectionStart) return;
        const rect = e.currentTarget.getBoundingClientRect();
        const x = e.clientX - rect.left;
        const y = e.clientY - rect.top;
        setSelectionEnd({ x, y });
      },
      [isSelectionMode, selectionStart],
    );

    if (!isSelectionMode) return null;

    return (
      <div
        className="absolute inset-0 cursor-crosshair"
        onMouseDown={handleSelectionStart}
        onMouseMove={handleSelectionMove}
        onMouseUp={handleCopySelection}
        onMouseLeave={() => {
          if (selectionStart) {
            setSelectionStart(null);
          }
        }}
        style={{
          backgroundColor: isCapturing ? 'transparent' : 'rgba(0, 0, 0, 0.1)',
          userSelect: 'none',
          WebkitUserSelect: 'none',
          pointerEvents: 'all',
          opacity: isCapturing ? 0 : 1,
          zIndex: 50,
          transition: 'opacity 0.1s ease-in-out',
        }}
      >
        {selectionStart && selectionEnd && !isCapturing && (
          <div
            className="absolute border-2 border-blue-500 bg-blue-200 bg-opacity-20"
            style={{
              left: Math.min(selectionStart.x, selectionEnd.x),
              top: Math.min(selectionStart.y, selectionEnd.y),
              width: Math.abs(selectionEnd.x - selectionStart.x),
              height: Math.abs(selectionEnd.y - selectionStart.y),
            }}
          />
        )}
      </div>
    );
  },
);