cli-framework/src/lib/util/ansi.ts

import { nanoseconds } from 'bun'

const encoder = new TextEncoder()
const emptyBytes = encoder.encode('')
type Writer = (x: Uint8Array | Array<Uint8Array>) => Promise<number>
const writer: Writer = x => Bun.write(Bun.stdout, x)

export const ANSI = {
  // Text Styles
  RESET: '\x1B[0m',
  BOLD: '\x1B[1m',
  DIM: '\x1B[2m',
  ITALIC: '\x1B[3m',
  UNDERLINE: '\x1B[4m',
  BLINK: '\x1B[5m',
  REVERSE: '\x1B[7m',
  HIDDEN: '\x1B[8m',

  // Foreground Colors
  BLACK: '\x1B[30m',
  RED: '\x1B[31m',
  GREEN: '\x1B[32m',
  YELLOW: '\x1B[33m',
  BLUE: '\x1B[34m',
  MAGENTA: '\x1B[35m',
  CYAN: '\x1B[36m',
  WHITE: '\x1B[37m',

  // Background Colors
  BG_BLACK: '\x1B[40m',
  BG_RED: '\x1B[41m',
  BG_GREEN: '\x1B[42m',
  BG_YELLOW: '\x1B[43m',
  BG_BLUE: '\x1B[44m',
  BG_MAGENTA: '\x1B[45m',
  BG_CYAN: '\x1B[46m',
  BG_WHITE: '\x1B[47m',

  // Cursor Control
  GET_CURSOR_POSITION: '\x1B[6n',
  SAVE_CURSOR: '\x1B[s',
  RESTORE_CURSOR: '\x1B[u',
  CURSOR_TO: (x: number, y: number) => `\x1B[${y};${x}H`,
  CURSOR_UP: (n = 1) => `\x1B[${n}A`,
  CURSOR_DOWN: (n = 1) => `\x1B[${n}B`,
  CURSOR_FORWARD: (n = 1) => `\x1B[${n}C`,
  CURSOR_BACK: (n = 1) => `\x1B[${n}D`,

  // Screen Control
  CLEAR_SCREEN: '\x1B[2J',
  CLEAR_LINE: '\x1B[2K',
  CLEAR_END_LINE: '\x1B[0K',
  CLEAR_START_LINE: '\x1B[1K',
  CLEAR_TO_END: '\x1B[0J', // From cursor to end of screen
  CLEAR_TO_START: '\x1B[1J', // From cursor to beginning of screen
  CLEAR_BELOW: '\x1B[J',
}

export const ANSI_BUFFERS = {
  // Static code buffers
  RESET: encoder.encode(ANSI.RESET),
  BOLD: encoder.encode(ANSI.BOLD),
  DIM: encoder.encode(ANSI.DIM),
  ITALIC: encoder.encode(ANSI.ITALIC),
  UNDERLINE: encoder.encode(ANSI.UNDERLINE),
  BLINK: encoder.encode(ANSI.BLINK),
  REVERSE: encoder.encode(ANSI.REVERSE),
  HIDDEN: encoder.encode(ANSI.HIDDEN),

  // Colors
  BLACK: encoder.encode(ANSI.BLACK),
  RED: encoder.encode(ANSI.RED),
  GREEN: encoder.encode(ANSI.GREEN),
  YELLOW: encoder.encode(ANSI.YELLOW),
  BLUE: encoder.encode(ANSI.BLUE),
  MAGENTA: encoder.encode(ANSI.MAGENTA),
  CYAN: encoder.encode(ANSI.CYAN),
  WHITE: encoder.encode(ANSI.WHITE),

  // Backgrounds
  BG_BLACK: encoder.encode(ANSI.BG_BLACK),
  BG_RED: encoder.encode(ANSI.BG_RED),
  BG_GREEN: encoder.encode(ANSI.BG_GREEN),
  BG_YELLOW: encoder.encode(ANSI.BG_YELLOW),
  BG_BLUE: encoder.encode(ANSI.BG_BLUE),
  BG_MAGENTA: encoder.encode(ANSI.BG_MAGENTA),
  BG_CYAN: encoder.encode(ANSI.BG_CYAN),
  BG_WHITE: encoder.encode(ANSI.BG_WHITE),

  // Cursor Control
  GET_CURSOR_POSITION: encoder.encode(ANSI.GET_CURSOR_POSITION),
  SAVE_CURSOR: encoder.encode(ANSI.SAVE_CURSOR),
  RESTORE_CURSOR: encoder.encode(ANSI.RESTORE_CURSOR),

  // Screen Control
  CLEAR_SCREEN: encoder.encode(ANSI.CLEAR_SCREEN),
  CLEAR_LINE: encoder.encode(ANSI.CLEAR_LINE),
  CLEAR_TO_END: encoder.encode(ANSI.CLEAR_TO_END),
  CLEAR_TO_START: encoder.encode(ANSI.CLEAR_TO_START),
  CLEAR_BELOW: encoder.encode(ANSI.CLEAR_BELOW),
}

