main.js

const { app, BrowserWindow, ipcMain, desktopCapturer, screen } = require('electron');
const { exec, execFile } = require('child_process');
const path = require('path');
const fs = require('fs');

let mainWindow;
let overlayWindow;

const pythonExe = '..\\.venv\\Scripts\\python.exe';

let ocrInProgress = false; 

function runLatexOCR(imagePath, callback) {
  if (ocrInProgress) {
    console.log('OCR already in progress, skipping...');
    callback(null);
    return;
  }

  ocrInProgress = true;
  const scriptPath = path.join(__dirname, 'ocr', 'ocr.py');

  if (!fs.existsSync(scriptPath)) {
    console.error('OCR script not found at:', scriptPath);
    ocrInProgress = false;
    callback(null);
    return;
  }

  if (!fs.existsSync(imagePath)) {
    console.error('Image file not found at:', imagePath);
    ocrInProgress = false;
    callback(null);
    return;
  }

  if (!fs.existsSync(pythonExe)) {
    console.error('Python executable not found at:', pythonExe);
    ocrInProgress = false;
    callback(null);
    return;
  }

  const stats = fs.statSync(imagePath);
  console.log('🔍 DEBUG: Image file size:', stats.size, 'bytes');
  console.log('🔍 DEBUG: Executing Python script...');
  
  const timeout = setTimeout(() => {
    console.error('❌ OCR timeout - killing process');
    ocrInProgress = false;
    callback(null);
  }, 30000);

  execFile(pythonExe, [scriptPath, imagePath], { timeout: 25000 }, (error, stdout, stderr) => {
    clearTimeout(timeout);
    ocrInProgress = false; 
    console.log('🔍 DEBUG: Python execution completed');
    
    if (error) {
      console.error('❌ OCR Error:', error.message);
      console.error('❌ Error code:', error.code);

      if (mainWindow?.webContents) {
        mainWindow.webContents.send('ocr-error', {
          message: 'OCR processing failed. Please try a different area or ensure the image contains clear text/math.',
          errorCode: error.code
        });
      }
      
      callback(null);
      return;
    }
    
    if (stderr && !stderr.includes('Pydantic serializer warnings') && !stderr.includes('Expected `dict')) {
      console.warn('⚠️ OCR Warning:', stderr);
    }
    
    if (!stdout || stdout.trim().length === 0) {
      console.error('❌ OCR returned empty result');

      if (mainWindow?.webContents) {
        mainWindow.webContents.send('ocr-error', {
          message: 'No text/math detected in the selected area. Please try selecting a larger area with clearer content.',
          errorCode: 'EMPTY_RESULT'
        });
      }
      
      callback(null);
      return;
    }
    
    console.log('✅ OCR Success - Result:', stdout.trim());
    callback(stdout.trim());
  });
}

function createWindow() {
  mainWindow = new BrowserWindow({
    width: 1000,
    height: 600,
    webPreferences: {
      preload: path.join(__dirname, 'preload.js'),
    },
  });

  const isDev = !app.isPackaged;
  const devUrl = 'http://localhost:5173';
  const prodUrl = path.join(__dirname, 'dist/index.html');

  if (isDev) {
    mainWindow.loadURL(devUrl);
  } else {
    mainWindow.loadFile(prodUrl);
  }
}

