import { mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { dirname, resolve } from "node:path";

const root = process.cwd();
const analyzeDataDirectory = resolve(
  root,
  process.env.BUNDLE_ANALYZE_DIR ?? ".next/diagnostics/analyze/data",
);
const outputPath = resolve(
  root,
  process.env.PERF_REPORT_PATH ??
    ".next/diagnostics/bundle-summary.json",
);
const routes = parseList(
  process.env.PERF_ROUTES ?? "/splash,/chat,/subscription,/tip",
);

const auditTargets = [
  {
    name: "Stripe React SDK",
    matches: (path) => path.includes("/node_modules/@stripe/react-stripe-js/"),
  },
  {
    name: "Stripe loader",
    matches: (path) => path.includes("/node_modules/@stripe/stripe-js/"),
  },
  {
    name: "Stripe dialog",
    matches: (path) => path.includes("/stripe-payment-dialog.tsx"),
  },
  {
    name: "Reply animation",
    matches: (path) => path.includes("/lottie-message-bubble"),
  },
  {
    name: "Lottie runtime",
    matches: (path) => path.includes("/node_modules/lottie-react/"),
  },
];

function main() {
  const modulesData = new FramedDataFile(
    resolve(analyzeDataDirectory, "modules.data"),
  );
  const asyncBoundariesByPath = getAsyncBoundariesByModulePath(modulesData);
  const routeReports = routes.map((route) =>
    analyzeRoute(route, asyncBoundariesByPath),
  );
  const packageJson = JSON.parse(
    readFileSync(resolve(root, "package.json"), "utf8"),
  );

  const report = {
    generatedAt: new Date().toISOString(),
    nextVersion: packageJson.dependencies.next,
    routes: routeReports,
  };

  mkdirSync(dirname(outputPath), { recursive: true });
  writeFileSync(outputPath, `${JSON.stringify(report, null, 2)}\n`);

  console.table(
    routeReports.map((route) => ({
      route: route.route,
      "eager client JS": formatBytes(route.eagerClientJs.compressedBytes),
      "async client JS": formatBytes(route.asyncClientJs.compressedBytes),
      "total client JS": formatBytes(route.totalClientJs.compressedBytes),
    })),
  );

  for (const route of routeReports) {
    const presentAudits = route.audits.filter(
      (audit) => audit.phase !== "absent",
    );
    if (presentAudits.length === 0) continue;
    console.log(`\n${route.route} audited modules`);
    console.table(
      presentAudits.map((audit) => ({
        module: audit.name,
        phase: audit.phase,
        compressed: formatBytes(audit.compressedBytes),
      })),
    );
  }

  console.log(`\nBundle summary written to ${outputPath}`);
}

function analyzeRoute(route, boundariesByPath) {
  const routeFile =
    route === "/"
      ? resolve(analyzeDataDirectory, "analyze.data")
      : resolve(
          analyzeDataDirectory,
          route.replace(/^\//, ""),
          "analyze.data",
        );
  const data = new FramedDataFile(routeFile);
  const sourceIndexes = collectRouteSourceIndexes(data);
  const clientParts = collectClientJavaScriptParts(data, sourceIndexes);
  const eagerParts = [];
  const asyncParts = [];

  for (const partIndex of clientParts) {
    const part = data.header.chunk_parts[partIndex];
    const sourcePath = normalizeSourcePath(
      data.getFullSourcePath(part.source_index),
    );
    const boundaryCount = boundariesByPath.get(sourcePath);
    if (boundaryCount !== undefined && boundaryCount > 0) {
      asyncParts.push(partIndex);
    } else {
      eagerParts.push(partIndex);
    }
  }

  return {
    route,
    eagerClientJs: sumParts(data, eagerParts),
    asyncClientJs: sumParts(data, asyncParts),
    totalClientJs: sumParts(data, clientParts),
    audits: auditTargets.map((target) => {
      const matchingParts = clientParts.filter((partIndex) => {
        const part = data.header.chunk_parts[partIndex];
        return target.matches(data.getFullSourcePath(part.source_index));
      });
      const boundaryCounts = matchingParts
        .map((partIndex) => {
          const part = data.header.chunk_parts[partIndex];
          return boundariesByPath.get(
            normalizeSourcePath(data.getFullSourcePath(part.source_index)),
          );
        })
        .filter((value) => value !== undefined);
      const sizes = sumParts(data, matchingParts);

      return {
        name: target.name,
        phase:
          matchingParts.length === 0
            ? "absent"
            : boundaryCounts.length > 0 && Math.min(...boundaryCounts) > 0
              ? "async"
              : "eager",
        ...sizes,
      };
    }),
  };
}

function collectRouteSourceIndexes(data) {
  const indexes = new Set();
  const pending = [...data.header.source_roots];

  while (pending.length > 0) {
    const index = pending.pop();
    if (index === undefined || indexes.has(index)) continue;
    indexes.add(index);
    pending.push(...data.readEdges(data.header.source_children, index));
  }

  return indexes;
}

function collectClientJavaScriptParts(data, sourceIndexes) {
  const parts = new Set();

  for (const sourceIndex of sourceIndexes) {
    for (const partIndex of data.readEdges(
      data.header.source_chunk_parts,
      sourceIndex,
    )) {
      const part = data.header.chunk_parts[partIndex];
      const output = data.header.output_files[part.output_file_index];
      if (
        output.filename.startsWith("[client-fs]/") &&
        output.filename.endsWith(".js")
      ) {
        parts.add(partIndex);
      }
    }
  }

  return [...parts];
}

function sumParts(data, partIndexes) {
  let bytes = 0;
  let compressedBytes = 0;

  for (const partIndex of partIndexes) {
    const part = data.header.chunk_parts[partIndex];
    bytes += part.size;
    compressedBytes += part.compressed_size;
  }

  return { bytes, compressedBytes };
}

function getAsyncBoundariesByModulePath(data) {
  const modules = data.header.modules;
  const entryDependents = [
    "next/dist/esm/build/templates/pages.js",
    "next/dist/esm/build/templates/pages-api.js",
    "next/dist/esm/build/templates/pages-edge-api.js",
    "next/dist/esm/build/templates/edge-ssr.js",
    "next/dist/esm/build/templates/app-route.js",
    "next/dist/esm/build/templates/edge-app-route.js",
    "next/dist/esm/build/templates/app-page.js",
    "next/dist/esm/build/templates/edge-ssr-app.js",
    "next/dist/esm/build/templates/middleware.js",
    "[next]/entry/page-loader.ts",
  ];
  const directEntries = [
    "next/dist/client/app-next-turbopack.js",
    "next/dist/client/next-turbopack.js",
  ];
  const entries = new Set();

  for (let index = 0; index < modules.length; index += 1) {
    const moduleRecord = modules[index];
    if (entryDependents.some((entry) => moduleRecord.ident.includes(entry))) {
      for (const dependency of data.readEdges(
        data.header.module_dependencies,
        index,
      )) {
        if (!modules[dependency].path.includes("next/dist/")) {
          entries.add(dependency);
        }
      }
    }
    if (directEntries.some((entry) => moduleRecord.ident.includes(entry))) {
      entries.add(index);
    }
  }

  const distances = Array(modules.length).fill(Number.POSITIVE_INFINITY);
  const buckets = [[...entries]];
  for (const entry of entries) distances[entry] = 0;

  for (let boundaryCount = 0; boundaryCount < buckets.length; boundaryCount += 1) {
    const bucket = buckets[boundaryCount] ?? [];
    for (let cursor = 0; cursor < bucket.length; cursor += 1) {
      const moduleIndex = bucket[cursor];
      if (distances[moduleIndex] !== boundaryCount) continue;

      relaxDependencies(
        data.header.module_dependencies,
        moduleIndex,
        boundaryCount,
      );
      relaxDependencies(
        data.header.async_module_dependencies,
        moduleIndex,
        boundaryCount + 1,
      );
    }
  }

  const boundariesByPath = new Map();
  for (let index = 0; index < modules.length; index += 1) {
    const moduleRecord = modules[index];
    if (!moduleRecord.ident.includes("[app-client]")) continue;
    const boundaryCount = distances[index];
    if (!Number.isFinite(boundaryCount)) continue;
    const current = boundariesByPath.get(moduleRecord.path);
    if (current === undefined || boundaryCount < current) {
      boundariesByPath.set(moduleRecord.path, boundaryCount);
    }
  }

  return boundariesByPath;

  function relaxDependencies(reference, moduleIndex, nextDistance) {
    for (const dependency of data.readEdges(reference, moduleIndex)) {
      if (nextDistance >= distances[dependency]) continue;
      distances[dependency] = nextDistance;
      buckets[nextDistance] ??= [];
      buckets[nextDistance].push(dependency);
    }
  }
}

function normalizeSourcePath(path) {
  return path.replace(/ \[app-(?:client|ssr)\].*$/, "");
}

function parseList(value) {
  return value
    .split(",")
    .map((item) => item.trim())
    .filter(Boolean)
    .map((route) => (route.startsWith("/") ? route : `/${route}`));
}

function formatBytes(bytes) {
  return `${(bytes / 1024).toFixed(1)} KiB`;
}

class FramedDataFile {
  constructor(filePath) {
    this.buffer = readFileSync(filePath);
    const headerLength = this.buffer.readUInt32BE(0);
    this.binaryStart = 4 + headerLength;
    this.header = JSON.parse(
      this.buffer.subarray(4, this.binaryStart).toString("utf8"),
    );
    this.sourcePathCache = new Map();
  }

  readEdges(reference, index) {
    if (reference.length === 0) return [];
    const referenceStart = this.binaryStart + reference.offset;
    const indexCount = this.buffer.readUInt32BE(referenceStart);
    if (index < 0 || index >= indexCount) return [];

    const offsetsStart = referenceStart + 4;
    const previousOffset =
      index === 0
        ? 0
        : this.buffer.readUInt32BE(offsetsStart + (index - 1) * 4);
    const nextOffset = this.buffer.readUInt32BE(offsetsStart + index * 4);
    const dataStart = offsetsStart + indexCount * 4;
    const edges = [];

    for (let cursor = previousOffset; cursor < nextOffset; cursor += 1) {
      edges.push(this.buffer.readUInt32BE(dataStart + cursor * 4));
    }
    return edges;
  }

  getFullSourcePath(index) {
    const cached = this.sourcePathCache.get(index);
    if (cached !== undefined) return cached;
    const source = this.header.sources[index];
    const path =
      source.parent_source_index === null
        ? source.path
        : `${this.getFullSourcePath(source.parent_source_index)}${source.path}`;
    this.sourcePathCache.set(index, path);
    return path;
  }
}

main();