import { createServer } from "node:http";
import { mkdtemp, readFile, rm, stat } from "node:fs/promises";
import { spawn } from "node:child_process";
import { createRequire } from "node:module";
import { tmpdir } from "node:os";
import { resolve, join, extname } from "node:path";
import { inflateSync } from "node:zlib";

const webDir = resolve(process.argv[2] ?? "zig-out/web");
const screenshotPath = resolve(process.argv[3] ?? "zig-out/browser-replay.png");

const mime = new Map([
  [".html", "text/html; charset=utf-8"],
  [".js", "text/javascript; charset=utf-8"],
  [".wasm", "application/wasm"],
  [".png", "image/png"],
]);

function run(cmd, args) {
  return new Promise((resolvePromise, reject) => {
    const child = spawn(cmd, args, { stdio: "inherit" });
    child.on("error", reject);
    child.on("exit", (code) => {
      if (code === 0) resolvePromise();
      else reject(new Error(`${cmd} ${args.join(" ")} exited with ${code}`));
    });
  });
}

async function loadPlaywright() {
  try {
    return { api: await import("playwright"), cleanup: async () => {} };
  } catch (err) {
    if (process.env.PARDES_BROWSER_SKIP_INSTALL === "1") throw err;
  }

  const dir = await mkdtemp(join(tmpdir(), "pardes-playwright-"));
  try {
    await run("npm", ["--prefix", dir, "install", "--silent", "playwright"]);
    await run("npm", ["--prefix", dir, "exec", "--", "playwright", "install", "chromium"]);
    const require = createRequire(join(dir, "package.json"));
    return {
      api: require("playwright"),
      cleanup: async () => {
        await rm(dir, { recursive: true, force: true });
      },
    };
  } catch (err) {
    await rm(dir, { recursive: true, force: true });
    throw err;
  }
}

async function serve(root) {
  await stat(join(root, "replay.html"));
  await stat(join(root, "pardes-replay.js"));

  const server = createServer(async (req, res) => {
    try {
      const url = new URL(req.url ?? "/", "http://127.0.0.1");
      const path = url.pathname === "/" ? "/replay.html" : url.pathname;
      if (path.includes("..")) throw new Error("bad path");
      const file = join(root, path.slice(1));
      const body = await readFile(file);
      res.writeHead(200, {
        "content-type": mime.get(extname(file)) ?? "application/octet-stream",
        "cache-control": "no-store",
      });
      res.end(body);
    } catch {
      res.writeHead(404);
      res.end("not found");
    }
  });

  await new Promise((resolvePromise) => server.listen(0, "127.0.0.1", resolvePromise));
  const address = server.address();
  if (!address || typeof address === "string") throw new Error("failed to bind test server");
  return { server, url: `http://127.0.0.1:${address.port}/replay.html` };
}

function parsePng(buf) {
  const sig = "89504e470d0a1a0a";
  if (buf.subarray(0, 8).toString("hex") !== sig) throw new Error("not a png");

  let offset = 8;
  let width = 0;
  let height = 0;
  let bitDepth = 0;
  let colorType = 0;
  const idat = [];
  while (offset < buf.length) {
    const len = buf.readUInt32BE(offset);
    const type = buf.subarray(offset + 4, offset + 8).toString("ascii");
    const data = buf.subarray(offset + 8, offset + 8 + len);
    offset += 12 + len;
    if (type === "IHDR") {
      width = data.readUInt32BE(0);
      height = data.readUInt32BE(4);
      bitDepth = data[8];
      colorType = data[9];
    } else if (type === "IDAT") {
      idat.push(data);
    } else if (type === "IEND") {
      break;
    }
  }
  if (bitDepth !== 8) throw new Error(`unsupported png bit depth ${bitDepth}`);
  const channels = colorType === 6 ? 4 : colorType === 2 ? 3 : colorType === 0 ? 1 : 0;
  if (channels === 0) throw new Error(`unsupported png color type ${colorType}`);

  const raw = inflateSync(Buffer.concat(idat));
  const stride = width * channels;
  const rows = Buffer.alloc(width * height * channels);
  let rawOffset = 0;
  for (let y = 0; y < height; y += 1) {
    const filter = raw[rawOffset];
    rawOffset += 1;
    const rowStart = y * stride;
    for (let x = 0; x < stride; x += 1) {
      const left = x >= channels ? rows[rowStart + x - channels] : 0;
      const up = y > 0 ? rows[rowStart + x - stride] : 0;
      const upLeft = y > 0 && x >= channels ? rows[rowStart + x - stride - channels] : 0;
      const value = raw[rawOffset + x];
      let recon;
      if (filter === 0) recon = value;
      else if (filter === 1) recon = value + left;
      else if (filter === 2) recon = value + up;
      else if (filter === 3) recon = value + Math.floor((left + up) / 2);
      else if (filter === 4) {
        const p = left + up - upLeft;
        const pa = Math.abs(p - left);
        const pb = Math.abs(p - up);
        const pc = Math.abs(p - upLeft);
        recon = value + (pa <= pb && pa <= pc ? left : pb <= pc ? up : upLeft);
      } else {
        throw new Error(`unsupported png filter ${filter}`);
      }
      rows[rowStart + x] = recon & 0xff;
    }
    rawOffset += stride;
  }
  return { width, height, channels, rows };
}

