server.js

const {
  StateGraph,
  Annotation,
  MemorySaver,
  START,
  END,
} = require('@langchain/langgraph');
const puppeteer = require('puppeteer');
const { chromium, firefox, webkit } = require('playwright');
const { Builder, Key } = require('selenium-webdriver');
const express = require('express');
const http = require('http');
const { Server } = require('socket.io');
const diff = require('diff');
const fs = require('fs');
const path = require('path');

const SCREENSHOTS_DIR = path.join(__dirname, 'screenshots');
if (!fs.existsSync(SCREENSHOTS_DIR)) {
  fs.mkdirSync(SCREENSHOTS_DIR, { recursive: true });
}

const app = express();
const server = http.createServer(app);
const io = new Server(server);

const PORT = process.env.PORT || 3000;

app.use(express.json());
app.use('/screenshots', express.static(SCREENSHOTS_DIR));

io.on('connection', (socket) => {
  console.log('Client connected');
  socket.on('disconnect', () => console.log('Client disconnected'));
});

const StateAnnotation = Annotation.Root({
  task: Annotation(),
  tool: Annotation(),
  browserEngine: Annotation(),
  url: Annotation(),
  editAction: Annotation(),
  beforeDom: Annotation(),
  afterDom: Annotation(),
  visuals: Annotation(),
  editsLog: Annotation(),
  analysis: Annotation(),
  networkLog: Annotation(),
  traces: Annotation(),
});

async function parseTask(state) {
  const task = state.task ?? '';
  const urlMatch = task.match(/on\s+(https?:\/\/[^\s,]+)/);
  const url = urlMatch ? urlMatch[1].replace(/,+$/, '') : 'https://example.com';
  const editMatch = task.match(/edit\s+(.+?)\s*,/);
  const editAction = editMatch ? editMatch[1] : 'placeholder text';
  const toolMatch = task.match(/use\s+(\w+)/);
  const tool = toolMatch ? toolMatch[1].toLowerCase() : 'playwright';
  const browserMatch = task.match(/browser\s+(chromium|firefox|webkit)/);
  const browserEngine = browserMatch ? browserMatch[1].toLowerCase() : 'chromium';
  io.emit('update', { stage: 'parse', message: `Parsed: ${url}, Tool: ${tool}, Browser: ${browserEngine}` });
  return { url, editAction, tool, browserEngine };
}

