shout/src/run.ts

import { mkdtemp, rm } from "node:fs/promises"
import { tmpdir } from "node:os"
import { join } from "node:path"

import type { Command, ShoutFile } from "./parse.ts"

export type CommandResult = {
  command: Command
  actual: string[]
  exitCode: number
}

export type FileResult = {
  file: ShoutFile
  results: CommandResult[]
  tmpDir: string
  error?: string
}

type RunOptions = {
  cleanEnv: boolean
  pathDirs?: string[]
  envVars?: Record<string, string>
  sourceDir?: string
  projectDir?: string
  timeout: number
  verbose: boolean
  onCommand?: (cmd: Command) => void
}

function killTree(pid: number): void {
  // Find any processes that escaped the process group (e.g. via setsid).
  // This assumes pid === pgid, which holds because the child is spawned
  // with detached: true (making it a process group leader).
  try {
    const result = Bun.spawnSync(["ps", "-eo", "pid,pgid"])
    const output = result.stdout.toString()
    const pgid = String(pid)
    for (const line of output.split("\n")) {
      const parts = line.trim().split(/\s+/)
      if (parts[1] === pgid) {
        const p = parseInt(parts[0]!, 10)
        if (!isNaN(p) && p !== pid && p > 1) {
          try { process.kill(p, "SIGKILL") } catch {}
        }
      }
    }
  } catch {}
  // Kill the process group
  try { process.kill(-pid, "SIGKILL") } catch {}
}

const SENTINEL_PREFIX = "__SHOUT_SENTINEL_"
const VERBOSE_MARKER = "__SHOUT_CMD_"

function buildScript(commands: Command[], sentinel: string, verbose: boolean): string {
  const lines: string[] = []

  if (verbose) {
    // Save original stderr to fd 3 before merging stderr into stdout
    lines.push("exec 3>&2 2>&1")
  } else {
    lines.push("exec 2>&1")
  }

  for (let i = 0; i < commands.length; i++) {
    const cmd = commands[i]!
    if (verbose) {
      lines.push(`printf '${VERBOSE_MARKER}${i}\\n' >&3`)
    }
    lines.push(cmd.command)
    // Sentinel: printf to avoid echo interpretation issues
    // Format: __SHOUT_SENTINEL_<exitcode>_<index>__
    lines.push(
      `printf '\\n${sentinel}%s_${i}__\\n' "$?"`,
    )
  }

  return lines.join("\n") + "\n"
}

function parseSentinelOutput(
  raw: string,
  sentinel: string,
  commandCount: number,
): { outputs: string[][]; exitCodes: number[] } {
  const outputs: string[][] = []
  const exitCodes: number[] = []

  // Split by sentinel lines
  const sentinelRegex = new RegExp(
    `${escapeRegex(sentinel)}(\\d+)_(\\d+)__`,
  )

  let remaining = raw
  for (let i = 0; i < commandCount; i++) {
    const match = remaining.match(sentinelRegex)
    if (!match) {
      // No sentinel found — rest is output for this command
      const lines = remaining.split("\n")
      // Remove leading empty line (from printf \n prefix)
      if (lines.length > 0 && lines[0] === "") lines.shift()
      outputs.push(trimTrailingEmpty(lines))
      exitCodes.push(1) // assume failure
      break
    }

    const idx = remaining.indexOf(match[0])
    const before = remaining.slice(0, idx)
    const afterSentinel = remaining.slice(idx + match[0].length)

    // Parse output lines
    let lines = before.split("\n")
    // Remove leading empty line from previous sentinel's trailing \n
    if (lines.length > 0 && lines[0] === "") lines.shift()
    // Remove trailing empty lines (from printf's \n prefix)
    lines = trimTrailingEmpty(lines)
    outputs.push(lines.length === 1 && lines[0] === "" ? [] : lines)

    exitCodes.push(parseInt(match[1]!, 10))

    // Skip past sentinel line (including trailing newline)
    remaining = afterSentinel.startsWith("\n")
      ? afterSentinel.slice(1)
      : afterSentinel
  }

  // Fill missing entries
  while (outputs.length < commandCount) {
    outputs.push([])
    exitCodes.push(1)
  }

  return { outputs, exitCodes }
}

function trimTrailingEmpty(lines: string[]): string[] {
  let end = lines.length
  while (end > 0 && lines[end - 1] === "") end--
  return lines.slice(0, end)
}

function escapeRegex(s: string): string {
  return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")
}

function streamVerboseMarkers(
  stderr: ReadableStream<Uint8Array>,
  commands: Command[],
  onCommand: (cmd: Command) => void,
): void {
  const reader = stderr.getReader()
  const decoder = new TextDecoder()
  let buffer = ""

  const pump = (): void => {
    reader.read().then(({ done, value }) => {
      if (done) return
      if (value) buffer += decoder.decode(value, { stream: true })
      let nlIdx: number
      while ((nlIdx = buffer.indexOf("\n")) !== -1) {
        const line = buffer.slice(0, nlIdx)
        buffer = buffer.slice(nlIdx + 1)
        if (line.startsWith(VERBOSE_MARKER)) {
          const i = parseInt(line.slice(VERBOSE_MARKER.length), 10)
          if (i >= 0 && i < commands.length) {
            onCommand(commands[i]!)
          }
        }
      }
      pump()
    }).catch(() => {})
  }
  pump()
}