app.whenReady().then(() => {
  exec('wsl.exe', (err) => {
    if (err) console.error('⚠️ Failed to start WSL:', err);
  });

  createWindow();

  ipcMain.handle('start-area-screenshot', () => {
    const { width, height } = screen.getPrimaryDisplay().bounds;

    overlayWindow = new BrowserWindow({
      width,
      height,
      transparent: true,
      frame: false,
      alwaysOnTop: true,
      fullscreen: true,
      skipTaskbar: true,
      webPreferences: {
        preload: path.join(__dirname, 'preload.js'),
        nodeIntegration: true,
        contextIsolation: false,
      },
    });

    overlayWindow.loadFile(path.join(__dirname, 'public', 'overlay.html'));
  });

  ipcMain.on('cancel-screenshot', () => {
    console.log('Screenshot cancelled by user');
    if (overlayWindow) {
      overlayWindow.close();
      overlayWindow = null;
    }
    
    if (mainWindow?.webContents) {
      mainWindow.webContents.send('screenshot-cancelled');
    }
  });

  ipcMain.on('area-selected', async (event, rect) => {
    if (overlayWindow) {
      overlayWindow.close();
      overlayWindow = null;
    }

    try {
      const display = screen.getPrimaryDisplay();
      const dpr = rect.dpr || display.scaleFactor || 1;
      const fullW = Math.round(display.bounds.width * dpr);
      const fullH = Math.round(display.bounds.height * dpr);

      const sources = await desktopCapturer.getSources({
        types: ['screen'],
        thumbnailSize: { width: fullW, height: fullH },
      });

      const screenSource = sources[0];
      if (!screenSource) {
        console.error('No screen source found');
        if (mainWindow?.webContents) {
          mainWindow.webContents.send('ocr-error', {
            message: 'Failed to capture screen. Please try again.',
            errorCode: 'NO_SCREEN_SOURCE'
          });
        }
        return;
      }


      let cropped;
      try {
        cropped = screenSource.thumbnail.crop({
          x: Math.round(rect.x * dpr),
          y: Math.round(rect.y * dpr),
          width: Math.round(rect.width * dpr),
          height: Math.round(rect.height * dpr),
        });
      } catch (e) {
        console.warn('Crop failed, using full screenshot');
        cropped = screenSource.thumbnail;
      }

      const buffer = cropped.toPNG();
      const tempImagePath = path.join(app.getPath('temp'), 'ocr_capture.png');
      
      console.log('🔍 DEBUG: Saving image to:', tempImagePath);
      console.log('🔍 DEBUG: Buffer size:', buffer.length, 'bytes');
      
      try {
        fs.writeFileSync(tempImagePath, buffer);
        console.log('Image saved successfully');
        
        if (fs.existsSync(tempImagePath)) {
          const stats = fs.statSync(tempImagePath);
          console.log('🔍 DEBUG: Saved file size:', stats.size, 'bytes');
        } else {
          console.error('Failed to save image file');
          if (mainWindow?.webContents) {
            mainWindow.webContents.send('ocr-error', {
              message: 'Failed to save screenshot. Please try again.',
              errorCode: 'SAVE_FAILED'
            });
          }
          return;
        }
      } catch (saveError) {
        console.error('Error saving image:', saveError);
        if (mainWindow?.webContents) {
          mainWindow.webContents.send('ocr-error', {
            message: 'Failed to save screenshot. Please try again.',
            errorCode: 'SAVE_ERROR'
          });
        }
        return;
      }

      runLatexOCR(tempImagePath, (latex) => {
        if (latex) {
          console.log('✅ Area screenshot OCR completed successfully');
          if (mainWindow?.webContents) {
            mainWindow.webContents.send('screenshot-captured', {
              img: `data:image/png;base64,${buffer.toString('base64')}`,
              latex,
            });
          } else {
            console.error('mainWindow not ready');
          }
        } else {
          console.error('Area screenshot OCR failed');
          // CRITICAL: Always send an event to reset loading state
          if (mainWindow?.webContents) {
            // If no error was already sent, send a generic one
            mainWindow.webContents.send('ocr-error', {
              message: 'No text or mathematical content detected in the selected area. Please try selecting a different area.',
              errorCode: 'NO_CONTENT_DETECTED'
            });
          }
        }
      });
    } catch (error) {
      console.error('Error during area selection processing:', error);
      if (mainWindow?.webContents) {
        mainWindow.webContents.send('ocr-error', {
          message: 'An unexpected error occurred. Please try again.',
          errorCode: 'PROCESSING_ERROR'
        });
      }
    }
  });

  ipcMain.handle('capture-area', async () => {
    const sources = await desktopCapturer.getSources({ types: ['screen'] });
    const screenSource = sources[0];
    if (!screenSource) return null;
    return screenSource.thumbnail.toDataURL();
  });

  app.on('activate', () => {
    if (BrowserWindow.getAllWindows().length === 0) createWindow();
  });
});

app.on('window-all-closed', () => {
  if (process.platform !== 'darwin') app.quit();
});