function validateScreenshot(png) {
  let bright = 0;
  let colored = 0;
  let nonDark = 0;
  const pixels = png.width * png.height;
  for (let i = 0; i < pixels; i += 1) {
    const off = i * png.channels;
    const r = png.rows[off];
    const g = png.channels === 1 ? r : png.rows[off + 1];
    const b = png.channels === 1 ? r : png.rows[off + 2];
    if (r > 90 || g > 90 || b > 90) bright += 1;
    if (Math.max(r, g, b) - Math.min(r, g, b) > 20) colored += 1;
    if (r > 35 || g > 35 || b > 35) nonDark += 1;
  }
  if (png.width < 640 || png.height < 360) {
    throw new Error(`screenshot too small: ${png.width}x${png.height}`);
  }
  if (bright < 500 || nonDark < 2000) {
    throw new Error(`browser replay appears blank: bright=${bright} nonDark=${nonDark} colored=${colored}`);
  }
  console.log(`browser replay screenshot ok: ${png.width}x${png.height}, bright=${bright}, nonDark=${nonDark}, colored=${colored}`);
}

const { server, url } = await serve(webDir);
let cleanupPlaywright = async () => {};
let browser;
try {
  const { api, cleanup } = await loadPlaywright();
  cleanupPlaywright = cleanup;
  const { chromium } = api;

  browser = await chromium.launch({
    headless: true,
    args: ["--use-angle=swiftshader", "--enable-unsafe-swiftshader"],
  });
  const page = await browser.newPage({ viewport: { width: 1280, height: 720 } });
  const consoleMessages = [];
  const pageErrors = [];
  page.on("console", (msg) => consoleMessages.push(`${msg.type()}: ${msg.text()}`));
  page.on("pageerror", (err) => pageErrors.push(err.stack ?? err.message));

  await page.addInitScript(() => {
    window.__pardesReplayE2E = { draws: 0, contextLost: 0 };
    window.addEventListener("webglcontextlost", () => {
      window.__pardesReplayE2E.contextLost += 1;
    }, true);

    const wrapDraws = (target) => {
      if (!target) return;
      for (const name of ["drawArraysInstanced", "drawArrays", "drawElements", "drawElementsInstanced"]) {
        if (typeof target[name] !== "function" || target[name].__pardesReplayE2E) continue;
        const oldDraw = target[name];
        target[name] = function(...args) {
          window.__pardesReplayE2E.draws += 1;
          return oldDraw.apply(this, args);
        };
        target[name].__pardesReplayE2E = true;
      }
    };

    wrapDraws(globalThis.WebGLRenderingContext?.prototype);
    wrapDraws(globalThis.WebGL2RenderingContext?.prototype);

    const oldGetContext = HTMLCanvasElement.prototype.getContext;
    HTMLCanvasElement.prototype.getContext = function(type, attrs) {
      const ctx = oldGetContext.call(this, type, attrs);
      if (ctx && !ctx.__pardesReplayE2E && (type === "webgl2" || type === "webgl" || type === "experimental-webgl")) {
        ctx.__pardesReplayE2E = true;
        wrapDraws(ctx);
      }
      return ctx;
    };
  });

  await page.goto(url, { waitUntil: "load" });
  await page.waitForSelector("canvas", { timeout: 5000 });
  await page.waitForFunction(() => window.__pardesReplayE2E?.draws >= 2, null, { timeout: 5000 });
  await page.waitForTimeout(250);
  const trace = await page.evaluate(() => window.__pardesReplayE2E);
  if (trace.contextLost !== 0) {
    throw new Error(`browser WebGL context was lost during replay: ${JSON.stringify(trace)}`);
  }
  if (pageErrors.length !== 0) {
    throw new Error(`browser page error: ${pageErrors[0]}`);
  }
  const lostMessage = consoleMessages.find((msg) => msg.includes("CONTEXT_LOST_WEBGL"));
  if (lostMessage) throw new Error(`browser WebGL context lost: ${lostMessage}`);

  await page.screenshot({ path: screenshotPath, fullPage: true });
  validateScreenshot(parsePng(await readFile(screenshotPath)));
} finally {
  if (browser) await browser.close();
  await cleanupPlaywright();
  server.close();
}