shrimp/bin/shrimp

#!/usr/bin/env bun

import { colors } from '../src/prelude'
import { treeToString } from '../src/utils/tree'
import { runCode, runFile, compileFile, parseCode } from '../src'
import { bytecodeToString } from 'reefvm'
import { readFileSync } from 'fs'
import { spawn } from 'child_process'
import { join } from 'path'

function showHelp() {
  console.log(`${colors.bright}${colors.magenta}🦐 Shrimp${colors.reset} is a scripting language in a shell.

${colors.bright}Usage:${colors.reset} shrimp <command> [options] [...args]

${colors.bright}Commands:${colors.reset}
  ${colors.cyan}run      ${colors.yellow}./my-file.sh${colors.reset}    Execute a file with Shrimp
  ${colors.cyan}parse    ${colors.yellow}./my-file.sh${colors.reset}    Print parse tree for Shrimp file
  ${colors.cyan}bytecode ${colors.yellow}./my-file.sh${colors.reset}    Print bytecode for Shrimp file
  ${colors.cyan}eval     ${colors.yellow}'some code'${colors.reset}     Evaluate a line of Shrimp code
  ${colors.cyan}print    ${colors.yellow}'some code'${colors.reset}     Evaluate a line of Shrimp code and print the result
  ${colors.cyan}repl${colors.reset}                     Start REPL
  ${colors.cyan}help${colors.reset}                     Print this help message
  ${colors.cyan}version${colors.reset}                  Print version

${colors.bright}Options:${colors.reset}
  ${colors.cyan}eval -I${colors.reset} ${colors.yellow}<module>${colors.reset}         Import module (can be repeated)
                           Example: shrimp -I math -e 'random | echo'
                           Example: shrimp -Imath -Istr -e 'random | echo'`)
}

function showVersion() {
  console.log('🦐 v0.0.1')
}

async function evalCode(code: string, imports: string[], globals?: Record<string, any>) {
  const importStatement = imports.length > 0 ? `import ${imports.join(' ')}` : ''
  if (importStatement) code = `${importStatement}; ${code}`
  return await runCode(code, globals)
}

async function main() {
  let args = process.argv.slice(2)

  if (args.length === 0) {
    showHelp()
    return
  }

  // cli command and arg
  let command = ''
  let cmdArg = ''

  // parse -I flags for imports. supports both "-I math" and "-Imath"
  const imports: string[] = []

  // shrimp -E args -- arg1 arg2
  let grabArgs = false
  let scriptArgs = []

  for (const arg of args) {
    // everything after -- goes into "args" for this script
    if (arg === '--') {
      grabArgs = true
      continue
    }

    // grab all args after --
    if (grabArgs) {
      scriptArgs.push(arg)
      continue
    }

    if (arg === '-I') {
      // "-I math" format
      if (args.length < 2) {
        console.log(`${colors.bright}error: -I requires a module name${colors.reset}`)
        process.exit(1)
      }
      imports.push(args[1])
      args = args.slice(2)
    } else if (arg.startsWith('-I')) {
      // "-Imath" format
      const moduleName = arg.slice(2)
      if (!moduleName) {
        console.log(`${colors.bright}error: -I requires a module name${colors.reset}`)
        process.exit(1)
      }
      imports.push(moduleName)
      args = args.slice(1)
    }

    if (!command) {
      command = arg
    } else {
      cmdArg = arg
    }
  }

  if (args.length === 0) {
    showHelp()
    return
  }

  if (['help', '-help', '--help', '-h'].includes(command)) {
    showHelp()
    return
  }

  if (['version', '-version', '--version', '-v'].includes(command)) {
    showVersion()
    return
  }

  if (['repl', '-repl', '--repl'].includes(command)) {
    const replPath = join(import.meta.dir, 'repl')
    const replArgs = args.slice(1)
    const repl = spawn('bun', [replPath, ...replArgs], { stdio: 'inherit' })
    repl.on('exit', code => process.exit(code || 0))
    return
  }

  if (['eval', '-eval', '--eval', '-e'].includes(command)) {
    const code = cmdArg
    if (!code) {
      console.log(`${colors.bright}usage: shrimp eval <code>${colors.reset}`)
      process.exit(1)
    }

    await evalCode(code, imports, { args: scriptArgs })
    return
  }

  if (['print', '-print', '--print', '-E'].includes(command)) {
    const code = cmdArg
    if (!code) {
      console.log(`${colors.bright}usage: shrimp print <code>${colors.reset}`)
      process.exit(1)
    }

    console.log(await evalCode(code, imports, { args: scriptArgs }))
    return
  }

  if (['bytecode', '-bytecode', '--bytecode', '-b'].includes(command)) {
    const file = cmdArg
    if (!file) {
      console.log(`${colors.bright}usage: shrimp bytecode <file>${colors.reset}`)
      process.exit(1)
    }
    console.log(bytecodeToString(compileFile(file)))
    return
  }

  if (['parse', '-parse', '--parse', '-p'].includes(command)) {
    const file = cmdArg
    if (!file) {
      console.log(`${colors.bright}usage: shrimp parse <file>${colors.reset}`)
      process.exit(1)
    }
    const input = readFileSync(file, 'utf-8')
    console.log(treeToString(parseCode(input), input))
    return
  }

  if (['run', '-run', '--run', '-r'].includes(command)) {
    const file = cmdArg
    if (!file) {
      console.log(`${colors.bright}usage: shrimp run <file>${colors.reset}`)
      process.exit(1)
    }
    await runFile(file)
    return
  }

  await runFile(command)
}

await main()