import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"
import { tmpdir } from "node:os"
import { resolve } from "node:path"
import { spawn } from "node:child_process"

type CacheInfo = {
  storeDir: string | undefined
  wantMassQuery: string | undefined
  priority: string | undefined
}

type CacheProbe =
  | {
      ok: true
      url: string
      elapsedMs: number
      info: CacheInfo
    }
  | {
      ok: false
      url: string
      elapsedMs: number
      error: string
    }

type Narinfo = {
  storePath: string
  references: string[]
  narSize: number | undefined
  fileSize: number | undefined
  system: string | undefined
}

type NarinfoProbe =
  | {
      ok: true
      url: string
      elapsedMs: number
      narinfo: Narinfo
    }
  | {
      ok: false
      url: string
      elapsedMs: number
      status: number | "error"
      error: string
    }

type Options = {
  urls: string[]
  timeoutMs: number
  guixTribesDir: string
  target: string
  system: string
  concurrency: number
  deriveTimeoutMs: number | undefined
  showMissing: boolean
}

type CacheCoverage = {
  url: string
  available: Set<string>
  missing: Set<string>
}

const DEFAULT_SUBSTITUTE_URLS = [
  "https://guix.tribe-one.org",
  "https://bordeaux.guix.gnu.org",
  "https://ci.guix.gnu.org"
]

const PLUGIN_TARGETS = new Map<string, string>([
  ["base", "base-edge-operating-system"],
  ["aether", "aether-edge-operating-system"],
  ["sender", "sender-edge-operating-system"],
  ["supertest", "supertest-edge-operating-system"]
])

function usage(): string {
  return `Usage: npm run preflight:substitutes -- [OPTIONS]

Check substitute coverage for the Guix system targets built by ../guix-tribes.

Options:
  --plugin NAME         Check a plugin-enabled node target: aether, sender, or supertest.
                        The default is the base Tribes node target.
  --target NAME         Check an explicit (tribes ci substitutes) OS target.
  --guix-tribes PATH    guix-tribes checkout. Default: ../guix-tribes.
  --url URL             Add one substitute URL. Can be repeated.
  --urls URLS           Space or comma separated substitute URLs.
                        Defaults to LEGION_GUIX_SUBSTITUTE_URLS or Legion's defaults.
  --timeout SECONDS     Per-request timeout. Defaults to LEGION_GUIX_SUBSTITUTE_HEALTH_TIMEOUT or 3.
  --system SYSTEM       Guix system. Default: x86_64-linux.
  --concurrency COUNT   Concurrent narinfo probes. Default: 24.
  --derive-timeout SEC  Optional timeout for the Guix time-machine derivation step.
  --show-missing        Print missing store paths per substitute server.
  -h, --help            Show this help.

Examples:
  npm run preflight:substitutes
  npm run preflight:substitutes -- --plugin sender
`
}

