#!/usr/bin/env node
import fs from "node:fs/promises";
import path from "node:path";
import { fileURLToPath } from "node:url";

function isPathInside(parentPath, candidatePath) {
  const relativePath = path.relative(parentPath, candidatePath);
  return relativePath !== ".." && !relativePath.startsWith(`..${path.sep}`) && !path.isAbsolute(relativePath);
}

function packageNameFromLink(nodeModulesPath, linkPath) {
  const relativePath = path.relative(nodeModulesPath, linkPath);
  const segments = relativePath.split(path.sep);
  return segments[0]?.startsWith("@") ? segments.slice(0, 2).join("/") : segments[0];
}

async function packageLinksAt(nodeModulesPath) {
  let entries;
  try {
    entries = await fs.readdir(nodeModulesPath, { withFileTypes: true });
  } catch (error) {
    if (error?.code === "ENOENT") return [];
    throw error;
  }

  const links = [];
  for (const entry of entries) {
    if (entry.name === ".bin" || entry.name === ".pnpm") continue;
    const entryPath = path.join(nodeModulesPath, entry.name);

    if (entry.isSymbolicLink()) {
      links.push({ linkPath: entryPath, packageName: entry.name });
      continue;
    }

    if (!entry.isDirectory() || !entry.name.startsWith("@")) continue;
    for (const scopedEntry of await fs.readdir(entryPath, { withFileTypes: true })) {
      if (!scopedEntry.isSymbolicLink()) continue;
      const linkPath = path.join(entryPath, scopedEntry.name);
      links.push({ linkPath, packageName: packageNameFromLink(nodeModulesPath, linkPath) });
    }
  }

  return links;
}

function virtualStoreEntryForPath(virtualStorePath, candidatePath) {
  if (!isPathInside(virtualStorePath, candidatePath)) return null;
  const [entryName] = path.relative(virtualStorePath, candidatePath).split(path.sep);
  return entryName && entryName !== "node_modules" ? entryName : null;
}

async function directoryStats(directoryPath) {
  let bytes = 0;
  let files = 0;
  let entries;
  try {
    entries = await fs.readdir(directoryPath, { withFileTypes: true });
  } catch (error) {
    if (error?.code === "ENOENT") return { bytes, files };
    throw error;
  }

  for (const entry of entries) {
    const entryPath = path.join(directoryPath, entry.name);
    if (entry.isDirectory()) {
      const nested = await directoryStats(entryPath);
      bytes += nested.bytes;
      files += nested.files;
    } else if (entry.isFile()) {
      bytes += (await fs.stat(entryPath)).size;
      files += 1;
    }
  }

  return { bytes, files };
}

async function removeBrokenLinks(directoryPath) {
  let removed = 0;
  let entries;
  try {
    entries = await fs.readdir(directoryPath, { withFileTypes: true });
  } catch (error) {
    if (error?.code === "ENOENT") return removed;
    throw error;
  }

  for (const entry of entries) {
    const entryPath = path.join(directoryPath, entry.name);
    if (entry.isSymbolicLink()) {
      try {
        await fs.realpath(entryPath);
      } catch (error) {
        if (error?.code !== "ENOENT") throw error;
        await fs.unlink(entryPath);
        removed += 1;
      }
    } else if (entry.isDirectory()) {
      removed += await removeBrokenLinks(entryPath);
    }
  }

  return removed;
}

export async function prunePnpmRuntimeClosure({ excludedPackageNames = [], nodeModulesPath }) {
  const resolvedNodeModulesPath = await fs.realpath(path.resolve(nodeModulesPath));
  const virtualStorePath = path.join(resolvedNodeModulesPath, ".pnpm");
  const virtualStoreEntries = await fs.readdir(virtualStorePath, { withFileTypes: true });
  const packageEntryNames = virtualStoreEntries
    .filter((entry) => entry.isDirectory() && entry.name !== "node_modules")
    .map((entry) => entry.name)
    .sort();
  const excluded = new Set(excludedPackageNames);
  const reachableEntries = new Set();
  const queue = [];

  async function addLink(link) {
    if (excluded.has(link.packageName)) return;

    let targetPath;
    try {
      targetPath = await fs.realpath(link.linkPath);
    } catch (error) {
      if (error?.code === "ENOENT") return;
      throw error;
    }

    const entryName = virtualStoreEntryForPath(virtualStorePath, targetPath);
    if (!entryName || reachableEntries.has(entryName)) return;
    reachableEntries.add(entryName);
    queue.push(entryName);
  }

  for (const link of await packageLinksAt(resolvedNodeModulesPath)) await addLink(link);
  while (queue.length > 0) {
    const entryName = queue.shift();
    const entryNodeModulesPath = path.join(virtualStorePath, entryName, "node_modules");
    for (const link of await packageLinksAt(entryNodeModulesPath)) await addLink(link);
  }

  if (packageEntryNames.length > 0 && reachableEntries.size === 0) {
    throw new Error(`refusing to prune ${packageEntryNames.length} pnpm entries without a reachable runtime root`);
  }

  const removedEntries = packageEntryNames.filter((entryName) => !reachableEntries.has(entryName));
  let bytesRemoved = 0;
  let filesRemoved = 0;
  for (const entryName of removedEntries) {
    const entryPath = path.join(virtualStorePath, entryName);
    const stats = await directoryStats(entryPath);
    bytesRemoved += stats.bytes;
    filesRemoved += stats.files;
    await fs.rm(entryPath, { force: true, recursive: true });
  }

  const brokenLinksRemoved = await removeBrokenLinks(resolvedNodeModulesPath);
  return {
    brokenLinksRemoved,
    bytesRemoved,
    filesRemoved,
    reachableEntries: reachableEntries.size,
    removedEntries,
  };
}

async function runCli() {
  const nodeModulesPath = process.argv[2] ?? "node_modules";
  const excludedPackageNames = process.argv.slice(3);
  const result = await prunePnpmRuntimeClosure({ excludedPackageNames, nodeModulesPath });
  console.log(
    `PRUNED pnpm-runtime reachable=${result.reachableEntries} entries=${result.removedEntries.length} files=${result.filesRemoved} bytes=${result.bytesRemoved} broken-links=${result.brokenLinksRemoved}`
  );
}

if (process.argv[1] && path.resolve(process.argv[1]) === fileURLToPath(import.meta.url)) {
  await runCli();
}