perfect-postcode/screenshot/src/screenshot.ts

import { chromium, type Browser, type BrowserContext, type Page, type Route } from 'playwright';
import { existsSync, readdirSync } from 'fs';
import { NetworkCache } from './network-cache.js';

const VIEWPORT = { width: 1200, height: 630 };
const NAVIGATION_TIMEOUT = 15_000;
const READY_TIMEOUT = 15_000;
const RENDER_BUFFER_MS = 200;
const POOL_SIZE = 3;

let browser: Browser | null = null;
let sharedContext: BrowserContext | null = null;
const pagePool: Page[] = [];
let warmingUp = false;
const networkCache = new NetworkCache();

function detectGpu(): boolean {
  try {
    if (existsSync('/dev/dri')) {
      const devices = readdirSync('/dev/dri');
      const hasGpu = devices.some((d) => d.startsWith('renderD') || d.startsWith('card'));
      if (hasGpu) {
        console.log(`GPU detected: /dev/dri contains ${devices.join(', ')}`);
        return true;
      }
    }
  } catch {
    // Fall through to software rendering
  }
  console.log('No GPU detected, using SwiftShader software rendering');
  return false;
}

const hasGpu = detectGpu();

function getBrowserArgs(): string[] {
  const baseArgs = [
    '--no-sandbox',
    '--disable-dev-shm-usage',
    '--enable-webgl',
    '--ignore-gpu-blocklist',
    // Prevent timer/renderer throttling in headless mode
    '--disable-background-timer-throttling',
    '--disable-backgrounding-occluded-windows',
    '--disable-renderer-backgrounding',
    // Ensure compositor finishes all stages before screenshot capture
    '--run-all-compositor-stages-before-draw',
    // Skip color profile conversion overhead
    '--force-color-profile=srgb',
    // Reduce non-essential browser features
    '--disable-domain-reliability',
    '--disable-features=TranslateUI',
  ];

  if (hasGpu) {
    return [...baseArgs, '--enable-gpu', '--use-gl=egl'];
  } else {
    return [
      ...baseArgs,
      '--disable-gpu',
      '--use-gl=swiftshader',
      // Run the GPU (SwiftShader) work in the browser process to eliminate
      // IPC overhead between browser and GPU processes — significant win
      // when all GL calls are CPU-bound anyway
      '--in-process-gpu',
    ];
  }
}

async function ensureBrowser(): Promise<Browser> {
  if (!browser || !browser.isConnected()) {
    const args = getBrowserArgs();
    console.log(`Launching browser with args: ${args.join(' ')}`);
    browser = await chromium.launch({ args });
  }
  return browser;
}

async function ensureContext(): Promise<BrowserContext> {
  if (sharedContext) return sharedContext;

  const instance = await ensureBrowser();
  sharedContext = await instance.newContext({
    viewport: VIEWPORT,
    deviceScaleFactor: 1,
  });

  // Set up response caching at context level — all pages share the cache.
  // This means tiles fetched during pre-warm or a previous screenshot are
  // served instantly for subsequent screenshots.
  await sharedContext.route('**/*', async (route: Route) => {
    const url = route.request().url();

    // Non-cacheable requests (API data, etc.) pass through directly
    if (!networkCache.shouldCache(url)) {
      await route.continue();
      return;
    }

    // Cache hit — fulfill from memory, zero network
    const cached = networkCache.get(url);
    if (cached) {
      await route.fulfill({
        status: cached.status,
        headers: cached.headers,
        body: cached.body,
      });
      return;
    }

    // Cache miss — fetch, cache, fulfill
    try {
      const response = await route.fetch();
      const body = await response.body();
      const entry = {
        body,
        headers: response.headers(),
        status: response.status(),
      };
      networkCache.set(url, entry);
      await route.fulfill({
        status: entry.status,
        headers: entry.headers,
        body: entry.body,
      });
    } catch {
      await route.continue().catch(() => {});
    }
  });

  return sharedContext;
}

async function createPage(): Promise<Page> {
  const context = await ensureContext();
  return context.newPage();
}

async function warmPool(): Promise<void> {
  if (warmingUp) return;
  warmingUp = true;
  try {
    while (pagePool.length < POOL_SIZE) {
      const page = await createPage();
      pagePool.push(page);
    }
  } finally {
    warmingUp = false;
  }
}

async function acquirePage(): Promise<Page> {
  const page = pagePool.shift();
  if (page && !page.isClosed()) {
    warmPool().catch(() => {});
    return page;
  }
  const newPage = await createPage();
  warmPool().catch(() => {});
  return newPage;
}

