perfect-postcode/screenshot/src/server.ts

import express, { type Request, type Response } from "express";
import { ScreenshotCache } from "./cache.js";
import {
  takeScreenshot,
  checkWebGL,
  closeBrowser,
  initialize,
} from "./screenshot.js";
import { buildScreenshotRequest, ValidationError } from "./validation.js";

const PORT = parseRequiredPositiveIntEnv("PORT");
const APP_URL = process.env.APP_URL;
const CACHE_DIR = process.env.CACHE_DIR;
const SCREENSHOT_CONCURRENCY = parseRequiredPositiveIntEnv(
  "SCREENSHOT_CONCURRENCY",
);
const RATE_LIMIT_WINDOW_MS = parseRequiredPositiveIntEnv(
  "SCREENSHOT_RATE_WINDOW_MS",
);
const RATE_LIMIT_MAX = parseRequiredPositiveIntEnv("SCREENSHOT_RATE_LIMIT");

if (!APP_URL) {
  console.error("Error: APP_URL environment variable is required");
  process.exit(1);
}

const CACHE_ENABLED = parseOptionalBoolEnv(
  "SCREENSHOT_CACHE_ENABLED",
  !isDevelopmentAppUrl(APP_URL),
);

if (CACHE_ENABLED && !CACHE_DIR) {
  console.error("Error: CACHE_DIR environment variable is required");
  process.exit(1);
}

const cache = CACHE_ENABLED ? new ScreenshotCache(CACHE_DIR as string) : null;
const app = express();
app.set("trust proxy", true);

let activeScreenshots = 0;
let lastRateLimitPrune = 0;
const rateLimitBuckets = new Map<string, { count: number; resetAt: number }>();
type ReleaseScreenshotSlot = () => void;
type PendingScreenshotSlot = {
  resolve: (release: ReleaseScreenshotSlot | null) => void;
  cleanup: () => void;
};
const screenshotSlotQueue: PendingScreenshotSlot[] = [];

function parseRequiredPositiveIntEnv(name: string): number {
  const raw = process.env[name];
  if (!raw) {
    throw new Error(`${name} environment variable is required`);
  }
  const value = Number.parseInt(raw, 10);
  if (!Number.isFinite(value) || value <= 0) {
    throw new Error(`${name} must be a positive integer`);
  }
  return value;
}

function parseOptionalBoolEnv(name: string, defaultValue: boolean): boolean {
  const raw = process.env[name];
  if (raw == null || raw === "") return defaultValue;
  if (raw === "1" || raw.toLowerCase() === "true") return true;
  if (raw === "0" || raw.toLowerCase() === "false") return false;
  throw new Error(`${name} must be true or false`);
}

function isDevelopmentAppUrl(rawUrl: string): boolean {
  try {
    const url = new URL(rawUrl);
    return (
      url.hostname === "localhost" ||
      url.hostname === "127.0.0.1" ||
      url.hostname === "frontend" ||
      url.port === "3001"
    );
  } catch {
    return false;
  }
}

function grantScreenshotSlot(): ReleaseScreenshotSlot {
  activeScreenshots += 1;
  let released = false;
  return () => {
    if (released) return;
    released = true;
    activeScreenshots = Math.max(0, activeScreenshots - 1);
    drainScreenshotSlotQueue();
  };
}

function drainScreenshotSlotQueue(): void {
  while (
    activeScreenshots < SCREENSHOT_CONCURRENCY &&
    screenshotSlotQueue.length > 0
  ) {
    const pending = screenshotSlotQueue.shift();
    if (!pending) return;
    pending.cleanup();
    pending.resolve(grantScreenshotSlot());
  }
}

function acquireScreenshotSlot(
  res: Response,
): Promise<ReleaseScreenshotSlot | null> {
  if (activeScreenshots < SCREENSHOT_CONCURRENCY) {
    return Promise.resolve(grantScreenshotSlot());
  }

  return new Promise((resolve) => {
    let pending: PendingScreenshotSlot;
    const onClose = () => {
      if (res.writableEnded) return;
      pending.cleanup();
      const index = screenshotSlotQueue.indexOf(pending);
      if (index !== -1) screenshotSlotQueue.splice(index, 1);
      resolve(null);
    };

    pending = {
      resolve,
      cleanup: () => res.off("close", onClose),
    };
    res.on("close", onClose);
    screenshotSlotQueue.push(pending);
    console.log(
      `Queued screenshot request; queue length: ${screenshotSlotQueue.length}`,
    );
  });
}

