import { execFileSync } from 'node:child_process';
import { readFile } from 'node:fs/promises';

import puppeteer from 'puppeteer-core';

const chromePath = process.env.CHROME_PATH;
const targetUrl = process.env.PROFILE_URL ?? 'http://127.0.0.1:4173/';
const label = process.env.PROFILE_LABEL ?? 'profile';
const steadySeconds = Number(process.env.STEADY_SECONDS ?? 12);
const memorySeconds = Number(process.env.MEMORY_SECONDS ?? 30);
const settleSeconds = Number(process.env.SETTLE_SECONDS ?? 0);

if (!chromePath) {
  throw new Error('Set CHROME_PATH to a Chromium/Chrome executable.');
}

const delay = (milliseconds) => new Promise((resolve) => setTimeout(resolve, milliseconds));

function traceSummary(trace) {
  const threadNames = new Map();

  for (const event of trace.traceEvents) {
    if (event.ph === 'M' && event.name === 'thread_name') {
      threadNames.set(`${event.pid}:${event.tid}`, event.args?.name);
    }
  }

  const mainThreads = new Set(
    [...threadNames.entries()]
      .filter(([, name]) => name === 'CrRendererMain')
      .map(([key]) => key),
  );
  const completeMainEvents = trace.traceEvents.filter(
    (event) =>
      event.ph === 'X' &&
      typeof event.dur === 'number' &&
      mainThreads.has(`${event.pid}:${event.tid}`),
  );

  const durationFor = (names) =>
    completeMainEvents
      .filter((event) => names.has(event.name))
      .reduce((total, event) => total + event.dur / 1000, 0);
  const countFor = (names) =>
    completeMainEvents.filter((event) => names.has(event.name)).length;

  const styleNames = new Set(['RecalculateStyles', 'UpdateLayoutTree']);
  const layoutNames = new Set(['Layout']);
  const paintNames = new Set(['Paint']);
  const prePaintNames = new Set(['PrePaint']);
  const runTasks = completeMainEvents.filter((event) => event.name === 'RunTask');
  const taskCandidates = runTasks.length
    ? runTasks
    : completeMainEvents.filter((event) => event.name.includes('RunTask'));
  const longTasks = taskCandidates
    .filter((event) => event.dur > 50_000)
    .map((event) => ({ name: event.name, durationMs: event.dur / 1000 }));

  return {
    mainThread: {
      style: { count: countFor(styleNames), durationMs: durationFor(styleNames) },
      layout: { count: countFor(layoutNames), durationMs: durationFor(layoutNames) },
      paint: { count: countFor(paintNames), durationMs: durationFor(paintNames) },
      prePaint: { count: countFor(prePaintNames), durationMs: durationFor(prePaintNames) },
    },
    longTasks,
    animationFrameCallbacks: {
      requested: completeMainEvents.filter((event) => event.name === 'RequestAnimationFrame').length,
      fired: completeMainEvents.filter((event) => event.name === 'FireAnimationFrame').length,
    },
  };
}

function rendererProcesses(browserPid) {
  const processTable = execFileSync('ps', ['-eo', 'pid=,ppid=,args='], { encoding: 'utf8' })
    .trim()
    .split('\n')
    .map((line) => {
      const match = line.trim().match(/^(\d+)\s+(\d+)\s+(.*)$/);
      return match ? { pid: Number(match[1]), ppid: Number(match[2]), args: match[3] } : null;
    })
    .filter(Boolean);
  const descendants = new Set([browserPid]);
  let changed = true;

  while (changed) {
    changed = false;
    for (const process of processTable) {
      if (descendants.has(process.ppid) && !descendants.has(process.pid)) {
        descendants.add(process.pid);
        changed = true;
      }
    }
  }

  return processTable.filter(
    (process) => descendants.has(process.pid) && process.args.includes('--type=renderer'),
  );
}

async function processCpuTicks(processes) {
  const samples = [];
  for (const process of processes) {
    try {
      const stat = (await readFile(`/proc/${process.pid}/stat`, 'utf8')).trim().split(' ');
      samples.push({
        pid: process.pid,
        ticks: Number(stat[13]) + Number(stat[14]),
        rssPages: Number(stat[23]),
      });
    } catch {
      // A renderer can exit between process discovery and sampling.
    }
  }
  return samples;
}

function cpuDelta(before, after, elapsedSeconds) {
  const ticksPerSecond = Number(execFileSync('getconf', ['CLK_TCK'], { encoding: 'utf8' }).trim());
  const afterByPid = new Map(after.map((sample) => [sample.pid, sample]));
  const totalTicks = before.reduce((total, sample) => {
    const end = afterByPid.get(sample.pid);
    return total + (end ? Math.max(0, end.ticks - sample.ticks) : 0);
  }, 0);
  return {
    rendererPids: before.map((sample) => sample.pid),
    cpuTimeMs: (totalTicks / ticksPerSecond) * 1000,
    cpuPercentOfOneCore: (totalTicks / ticksPerSecond / elapsedSeconds) * 100,
  };
}

