shrimp/bin/shrimp

#!/usr/bin/env bun

import { colors, globals as prelude } from '../src/prelude'
import { treeToString2 } from '../src/utils/tree'
import { runCode, runFile, compileFile, parseCode } from '../src'
import { resolve } from 'path'
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 (non-lezer parser)')
}

async function evalCode(code: string, imports: string[]) {
  const idx = Bun.argv.indexOf('--')
  prelude.$.args = idx >= 0 ? Bun.argv.slice(idx + 1) : []

  const importStatement = imports.length > 0 ? `import ${imports.join(' ')}` : ''
  if (importStatement) code = `${importStatement}; ${code}`
  return await runCode(code)
}

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

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

  // Parse -I flags for imports (supports both "-I math" and "-Imath")
  const imports: string[] = []

  while (args.length > 0) {
    const arg = args[0]

    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)
    } else {
      break
    }
  }

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

  const command = args[0]

  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 = args[1]
    if (!code) {
      console.log(`${colors.bright}usage: shrimp eval <code>${colors.reset}`)
      process.exit(1)
    }

    await evalCode(code, imports)
    return
  }

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

    console.log(await evalCode(code, imports))
    return
  }

  if (['bytecode', '-bytecode', '--bytecode', '-b'].includes(command)) {
    const file = args[1]
    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 = args[1]
    if (!file) {
      console.log(`${colors.bright}usage: shrimp parse <file>${colors.reset}`)
      process.exit(1)
    }
    const input = readFileSync(file, 'utf-8')
    console.log(treeToString2(parseCode(input).topNode, input))
    return
  }

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

  prelude.$.script.path = resolve(command)
  await runFile(command)
}

await main()