async function executeBrowser(state) {
  const { url, editAction, tool } = state;
  let browser;
  let driver;
  let beforeDom = '';
  let afterDom = '';
  const visuals = [];
  const networkLog = [];
  const traces = [];

  try {
    if (tool === 'playwright') {
      const playwrightBrowsers = { chromium, firefox, webkit };
      const selectedBrowser = playwrightBrowsers[state.browserEngine] || chromium;
      browser = await selectedBrowser.launch({ headless: process.env.HEADLESS !== '0' });
      const VIDEOS_DIR = path.join(SCREENSHOTS_DIR, 'videos');
      if (!fs.existsSync(VIDEOS_DIR)) fs.mkdirSync(VIDEOS_DIR, { recursive: true });
      const context = await browser.newContext({ recordVideo: { dir: VIDEOS_DIR } });
      await context.tracing.start({ screenshots: true, snapshots: true });
      const page = await context.newPage();
      await page.route('**/*', async (route) => {
        networkLog.push({ url: route.request().url(), method: route.request().method() });
        await route.continue();
      });
      await page.goto(url, { waitUntil: 'networkidle' });
      beforeDom = await page.evaluate(() => document.documentElement.outerHTML);
      await page.screenshot({ path: path.join(SCREENSHOTS_DIR, 'before.png') });
      const inputLocator = page.locator('input').first();
      if (await inputLocator.count() > 0) {
        await inputLocator.evaluate((el, edit) => { el.placeholder = `Edited: ${edit}`; }, editAction);
      }
      await page.keyboard.down('Control');
      await page.keyboard.press('KeyA');
      await page.keyboard.up('Control');
      afterDom = await page.evaluate(() => document.documentElement.outerHTML);
      await page.screenshot({ path: path.join(SCREENSHOTS_DIR, 'after.png') });
      visuals.push('screenshots/before.png', 'screenshots/after.png');
      const tracePath = path.join(SCREENSHOTS_DIR, 'trace.zip');
      await context.tracing.stop({ path: tracePath });
      traces.push('screenshots/trace.zip');
    } else if (tool === 'puppeteer') {
      browser = await puppeteer.launch({ headless: process.env.HEADLESS !== '0' });
      const page = await browser.newPage();
      await page.goto(url, { waitUntil: 'networkidle2' });
      beforeDom = await page.evaluate(() => document.documentElement.outerHTML);
      await page.screenshot({ path: path.join(SCREENSHOTS_DIR, 'before.png') });
      await page.evaluate((edit) => {
        const input = document.querySelector('input');
        if (input) input.placeholder = `Edited: ${edit}`;
      }, editAction);
      await page.keyboard.down('Control');
      await page.keyboard.press('A');
      await page.keyboard.up('Control');
      afterDom = await page.evaluate(() => document.documentElement.outerHTML);
      await page.screenshot({ path: path.join(SCREENSHOTS_DIR, 'after.png') });
      visuals.push('screenshots/before.png', 'screenshots/after.png');
    } else if (tool === 'selenium') {
      driver = await new Builder().forBrowser('chrome').build();
      await driver.get(url);
      beforeDom = await driver.executeScript('return document.documentElement.outerHTML');
      await driver.takeScreenshot().then((data) =>
        fs.writeFileSync(path.join(SCREENSHOTS_DIR, 'before.png'), Buffer.from(data, 'base64'))
      );
      await driver.executeScript(
        `const el = document.querySelector('input'); if (el) el.placeholder = 'Edited: ${editAction.replace(/'/g, "\\'")}';`
      );
      const actions = driver.actions();
      await actions.keyDown(Key.CONTROL).sendKeys('a').keyUp(Key.CONTROL).perform();
      afterDom = await driver.executeScript('return document.documentElement.outerHTML');
      await driver.takeScreenshot().then((data) =>
        fs.writeFileSync(path.join(SCREENSHOTS_DIR, 'after.png'), Buffer.from(data, 'base64'))
      );
      visuals.push('screenshots/before.png', 'screenshots/after.png');
    }
    io.emit('update', { stage: 'execute', visuals, networkLog, traces });
  } finally {
    if (browser) await browser.close();
    if (driver) await driver.quit();
  }

  return { beforeDom, afterDom, visuals, networkLog, traces };
}

async function analyzeResults(state) {
  const { beforeDom, afterDom } = state;
  const editsLog = diff.diffLines(beforeDom ?? '', afterDom ?? '');
  const analysis = {
    dom_changes: editsLog.filter((part) => part.added || part.removed),
    performance: {},
    suggestions: ['Enhance with AI selectors'],
    network_requests: state.networkLog ?? [],
    traces: state.traces ?? [],
  };
  io.emit('update', { stage: 'analyze', analysis });
  console.log('Session logged for billing');
  return { editsLog, analysis };
}

const graph = new StateGraph(StateAnnotation)
  .addNode('parseTask', parseTask)
  .addNode('executeBrowser', executeBrowser)
  .addNode('analyzeResults', analyzeResults)
  .addEdge(START, 'parseTask')
  .addEdge('parseTask', 'executeBrowser')
  .addEdge('executeBrowser', 'analyzeResults')
  .addEdge('analyzeResults', END);

const memory = new MemorySaver();
const runnableGraph = graph.compile({ checkpointer: memory });

app.post('/analyze', async (req, res) => {
  const { task } = req.body;
  if (!task) return res.status(400).json({ error: 'Task required' });

  try {
    const config = { configurable: { thread_id: 'dom-thread' } };
    const result = await runnableGraph.invoke({ task }, config);
    res.json(result);
  } catch (error) {
    res.status(500).json({ error: error.message });
  }
});

server.listen(PORT, () =>
  console.log(`DOMAnalyzer Swarm v4.0 (LangGraph) on port ${PORT}`)
);