function rateLimitKey(req: Request): string {
  const forwardedFor = req.get("x-forwarded-for")?.split(",")[0]?.trim();
  const key = forwardedFor || req.ip || req.socket.remoteAddress;
  if (!key) {
    throw new Error("Unable to determine request IP for rate limiting");
  }
  return key;
}

function allowScreenshotRequest(req: Request): boolean {
  const now = Date.now();
  if (now - lastRateLimitPrune > RATE_LIMIT_WINDOW_MS) {
    lastRateLimitPrune = now;
    for (const [key, bucket] of rateLimitBuckets) {
      if (bucket.resetAt <= now) {
        rateLimitBuckets.delete(key);
      }
    }
  }

  const key = rateLimitKey(req);
  let bucket = rateLimitBuckets.get(key);
  if (!bucket || bucket.resetAt <= now) {
    bucket = { count: 0, resetAt: now + RATE_LIMIT_WINDOW_MS };
    rateLimitBuckets.set(key, bucket);
  }
  if (bucket.count >= RATE_LIMIT_MAX) {
    return false;
  }
  bucket.count += 1;
  return true;
}

app.get("/health", (_req, res) => {
  res.status(200).send("ok");
});

app.get("/debug", async (_req, res) => {
  try {
    const info = await checkWebGL();
    res.json(info);
  } catch (err) {
    res.status(500).json({ error: String(err) });
  }
});

app.get("/screenshot", async (req, res) => {
  let releaseSlot: (() => void) | null = null;
  try {
    const { pagePath, qs } = buildScreenshotRequest(
      req.query as Record<string, unknown>,
    );
    const authHeader = req.headers.authorization;

    let cacheKey: string | null = null;
    if (cache) {
      const cacheParams = new URLSearchParams(qs);
      if (pagePath !== "/") cacheParams.set("path", pagePath);
      // Include auth status in cache key so authenticated screenshots
      // (with hexagons outside free zone) are cached separately
      if (authHeader) cacheParams.set("_auth", "1");
      cacheKey = cache.buildKey(cacheParams);

      // Check cache first
      const cached = cache.get(cacheKey);
      if (cached) {
        res.setHeader("Content-Type", "image/jpeg");
        res.setHeader("Cache-Control", "public, max-age=86400");
        res.setHeader("X-Cache", "HIT");
        cached.pipe(res);
        return;
      }
    }

    if (!allowScreenshotRequest(req)) {
      res.status(429).json({ error: "Screenshot rate limit exceeded" });
      return;
    }

    releaseSlot = await acquireScreenshotSlot(res);
    if (!releaseSlot) {
      return;
    }

    // Build the URL for the frontend in screenshot mode
    qs.set("screenshot", "1");
    const url = `${APP_URL}${pagePath}?${qs}`;

    console.log(
      `Taking screenshot: ${url}${authHeader ? " (authenticated)" : ""}`,
    );
    const jpeg = await takeScreenshot(url, authHeader);

    // Cache it
    if (cache && cacheKey) cache.set(cacheKey, jpeg);

    res.setHeader("Content-Type", "image/jpeg");
    res.setHeader("Cache-Control", cache ? "public, max-age=86400" : "no-store");
    res.setHeader("X-Cache", cache ? "MISS" : "BYPASS");
    res.send(jpeg);
  } catch (err) {
    if (err instanceof ValidationError) {
      res.status(err.status).json({ error: err.message });
      return;
    }
    console.error("Screenshot failed:", err);
    res.status(500).json({ error: "Screenshot failed" });
  } finally {
    releaseSlot?.();
  }
});

const server = app.listen(PORT, () => {
  console.log(`Screenshot service listening on port ${PORT}`);
  console.log(`  APP_URL: ${APP_URL}`);
  console.log(`  CACHE_ENABLED: ${CACHE_ENABLED}`);
  if (cache) console.log(`  CACHE_DIR: ${CACHE_DIR}`);
  console.log(`  SCREENSHOT_CONCURRENCY: ${SCREENSHOT_CONCURRENCY}`);
  console.log(
    `  SCREENSHOT_RATE_LIMIT: ${RATE_LIMIT_MAX}/${RATE_LIMIT_WINDOW_MS}ms`,
  );

  // Pre-warm browser and populate network cache in background.
  // The health endpoint is available immediately; screenshot requests
  // during warm-up will still work (just slower on the first call).
  initialize(APP_URL).catch((err) => {
    console.error("Initialization failed:", err);
  });
});

// Graceful shutdown
for (const signal of ["SIGTERM", "SIGINT"]) {
  process.on(signal, async () => {
    console.log(`Received ${signal}, shutting down...`);
    server.close();
    await closeBrowser();
    process.exit(0);
  });
}