export async function runFile(
  file: ShoutFile,
  options: RunOptions,
): Promise<FileResult> {
  const tmpDir = await mkdtemp(join(tmpdir(), "shout-"))

  if (file.commands.length === 0) {
    return { file, results: [], tmpDir }
  }

  const sentinel = SENTINEL_PREFIX
  const verbose = options.verbose && !!options.onCommand
  const script = buildScript(file.commands, sentinel, verbose)

  const env: Record<string, string> = options.cleanEnv
    ? {}
    : { ...process.env as Record<string, string> }

  env["HOME"] = tmpDir
  env["SHOUT_DIR"] = tmpDir
  if (options.sourceDir) {
    env["SHOUT_SOURCE_DIR"] = options.sourceDir
  }
  if (options.projectDir) {
    env["SHOUT_PROJECT_DIR"] = options.projectDir
  }

  if (options.envVars) {
    Object.assign(env, options.envVars)
  }

  if (options.pathDirs?.length) {
    env["PATH"] = options.pathDirs.join(":") + ":" + (env["PATH"] ?? "")
  }

  const proc = Bun.spawn(["/bin/sh"], {
    detached: true,
    stdin: "pipe",
    stdout: "pipe",
    stderr: "pipe",
    cwd: tmpDir,
    env,
  })

  try {
    if (verbose) {
      // Stream stderr for verbose command markers before writing script
      streamVerboseMarkers(proc.stderr, file.commands, options.onCommand!)
    }

    proc.stdin.write(script)
    proc.stdin.end()

    const totalTimeout = options.timeout * file.commands.length
    const lastSentinelSuffix = `_${file.commands.length - 1}__`
    const stdout = await readUntilSentinel(proc.stdout, sentinel, lastSentinelSuffix, totalTimeout)
    if (!verbose) {
      await readWithTimeout(proc.stderr, 1000).catch(() => "")
    }

    const { outputs, exitCodes } = parseSentinelOutput(
      stdout,
      sentinel,
      file.commands.length,
    )

    const results: CommandResult[] = file.commands.map((cmd, i) => ({
      command: cmd,
      actual: outputs[i] ?? [],
      exitCode: exitCodes[i] ?? 1,
    }))

    return { file, results, tmpDir }
  } catch (err) {
    return {
      file,
      results: [],
      tmpDir,
      error: err instanceof Error ? err.message : String(err),
    }
  } finally {
    if (proc.pid) {
      killTree(proc.pid)
    }
  }
}

async function readUntilSentinel(
  stream: ReadableStream<Uint8Array>,
  sentinelPrefix: string,
  sentinelSuffix: string,
  timeoutMs: number,
): Promise<string> {
  const reader = stream.getReader()
  const decoder = new TextDecoder()
  let accumulated = ""

  let timerId: ReturnType<typeof setTimeout>
  const timeout = new Promise<never>((_, reject) =>
    timerId = setTimeout(() => reject(new Error("Timeout reading output")), timeoutMs),
  )

  try {
    while (true) {
      const { done, value } = await Promise.race([reader.read(), timeout]) as ReadableStreamReadResult<Uint8Array>
      if (done) break
      if (value) {
        accumulated += decoder.decode(value, { stream: true })
        // Check if the last sentinel has appeared (prefix + exitcode + suffix)
        const prefixIdx = accumulated.lastIndexOf(sentinelPrefix)
        if (prefixIdx !== -1 && accumulated.indexOf(sentinelSuffix, prefixIdx) !== -1) break
      }
    }
  } finally {
    clearTimeout(timerId!)
    reader.releaseLock()
  }

  return accumulated + decoder.decode()
}

async function readWithTimeout(
  stream: ReadableStream<Uint8Array>,
  timeoutMs: number,
): Promise<string> {
  const reader = stream.getReader()
  const chunks: Uint8Array[] = []

  let timerId: ReturnType<typeof setTimeout>
  const timeout = new Promise<never>((_, reject) =>
    timerId = setTimeout(() => reject(new Error("Timeout reading output")), timeoutMs),
  )

  try {
    while (true) {
      const { done, value } = await Promise.race([reader.read(), timeout]) as ReadableStreamReadResult<Uint8Array>
      if (done) break
      if (value) chunks.push(value)
    }
  } finally {
    clearTimeout(timerId!)
    reader.releaseLock()
  }

  const decoder = new TextDecoder()
  return chunks.map(c => decoder.decode(c, { stream: true })).join("") +
    decoder.decode()
}

export async function cleanupTmpDir(dir: string): Promise<void> {
  await rm(dir, { recursive: true, force: true })
}