import { gzipSync } from 'node:zlib'
import { readdirSync, readFileSync } from 'node:fs'
import { join } from 'node:path'

const chunksDirectory = join(process.cwd(), '.next/static/chunks')
const routeStatsPath = join(process.cwd(), '.next/diagnostics/route-bundle-stats.json')
const files = readdirSync(chunksDirectory).filter((file) => file.endsWith('.js'))

// User-facing budgets are based on what a route downloads on first load. The
// sum of every generated chunk is informational only: adding a lazy route or
// splitting a feature creates more files without increasing a user's payload.
const MAX_ROUTE_FIRST_LOAD_GZIP_KIB = 720
const MAX_SHARED_FIRST_LOAD_GZIP_KIB = 330
const MAX_SINGLE_RAW_KIB = 900
const chunkMetrics = new Map()
let totalGzipBytes = 0
let largest = { bytes: 0, file: '' }

const getChunkMetrics = (chunkPath) => {
  const cached = chunkMetrics.get(chunkPath)

  if (cached) return cached

  const source = readFileSync(join(process.cwd(), chunkPath))
  const metrics = { gzipBytes: gzipSync(source).byteLength, rawBytes: source.byteLength }

  chunkMetrics.set(chunkPath, metrics)

  return metrics
}

for (const file of files) {
  const source = readFileSync(join(chunksDirectory, file))

  totalGzipBytes += gzipSync(source).byteLength
  if (source.byteLength > largest.bytes) largest = { bytes: source.byteLength, file }
}

const routeStats = JSON.parse(readFileSync(routeStatsPath, 'utf8'))

if (!Array.isArray(routeStats) || routeStats.length === 0) {
  throw new Error('Next.js route bundle statistics are missing or empty')
}

const routeUsage = new Map()
let largestRoute = { gzipBytes: 0, route: '' }

for (const entry of routeStats) {
  if (typeof entry.route !== 'string' || !Array.isArray(entry.firstLoadChunkPaths)) {
    throw new Error('Next.js route bundle statistics have an unexpected shape')
  }

  const uniqueChunkPaths = new Set(entry.firstLoadChunkPaths)
  let routeGzipBytes = 0

  for (const chunkPath of uniqueChunkPaths) {
    routeGzipBytes += getChunkMetrics(chunkPath).gzipBytes
    routeUsage.set(chunkPath, (routeUsage.get(chunkPath) ?? 0) + 1)
  }

  if (routeGzipBytes > largestRoute.gzipBytes) largestRoute = { gzipBytes: routeGzipBytes, route: entry.route }
}

const sharedGzipBytes = [...routeUsage.entries()]
  .filter(([, routeCount]) => routeCount === routeStats.length)
  .reduce((total, [chunkPath]) => total + getChunkMetrics(chunkPath).gzipBytes, 0)

const totalGzipKiB = Math.ceil(totalGzipBytes / 1024)
const largestRawKiB = Math.ceil(largest.bytes / 1024)
const largestRouteGzipKiB = Math.ceil(largestRoute.gzipBytes / 1024)
const sharedGzipKiB = Math.ceil(sharedGzipBytes / 1024)
const failures = []

if (largestRouteGzipKiB > MAX_ROUTE_FIRST_LOAD_GZIP_KIB) {
  failures.push(`route first-load ${largestRoute.route} ${largestRouteGzipKiB} KiB gzip > ${MAX_ROUTE_FIRST_LOAD_GZIP_KIB} KiB`)
}
if (sharedGzipKiB > MAX_SHARED_FIRST_LOAD_GZIP_KIB) {
  failures.push(`shared first-load ${sharedGzipKiB} KiB gzip > ${MAX_SHARED_FIRST_LOAD_GZIP_KIB} KiB`)
}
if (largestRawKiB > MAX_SINGLE_RAW_KIB) failures.push(`largest chunk ${largest.file} ${largestRawKiB} KiB > ${MAX_SINGLE_RAW_KIB} KiB`)

console.log(
  `Bundle budget: route ${largestRoute.route} ${largestRouteGzipKiB}/${MAX_ROUTE_FIRST_LOAD_GZIP_KIB} KiB gzip; shared ${sharedGzipKiB}/${MAX_SHARED_FIRST_LOAD_GZIP_KIB} KiB gzip; largest ${largestRawKiB}/${MAX_SINGLE_RAW_KIB} KiB raw (${largest.file})`
)
console.log(`All generated JavaScript: ${totalGzipKiB} KiB gzip (informational, includes lazy chunks for every route)`)

if (failures.length) {
  console.error(failures.join('\n'))
  process.exit(1)
}