demo-platform/apps/api/src/rasterizer.ts

import { chromium, type Browser } from "playwright-core";
import { createHash } from "node:crypto";
import { mkdirSync, writeFileSync } from "node:fs";
import { join } from "node:path";
import { pool } from "./db.js";

const DATA_DIR = process.env.DATA_DIR ?? "/data";
const EXEC =
  process.env.CHROMIUM_PATH ??
  ["/usr/bin/chromium-browser", "/usr/bin/chromium"].find(Boolean)!;

let browser: Browser | null = null;
let launching: Promise<Browser> | null = null;

async function getBrowser(): Promise<Browser> {
  if (browser?.isConnected()) return browser;
  if (!launching) {
    launching = (async () => {
      const candidates = process.env.CHROMIUM_PATH
        ? [process.env.CHROMIUM_PATH]
        : ["/usr/bin/chromium-browser", "/usr/bin/chromium"];
      let lastErr: unknown;
      for (const p of candidates) {
        try {
          browser = await chromium.launch({
            executablePath: p,
            args: ["--no-sandbox", "--disable-dev-shm-usage", "--disable-gpu"],
          });
          return browser;
        } catch (e) { lastErr = e; }
      }
      throw lastErr;
    })().finally(() => { launching = null; });
  }
  return launching;
}

async function storeAsset(buf: Buffer, mime: string): Promise<string> {
  const sha = createHash("sha256").update(buf).digest("hex");
  const existing = await pool.query("SELECT id FROM assets WHERE sha256=$1", [sha]);
  if (existing.rowCount) return existing.rows[0].id;
  const dir = join(DATA_DIR, "assets", sha.slice(0, 2));
  mkdirSync(dir, { recursive: true });
  const path = join(dir, sha);
  writeFileSync(path, buf);
  const r = await pool.query(
    "INSERT INTO assets (kind, path, mime, bytes, sha256) VALUES ('image',$1,$2,$3,$4) RETURNING id",
    [path, mime, buf.length, sha],
  );
  return r.rows[0].id;
}

/**
 * R4: rasterize a dom step to a PNG asset (1280x800 viewport), cached on the
 * step. dom_html is immutable post-creation, so the cache never invalidates.
 * The page runs with JavaScript disabled — snapshots are inert documents.
 */
export async function rasterizeStep(stepId: string): Promise<string | null> {
  const st = await pool.query(
    "SELECT snapshot_type, dom_html, raster_asset_id, asset_id FROM steps WHERE id=$1",
    [stepId],
  );
  if (!st.rowCount) return null;
  const row = st.rows[0];
  if (row.snapshot_type === "image") return row.asset_id;
  if (row.raster_asset_id) return row.raster_asset_id;
  if (!row.dom_html) return null;

  const b = await getBrowser();
  const ctx = await b.newContext({
    viewport: { width: 1280, height: 800 },
    javaScriptEnabled: false,
    baseURL: process.env.PUBLIC_BASE_URL ?? `http://127.0.0.1:${process.env.PORT ?? 3000}`,
  });
  try {
    const page = await ctx.newPage();
    // resolve /assets/... references against our own server
    const base = process.env.PUBLIC_BASE_URL ?? `http://127.0.0.1:${process.env.PORT ?? 3000}`;
    const html = row.dom_html.replace(/(src|href)="\/assets\//g, `$1="${base}/assets/`);
    await page.setContent(html, { waitUntil: "networkidle", timeout: 20_000 }).catch(() => {});
    const buf = await page.screenshot({ type: "png" });
    const assetId = await storeAsset(Buffer.from(buf), "image/png");
    await pool.query("UPDATE steps SET raster_asset_id=$2 WHERE id=$1", [stepId, assetId]);
    return assetId;
  } finally {
    await ctx.close();
  }
}

export async function closeRasterizer(): Promise<void> {
  await browser?.close().catch(() => {});
  browser = null;
}