async function releasePage(page: Page): Promise<void> {
  try {
    if (page.isClosed()) return;
    // Navigate to blank to free rendered page resources (DOM, WebGL context)
    // while keeping the page object alive for V8 code cache reuse
    await page.goto('about:blank', { timeout: 5000 }).catch(() => {});
    if (!page.isClosed() && pagePool.length < POOL_SIZE) {
      pagePool.push(page);
    } else {
      await page.close().catch(() => {});
    }
  } catch {
    await page.close().catch(() => {});
  }
}

/**
 * Pre-warm the browser and populate the network cache by loading the app once.
 * Subsequent screenshots benefit from cached JS/CSS bundles, map tiles,
 * and V8 compiled bytecode — eliminating cold-start latency.
 */
export async function initialize(appUrl: string): Promise<void> {
  console.log('Pre-warming browser and caches...');
  const page = await createPage();
  try {
    await page.goto(`${appUrl}/?screenshot=1`, {
      waitUntil: 'load',
      timeout: 30_000,
    });
    // Wait for the app to fully load and cache all resources
    try {
      await page.waitForFunction('window.__screenshot_ready === true', {
        timeout: READY_TIMEOUT,
      });
    } catch {
      // Non-fatal — cache will still have JS/CSS/tiles from the partial load
    }
    console.log(`Pre-warm complete. Cache: ${networkCache.stats()}`);
  } catch (err) {
    console.warn('Pre-warm failed (non-fatal):', err);
  } finally {
    await page.close().catch(() => {});
  }
  await warmPool();
}

export async function takeScreenshot(url: string): Promise<Buffer> {
  const page = await acquirePage();
  const t0 = performance.now();

  try {
    const response = await page.goto(url, {
      waitUntil: 'domcontentloaded',
      timeout: NAVIGATION_TIMEOUT,
    });
    const t1 = performance.now();
    if (response) {
      console.log(`  Navigate: ${(t1 - t0).toFixed(0)}ms (status ${response.status()})`);
    }

    // Wait for the frontend to signal that data is loaded and layers created
    try {
      await page.waitForFunction('window.__screenshot_ready === true', {
        timeout: READY_TIMEOUT,
      });
    } catch {
      console.warn('  Timed out waiting for __screenshot_ready');
    }
    const t2 = performance.now();
    console.log(`  Ready: ${(t2 - t1).toFixed(0)}ms`);

    // Brief buffer for SwiftShader to finish rendering the WebGL frame.
    // Reduced from 500ms → 200ms since tiles now load from the in-memory
    // cache and don't need network round-trips.
    await page.waitForTimeout(RENDER_BUFFER_MS);

    // JPEG at quality 85: ~3-5x faster encoding than PNG with negligible
    // visual difference for screenshot content (map + text overlay)
    const screenshot = await page.screenshot({ type: 'jpeg', quality: 85 });
    const t3 = performance.now();
    console.log(`  Capture: ${(t3 - t2).toFixed(0)}ms | Total: ${(t3 - t0).toFixed(0)}ms`);
    console.log(`  Cache: ${networkCache.stats()}`);

    return Buffer.from(screenshot);
  } finally {
    await releasePage(page);
  }
}

export async function checkWebGL(): Promise<Record<string, unknown>> {
  const page = await acquirePage();
  try {
    await page.setContent('<canvas id="c" width="1" height="1"></canvas>');
    const info = await page.evaluate(() => {
      const canvas = document.getElementById('c') as HTMLCanvasElement;
      const gl =
        canvas.getContext('webgl2') || canvas.getContext('webgl') || canvas.getContext('experimental-webgl');
      if (!gl) return { webgl: false, error: 'No WebGL context available' };
      const g = gl as WebGLRenderingContext;
      const debugExt = g.getExtension('WEBGL_debug_renderer_info');
      return {
        webgl: true,
        version: g.getParameter(g.VERSION),
        renderer: debugExt ? g.getParameter(debugExt.UNMASKED_RENDERER_WEBGL) : 'unknown',
        vendor: debugExt ? g.getParameter(debugExt.UNMASKED_VENDOR_WEBGL) : 'unknown',
      };
    });
    return { ...info, cache: networkCache.stats() };
  } finally {
    await releasePage(page);
  }
}

export async function closeBrowser(): Promise<void> {
  for (const page of pagePool) {
    await page.close().catch(() => {});
  }
  pagePool.length = 0;

  if (sharedContext) {
    await sharedContext.close().catch(() => {});
    sharedContext = null;
  }

  if (browser) {
    await browser.close().catch(() => {});
    browser = null;
  }
}