AI-Grading-Assistant/electron/main/screenshot.ts

import { BrowserWindow, desktopCapturer, screen } from 'electron'
import * as path from 'path'
import * as fs from 'fs'
import log from 'electron-log'

let screenshotWindow: BrowserWindow | null = null
let capturedImage: Buffer | null = null
let capturedBase64: string | null = null
let isScreenshotActive = false
let lastScreenshotTime = 0
const SCREENSHOT_COOLDOWN = 2000

export function startScreenshot(mainWin: BrowserWindow): void {
  const now = Date.now()
  if (now - lastScreenshotTime < SCREENSHOT_COOLDOWN) {
    log.warn('Screenshot on cooldown, ignoring')
    return
  }
  if (isScreenshotActive) {
    log.warn('Screenshot already active, ignoring')
    return
  }
  isScreenshotActive = true
  lastScreenshotTime = now
  require('./shortcut').disableShortcut()
  mainWin.hide()

  const cursor = screen.getCursorScreenPoint()
  const displays = screen.getAllDisplays()
  const targetDisplay = displays.find(d => {
    return cursor.x >= d.bounds.x && cursor.x <= d.bounds.x + d.bounds.width &&
           cursor.y >= d.bounds.y && cursor.y <= d.bounds.y + d.bounds.height
  }) || displays[0]

  const bounds = targetDisplay.bounds

  screenshotWindow = new BrowserWindow({
    x: bounds.x,
    y: bounds.y,
    width: bounds.width,
    height: bounds.height,
    frame: false,
    transparent: true,
    alwaysOnTop: true,
    skipTaskbar: true,
    resizable: false,
    webPreferences: {
      preload: path.join(__dirname, '../preload/preload.js'),
      contextIsolation: true,
      nodeIntegration: false
    }
  })

  if (process.env.NODE_ENV === 'development' || process.argv.includes('--dev')) {
    screenshotWindow.loadURL(`http://localhost:5173/#/screenshot`)
  } else {
    screenshotWindow.loadFile(path.join(__dirname, '../../dist/index.html'), { hash: '/screenshot' })
  }

  screenshotWindow.on('closed', () => {
    isScreenshotActive = false
    screenshotWindow = null
  })
}

export async function captureScreen(region: { x: number; y: number; width: number; height: number }): Promise<string | null> {
  try {
    const displays = screen.getAllDisplays()
    const sources = await desktopCapturer.getSources({
      types: ['screen'],
      thumbnailSize: { width: 4096, height: 2160 }
    })

    if (sources.length === 0) return null

    const cursor = screen.getCursorScreenPoint()
    const targetDisplay = displays.find(d => {
      return cursor.x >= d.bounds.x && cursor.x <= d.bounds.x + d.bounds.width &&
             cursor.y >= d.bounds.y && cursor.y <= d.bounds.y + d.bounds.height
    }) || displays[0]

    const sourceIndex = displays.findIndex(d => d.id === targetDisplay.id)
    const source = sources[Math.min(sourceIndex, sources.length - 1)]

    const fullImage = source.thumbnail
    const displayWidth = targetDisplay.bounds.width
    const displayHeight = targetDisplay.bounds.height
    const scaleX = fullImage.getSize().width / displayWidth
    const scaleY = fullImage.getSize().height / displayHeight

    const cropX = Math.round(region.x * scaleX)
    const cropY = Math.round(region.y * scaleY)
    const cropW = Math.round(region.width * scaleX)
    const cropH = Math.round(region.height * scaleY)

    const cropped = fullImage.crop({
      x: Math.max(0, cropX),
      y: Math.max(0, cropY),
      width: Math.min(cropW, fullImage.getSize().width - cropX),
      height: Math.min(cropH, fullImage.getSize().height - cropY)
    })

    const pngBuffer = cropped.toPNG()
    const base64 = pngBuffer.toString('base64')
    capturedImage = pngBuffer
    capturedBase64 = base64

    try {
      const Store = require('electron-store').default
      const store = new Store()
      const savePath = store.get('screenshotSavePath', '')
      if (savePath) {
        const timestamp = Date.now()
        const filename = `screenshot_${timestamp}.png`
        const fullPath = path.join(savePath, filename)
        fs.writeFileSync(fullPath, pngBuffer)
        log.info('Screenshot saved to', fullPath)
      }
    } catch (e) {
      log.error('Failed to save screenshot:', e)
    }

    log.info('Screenshot captured', { width: region.width, height: region.height })
    return base64
  } catch (err) {
    log.error('Screenshot capture error:', err)
    return null
  }
}

export function closeScreenshot(): void {
  if (screenshotWindow && !screenshotWindow.isDestroyed()) {
    screenshotWindow.close()
    if (!screenshotWindow.isDestroyed()) {
      screenshotWindow.destroy()
    }
  }
  screenshotWindow = null
  isScreenshotActive = false
}

export function getCapturedImage(): Buffer | null {
  return capturedImage
}

export function getCapturedBase64(): string | null {
  return capturedBase64
}

export function clearCapturedImage(): void {
  capturedImage = null
  capturedBase64 = null
}