function parseArgs(args: string[], env: NodeJS.ProcessEnv): Options {
  const urls: string[] = []
  let timeoutSeconds = Number(env["LEGION_GUIX_SUBSTITUTE_HEALTH_TIMEOUT"] || "3")
  let guixTribesDir = env["SUPERTEST_GUIX_TRIBES_REPO"] || "../guix-tribes"
  let target = PLUGIN_TARGETS.get("base") ?? "base-edge-operating-system"
  let system = "x86_64-linux"
  let concurrency = 24
  let deriveTimeoutSeconds =
    env["SUPERTEST_SUBSTITUTE_DERIVE_TIMEOUT"] === undefined
      ? undefined
      : Number(env["SUPERTEST_SUBSTITUTE_DERIVE_TIMEOUT"])
  let showMissing = false

  for (let index = 0; index < args.length; index += 1) {
    const arg = args[index]
    if (!arg) {
      throw new Error("Unexpected empty argument")
    }

    if (arg === "-h" || arg === "--help") {
      process.stdout.write(usage())
      process.exit(0)
    }

    if (arg === "--plugin") {
      target = targetForPlugin(requireValue(args, (index += 1), arg))
      continue
    }

    if (arg === "--target") {
      target = normalizeTargetName(requireValue(args, (index += 1), arg))
      continue
    }

    if (arg === "--guix-tribes") {
      guixTribesDir = requireValue(args, (index += 1), arg)
      continue
    }

    if (arg === "--url") {
      urls.push(requireValue(args, (index += 1), arg))
      continue
    }

    if (arg === "--urls") {
      urls.push(...splitUrls(requireValue(args, (index += 1), arg)))
      continue
    }

    if (arg === "--timeout") {
      timeoutSeconds = Number(requireValue(args, (index += 1), arg))
      continue
    }

    if (arg === "--system") {
      system = requireValue(args, (index += 1), arg)
      continue
    }

    if (arg === "--concurrency") {
      concurrency = Number(requireValue(args, (index += 1), arg))
      continue
    }

    if (arg === "--derive-timeout") {
      deriveTimeoutSeconds = Number(requireValue(args, (index += 1), arg))
      continue
    }

    if (arg === "--show-missing") {
      showMissing = true
      continue
    }

    if (arg.startsWith("--plugin=")) {
      target = targetForPlugin(arg.slice("--plugin=".length))
      continue
    }

    if (arg.startsWith("--target=")) {
      target = normalizeTargetName(arg.slice("--target=".length))
      continue
    }

    if (arg.startsWith("--guix-tribes=")) {
      guixTribesDir = arg.slice("--guix-tribes=".length)
      continue
    }

    if (arg.startsWith("--url=")) {
      urls.push(arg.slice("--url=".length))
      continue
    }

    if (arg.startsWith("--urls=")) {
      urls.push(...splitUrls(arg.slice("--urls=".length)))
      continue
    }

    if (arg.startsWith("--timeout=")) {
      timeoutSeconds = Number(arg.slice("--timeout=".length))
      continue
    }

    if (arg.startsWith("--system=")) {
      system = arg.slice("--system=".length)
      continue
    }

    if (arg.startsWith("--concurrency=")) {
      concurrency = Number(arg.slice("--concurrency=".length))
      continue
    }

    if (arg.startsWith("--derive-timeout=")) {
      deriveTimeoutSeconds = Number(arg.slice("--derive-timeout=".length))
      continue
    }

    throw new Error(`Unexpected argument: ${arg}\n\n${usage()}`)
  }

  const envUrls = splitUrls(env["LEGION_GUIX_SUBSTITUTE_URLS"] || "")
  const selectedUrls = dedupe(
    urls.length > 0 ? urls : envUrls.length > 0 ? envUrls : DEFAULT_SUBSTITUTE_URLS
  )

  if (!Number.isFinite(timeoutSeconds) || timeoutSeconds <= 0) {
    throw new Error(`Invalid --timeout value: ${timeoutSeconds}`)
  }

  if (!Number.isInteger(concurrency) || concurrency <= 0) {
    throw new Error(`Invalid --concurrency value: ${concurrency}`)
  }

  if (
    deriveTimeoutSeconds !== undefined &&
    (!Number.isFinite(deriveTimeoutSeconds) || deriveTimeoutSeconds <= 0)
  ) {
    throw new Error(`Invalid --derive-timeout value: ${deriveTimeoutSeconds}`)
  }

  return {
    urls: selectedUrls,
    timeoutMs: Math.round(timeoutSeconds * 1000),
    guixTribesDir: resolve(process.cwd(), guixTribesDir),
    target,
    system,
    concurrency,
    deriveTimeoutMs:
      deriveTimeoutSeconds === undefined ? undefined : Math.round(deriveTimeoutSeconds * 1000),
    showMissing
  }
}

function requireValue(args: string[], index: number, option: string): string {
  const value = args[index]
  if (!value || value.startsWith("--")) {
    throw new Error(`Missing value for ${option}`)
  }
  return value
}

function targetForPlugin(plugin: string): string {
  const target = PLUGIN_TARGETS.get(plugin)
  if (!target || plugin === "base") {
    throw new Error(`Unknown plugin '${plugin}'. Expected one of: aether, sender, supertest.`)
  }
  return target
}

function normalizeTargetName(target: string): string {
  return target.endsWith("-operating-system") ? target : `${target}-operating-system`
}

function splitUrls(value: string): string[] {
  return value
    .split(/[,\s]+/)
    .map((entry) => entry.trim())
    .filter(Boolean)
}

function dedupe(values: string[]): string[] {
  return [...new Set(values)]
}

function cacheInfoUrl(url: string): string {
  return `${url.replace(/\/+$/, "")}/nix-cache-info`
}

function narinfoUrl(url: string, storePath: string): string {
  return `${url.replace(/\/+$/, "")}/${storePathHash(storePath)}.narinfo`
}