export const ANSI_DYNAMIC = {
  // Cursor Control
  CURSOR_TO: (x: number, y: number) => encoder.encode(ANSI.CURSOR_TO(x, y)),
  CURSOR_UP: (x = 1) => encoder.encode(ANSI.CURSOR_UP(x)),
  CURSOR_DOWN: (x = 1) => encoder.encode(ANSI.CURSOR_DOWN(x)),
  CURSOR_FORWARD: (x = 1) => encoder.encode(ANSI.CURSOR_FORWARD(x)),
  CURSOR_BACK: (x = 1) => encoder.encode(ANSI.CURSOR_BACK(x)),
}

type AnsiBufferKey = keyof typeof ANSI_BUFFERS

type WriteAnsi = (
  c: Array<(AnsiBufferKey | Uint8Array)>,
  t?: string,
  a?: boolean,
) => Promise<number>
const writeAnsi: WriteAnsi = async (codes, text?, autoReset = true) => {
  const arr = [
    ...codes,
    !text ? emptyBytes : encoder.encode(text),
    autoReset ? ANSI_BUFFERS['RESET'] : emptyBytes,
  ].map(x => x instanceof Uint8Array ? x : ANSI_BUFFERS[x])

  return writer(Buffer.concat(arr))
}

const writeAnsi1: WriteAnsi = async (codes, text?, autoReset = true) => {
  const buffers = [
    ...codes,
    !text ? emptyBytes : encoder.encode(text),
    autoReset ? ANSI_BUFFERS['RESET'] : emptyBytes,
  ].map(x => x instanceof Uint8Array ? x : ANSI_BUFFERS[x])

  // Calculate total length and create combined buffer
  const totalLength = buffers.reduce((sum, buf) => sum + buf.length, 0)
  const combined = new Uint8Array(totalLength)

  // Copy all buffers into the combined array
  let offset = 0
  for (const buffer of buffers) {
    combined.set(buffer, offset)
    offset += buffer.length
  }

  return writer(combined)
}

// For CommonJS compatibility
export default {ANSI, ANSI_BUFFERS, writeAnsi}

const bench = async (x, n) => {
  const a = Bun.nanoseconds()
  let i = 0
  while (i < n) {
    await x()
    i++
  }
  const b = Bun.nanoseconds()
  return (b - a) / 100000
}

const one = async () => {
  await writeAnsi(
    ['BG_CYAN', 'BLACK', 'ITALIC'],
    'This is the first option I want',
  )
  await writeAnsi(
    ['BG_BLUE', 'BLACK', 'BOLD'],
    'This is the second option I want',
  )
  await writeAnsi(
    ['BG_YELLOW', 'BLACK', 'BLINK'],
    'Now, this is the third and final option',
  )
}

const two = async () => {
  await writeAnsi1(
    ['BG_CYAN', 'BLACK', 'ITALIC'],
    'This is the first option I want',
  )
  await writeAnsi1(
    ['BG_BLUE', 'BLACK', 'BOLD'],
    'This is the second option I want',
  )
  await writeAnsi1(
    ['BG_YELLOW', 'BLACK', 'BLINK'],
    'Now, this is the third and final option',
  )
}

const test = async () => {
  const testN = 100000

  await bench(one, 100000)
  await bench(two, 100000)

  const testTwo = await bench(two, testN)
  const testOne = await bench(one, testN)

  await writeAnsi(
    ['RESET', ANSI_DYNAMIC.CURSOR_TO(0, 0)],
    `${ANSI.CLEAR_SCREEN}\n\n`,
  )
  await writeAnsi(['BG_YELLOW', 'BLACK'], '---------------\n')
  await writeAnsi(['BG_YELLOW', 'BOLD', 'BLACK'], '***************\n')
  await writeAnsi(['BG_YELLOW', 'BLACK'], '---------------\n')
  await writeAnsi(['BG_YELLOW', 'BLACK'], 'ONE:')
  await writeAnsi(['BG_YELLOW', 'GREEN'], testOne.toString() + '\n')
  await writeAnsi(['BG_YELLOW', 'BLACK'], 'TWO:')
  await writeAnsi(['BG_YELLOW', 'GREEN'], testTwo.toString() + '\n')
  await writeAnsi(['BG_YELLOW', 'BLACK'], '---------------\n')
  await writeAnsi(['BG_YELLOW', 'BOLD', 'BLACK'], '***************\n')
  await writeAnsi(['BG_YELLOW', 'BLACK'], '---------------')
  await writeAnsi(['RESET'], '\n\n')
}

test()