async function recordTrace(page, path, action) {
  await page.tracing.start({
    path,
    categories: [
      '-*',
      'blink',
      'cc',
      'devtools.timeline',
      'disabled-by-default-devtools.timeline',
      'disabled-by-default-devtools.timeline.frame',
      'toplevel',
      'viz',
    ],
  });
  await action();
  await page.tracing.stop();
  return traceSummary(JSON.parse(await readFile(path, 'utf8')));
}

const browser = await puppeteer.launch({
  executablePath: chromePath,
  headless: true,
  args: [
    '--no-sandbox',
    '--disable-dev-shm-usage',
    '--disable-background-timer-throttling',
    '--disable-backgrounding-occluded-windows',
    '--disable-renderer-backgrounding',
    '--enable-precise-memory-info',
  ],
});

try {
  const page = await browser.newPage();
  await page.setViewport({ width: 390, height: 844, deviceScaleFactor: 2, isMobile: true });
  const layerSession = await page.createCDPSession();
  let lastLayers = [];
  layerSession.on('LayerTree.layerTreeDidChange', ({ layers }) => {
    lastLayers = layers ?? [];
  });
  await layerSession.send('LayerTree.enable');

  const loadTrace = await recordTrace(page, `/tmp/${label}-load-trace.json`, async () => {
    await page.goto(targetUrl, { waitUntil: 'networkidle0' });
    await delay(3000);
  });

  const paintEntries = await page.evaluate(() =>
    performance.getEntriesByType('paint').map(({ name, startTime }) => ({ name, startTime })),
  );
  const animationsAfterLoad = await page.evaluate(() =>
    document.getAnimations().map((animation) => ({
      animationName: animation.animationName,
      playState: animation.playState,
      iterations: Number.isFinite(animation.effect?.getTiming().iterations)
        ? animation.effect?.getTiming().iterations
        : 'Infinity',
    })),
  );

  await delay(settleSeconds * 1000);
  const animationsAtSteadyStart = await page.evaluate(() =>
    document.getAnimations().map((animation) => ({
      animationName: animation.animationName,
      playState: animation.playState,
      iterations: Number.isFinite(animation.effect?.getTiming().iterations)
        ? animation.effect?.getTiming().iterations
        : 'Infinity',
    })),
  );

  const processes = rendererProcesses(browser.process().pid);
  const cpuBefore = await processCpuTicks(processes);
  const steadyTrace = await recordTrace(page, `/tmp/${label}-steady-trace.json`, async () => {
    await delay(steadySeconds * 1000);
  });
  const cpuAfter = await processCpuTicks(processes);

  await layerSession.send('HeapProfiler.collectGarbage');
  const heapSamples = [];
  const heapStart = Date.now();
  while ((Date.now() - heapStart) / 1000 <= memorySeconds) {
    const { usedSize, totalSize } = await layerSession.send('Runtime.getHeapUsage');
    heapSamples.push({ elapsedSeconds: (Date.now() - heapStart) / 1000, usedSize, totalSize });
    await delay(Math.min(5000, Math.max(0, memorySeconds * 1000 - (Date.now() - heapStart))));
    if ((Date.now() - heapStart) / 1000 >= memorySeconds) break;
  }
  await layerSession.send('HeapProfiler.collectGarbage');
  const heapEnd = await layerSession.send('Runtime.getHeapUsage');
  heapSamples.push({ elapsedSeconds: (Date.now() - heapStart) / 1000, ...heapEnd });

  const styleLayoutPaintMs =
    steadyTrace.mainThread.style.durationMs +
    steadyTrace.mainThread.layout.durationMs +
    steadyTrace.mainThread.paint.durationMs;
  const result = {
    targetUrl,
    chromiumVersion: await browser.version(),
    viewport: { width: 390, height: 844, deviceScaleFactor: 2, mobile: true },
    loadWindowSeconds: 3,
    settleWindowSeconds: settleSeconds,
    steadyWindowSeconds: steadySeconds,
    memoryWindowSeconds: memorySeconds,
    paintEntries,
    animationsAfterLoad,
    animationsAtSteadyStart,
    loadTrace,
    steadyTrace: {
      ...steadyTrace,
      styleLayoutPaintMsPerSecond: styleLayoutPaintMs / steadySeconds,
    },
    layers: {
      total: lastLayers.length,
      drawsContent: lastLayers.filter((layer) => layer.drawsContent).length,
    },
    rendererCpu: cpuDelta(cpuBefore, cpuAfter, steadySeconds),
    heap: {
      samples: heapSamples,
      startUsedBytes: heapSamples[0].usedSize,
      endUsedBytes: heapEnd.usedSize,
      growthBytes: heapEnd.usedSize - heapSamples[0].usedSize,
      minUsedBytes: Math.min(...heapSamples.map((sample) => sample.usedSize)),
      maxUsedBytes: Math.max(...heapSamples.map((sample) => sample.usedSize)),
    },
  };
  console.log(JSON.stringify(result, null, 2));
} finally {
  await browser.close();
}