async function fetchText(
  url: string,
  timeoutMs: number
): Promise<{ status: number; text: string }> {
  const controller = new AbortController()
  const timeout = setTimeout(() => controller.abort(), timeoutMs)
  try {
    const response = await fetch(url, { signal: controller.signal })
    return { status: response.status, text: await response.text() }
  } finally {
    clearTimeout(timeout)
  }
}

async function probeCache(url: string, timeoutMs: number): Promise<CacheProbe> {
  const startedAt = Date.now()
  try {
    const response = await fetchText(cacheInfoUrl(url), timeoutMs)
    const elapsedMs = Date.now() - startedAt
    if (response.status < 200 || response.status > 299) {
      return { ok: false, url, elapsedMs, error: `HTTP ${response.status}` }
    }

    const info = parseCacheInfo(response.text)
    if (info.storeDir !== "/gnu/store") {
      return {
        ok: false,
        url,
        elapsedMs,
        error: `unexpected StoreDir: ${info.storeDir || "<missing>"}`
      }
    }

    return { ok: true, url, elapsedMs, info }
  } catch (error) {
    return {
      ok: false,
      url,
      elapsedMs: Date.now() - startedAt,
      error: error instanceof Error ? error.message : String(error)
    }
  }
}

function parseCacheInfo(raw: string): CacheInfo {
  const fields = parseColonFields(raw)
  return {
    storeDir: fields.get("StoreDir"),
    wantMassQuery: fields.get("WantMassQuery"),
    priority: fields.get("Priority")
  }
}

async function probeNarinfo(
  url: string,
  storePath: string,
  timeoutMs: number
): Promise<NarinfoProbe> {
  const startedAt = Date.now()
  try {
    const response = await fetchText(narinfoUrl(url, storePath), timeoutMs)
    const elapsedMs = Date.now() - startedAt
    if (response.status !== 200) {
      return {
        ok: false,
        url,
        elapsedMs,
        status: response.status,
        error: response.status === 404 ? "missing" : `HTTP ${response.status}`
      }
    }

    return { ok: true, url, elapsedMs, narinfo: parseNarinfo(response.text) }
  } catch (error) {
    return {
      ok: false,
      url,
      elapsedMs: Date.now() - startedAt,
      status: "error",
      error: error instanceof Error ? error.message : String(error)
    }
  }
}

function parseNarinfo(raw: string): Narinfo {
  const fields = parseColonFields(raw)
  const storePath = fields.get("StorePath")
  if (!storePath) {
    throw new Error("narinfo missing StorePath")
  }

  return {
    storePath,
    references: (fields.get("References") ?? "")
      .split(/\s+/)
      .map((entry) => entry.trim())
      .filter(Boolean)
      .map((entry) => (entry.startsWith("/gnu/store/") ? entry : `/gnu/store/${entry}`)),
    narSize: parseOptionalNumber(fields.get("NarSize")),
    fileSize: parseOptionalNumber(fields.get("FileSize")),
    system: fields.get("System")
  }
}

function parseColonFields(raw: string): Map<string, string> {
  const fields = new Map<string, string>()
  for (const line of raw.split(/\r?\n/)) {
    const match = /^([^:]+):\s*(.*)$/.exec(line)
    if (match?.[1] !== undefined && match[2] !== undefined) {
      fields.set(match[1], match[2])
    }
  }
  return fields
}

function parseOptionalNumber(value: string | undefined): number | undefined {
  if (!value) {
    return undefined
  }
  const parsed = Number(value)
  return Number.isFinite(parsed) ? parsed : undefined
}

function storePathHash(path: string): string {
  const match = /^\/gnu\/store\/([0-9a-df-np-sv-z]{32})-[^/]+$/.exec(path.trim())
  const hash = match?.[1]
  if (!hash) {
    throw new Error(`Not a valid /gnu/store path: ${path}`)
  }
  return hash
}

async function deriveSystemRootPath(
  options: Options
): Promise<{ drvPath: string; rootPath: string }> {
  const channelContext = await writePinnedChannelsFile(options.guixTribesDir)
  const expression = `(@ (tribes ci substitutes) ${options.target})`
  try {
    const result = await runCommand(
      "guix",
      [
        "time-machine",
        "-C",
        channelContext.channelsPath,
        "--",
        "system",
        "build",
        "--derivation",
        "--verbosity=0",
        `--system=${options.system}`,
        `--expression=${expression}`
      ],
      options.deriveTimeoutMs
    )

    if (result.exitCode !== 0) {
      const timeoutNote = result.timedOut
        ? `\nThe derivation command timed out after ${(options.deriveTimeoutMs ?? 0) / 1000}s.`
        : ""
      throw new Error(
        `Failed to derive ${expression} through guix time-machine.${timeoutNote}\nchannels: ${channelContext.channelsPath}\nstdout:\n${tail(result.stdout)}\nstderr:\n${tail(result.stderr)}`
      )
    }

    const drvPath = result.stdout
      .split(/\r?\n/)
      .map((line) => line.trim())
      .findLast((line) => /^\/gnu\/store\/.+\.drv$/.test(line))
    if (!drvPath) {
      throw new Error(`Could not find derivation path in guix output:\n${result.stdout}`)
    }

    const drv = await readFile(drvPath, "utf8")
    const rootPath = firstDerivationOutputPath(drv)
    return { drvPath, rootPath }
  } finally {
    await rm(channelContext.tempDir, { recursive: true, force: true })
  }
}

async function writePinnedChannelsFile(
  guixTribesDir: string
): Promise<{ tempDir: string; channelsPath: string }> {
  const sourcePath = resolve(guixTribesDir, "pins/legion-channels.sexp")
  const commitResult = await runCommand("git", ["-C", guixTribesDir, "rev-parse", "HEAD"])
  if (commitResult.exitCode !== 0) {
    throw new Error(`Failed to read guix-tribes HEAD:\n${tail(commitResult.stderr)}`)
  }

  const commit = commitResult.stdout.trim()
  if (!/^[0-9a-f]{40}$/.test(commit)) {
    throw new Error(`Unexpected guix-tribes HEAD commit: ${commit}`)
  }

  const source = await readFile(sourcePath, "utf8")
  const rendered = insertTribesChannelCommit(source, commit)
  const tempDir = await mkdtemp(resolve(tmpdir(), "tribes-supertest-channels-"))
  const channelsPath = resolve(tempDir, "channels.scm")
  await writeFile(channelsPath, rendered, "utf8")
  return { tempDir, channelsPath }
}

function insertTribesChannelCommit(channels: string, commit: string): string {
  const marker = /(name 'tribes\)[\s\S]*?\(branch\s+"[^"]+"\))/
  if (!marker.test(channels)) {
    throw new Error("Could not find tribes channel branch in legion-channels.sexp")
  }

  return channels.replace(marker, (match) => {
    if (/\(commit\s+"[0-9a-f]{40,}"\)/.test(match)) {
      return match.replace(/\(commit\s+"[0-9a-f]{40,}"\)/, `(commit "${commit}")`)
    }
    return `${match}\n    (commit "${commit}")`
  })
}

function firstDerivationOutputPath(drv: string): string {
  const match = /^Derive\(\[\("[^"]+","([^"]+)"/.exec(drv)
  const output = match?.[1]
  if (!output?.startsWith("/gnu/store/")) {
    throw new Error("Could not parse output path from derivation")
  }
  return output
}

type CommandResult = {
  exitCode: number
  stdout: string
  stderr: string
  timedOut: boolean
}

async function runCommand(
  command: string,
  args: string[],
  timeoutMs?: number
): Promise<CommandResult> {
  return await new Promise((resolvePromise, reject) => {
    const child = spawn(command, args, {
      cwd: process.cwd(),
      env: process.env,
      stdio: ["ignore", "pipe", "pipe"]
    })
    const stdoutChunks: Buffer[] = []
    const stderrChunks: Buffer[] = []
    let timedOut = false
    const timeout =
      timeoutMs === undefined
        ? undefined
        : setTimeout(() => {
            timedOut = true
            child.kill("SIGTERM")
          }, timeoutMs)

    child.stdout.on("data", (chunk: Buffer) => {
      stdoutChunks.push(chunk)
    })
    child.stderr.on("data", (chunk: Buffer) => {
      stderrChunks.push(chunk)
    })
    child.on("error", reject)
    child.on("close", (exitCode) => {
      if (timeout) {
        clearTimeout(timeout)
      }
      resolvePromise({
        exitCode: exitCode ?? 1,
        stdout: Buffer.concat(stdoutChunks).toString("utf8"),
        stderr: Buffer.concat(stderrChunks).toString("utf8"),
        timedOut
      })
    })
  })
}

function tail(text: string, maxLines = 80): string {
  const lines = text.trimEnd().split(/\r?\n/)
  return lines.slice(-maxLines).join("\n")
}

async function discoverClosure(
  rootPath: string,
  urls: string[],
  timeoutMs: number,
  concurrency: number
): Promise<{ closure: Set<string>; coverage: CacheCoverage[]; unresolved: Set<string> }> {
  const closure = new Set<string>([rootPath])
  const queued = [rootPath]
  const resolved = new Set<string>()
  const unresolved = new Set<string>()
  const coverage = urls.map((url) => ({
    url,
    available: new Set<string>(),
    missing: new Set<string>()
  }))

  async function processOne(path: string): Promise<void> {
    const probes = await Promise.all(urls.map((url) => probeNarinfo(url, path, timeoutMs)))
    let firstNarinfo: Narinfo | undefined

    for (const probe of probes) {
      const entry = coverage.find((item) => item.url === probe.url)
      if (!entry) {
        continue
      }
      if (probe.ok) {
        entry.available.add(path)
        firstNarinfo ??= probe.narinfo
      } else {
        entry.missing.add(path)
      }
    }

    if (!firstNarinfo) {
      unresolved.add(path)
      return
    }

    for (const reference of firstNarinfo.references) {
      if (!closure.has(reference)) {
        closure.add(reference)
        queued.push(reference)
      }
    }
  }

  while (queued.length > 0) {
    const batch = queued.splice(0, concurrency).filter((path) => !resolved.has(path))
    await Promise.all(
      batch.map(async (path) => {
        await processOne(path)
        resolved.add(path)
      })
    )
  }
  return { closure, coverage, unresolved }
}

function printCacheProbe(probe: CacheProbe): void {
  if (probe.ok) {
    const details = [
      `StoreDir=${probe.info.storeDir}`,
      `WantMassQuery=${probe.info.wantMassQuery ?? "<missing>"}`,
      `Priority=${probe.info.priority ?? "<missing>"}`
    ].join(" ")
    console.log(`ok   ${probe.url} (${probe.elapsedMs} ms) ${details}`)
    return
  }

  console.log(`fail ${probe.url} (${probe.elapsedMs} ms) ${probe.error}`)
}

function formatPercent(value: number, total: number): string {
  if (total === 0) {
    return "0%"
  }
  return `${Math.round((value / total) * 100)}%`
}

async function main(): Promise<void> {
  const options = parseArgs(process.argv.slice(2), process.env)

  console.log("Guix substitute preflight")
  console.log(`guix-tribes: ${options.guixTribesDir}`)
  console.log(`target: ${options.target}`)
  console.log(`system: ${options.system}`)
  console.log(`timeout: ${options.timeoutMs / 1000}s`)
  console.log("")
  console.log("Substitute cache health:")

  const cacheProbes = await Promise.all(
    options.urls.map((url) => probeCache(url, options.timeoutMs))
  )
  for (const probe of cacheProbes) {
    printCacheProbe(probe)
  }

  const healthyUrls = cacheProbes.filter((probe) => probe.ok).map((probe) => probe.url)
  if (healthyUrls.length === 0) {
    throw new Error("No healthy substitute caches remain.")
  }

  console.log("")
  console.log("Deriving target output path through guix time-machine...")
  const { drvPath, rootPath } = await deriveSystemRootPath(options)
  console.log(`derivation: ${drvPath}`)
  console.log(`root: ${rootPath}`)

  console.log("")
  console.log("Walking narinfo references...")
  const { closure, coverage, unresolved } = await discoverClosure(
    rootPath,
    healthyUrls,
    options.timeoutMs,
    options.concurrency
  )
  const total = closure.size
  console.log("")
  console.log(`${total} total required NARs`)
  for (const entry of coverage) {
    const count = entry.available.size
    console.log(`${entry.url} has ${count} (${formatPercent(count, total)})`)
  }

  if (unresolved.size > 0) {
    console.log("")
    console.log(
      `warning: ${unresolved.size} path(s) were missing from all healthy caches; their references could not be discovered.`
    )
  }

  if (options.showMissing) {
    console.log("")
    for (const entry of coverage) {
      const missing = [...closure].filter((path) => !entry.available.has(path)).sort()
      console.log(`Missing on ${entry.url}: ${missing.length}`)
      for (const path of missing) {
        console.log(`  ${path}`)
      }
    }
  }
}

main().catch((error: unknown) => {
  console.error(error instanceof Error ? error.message : String(error))
  process.exitCode = 1
})