shrimp/bin/shrimp

#!/usr/bin/env bun

import { Compiler } from '../src/compiler/compiler'
import { VM, toValue, fromValue, bytecodeToString } from 'reefvm'
import { readFileSync, writeFileSync, mkdirSync } from 'fs'
import { randomUUID } from "crypto"
import { spawn } from 'child_process'
import { join } from 'path'

const colors = {
  reset: '\x1b[0m',
  bright: '\x1b[1m',
  dim: '\x1b[2m',
  red: '\x1b[31m',
  yellow: '\x1b[33m',
  cyan: '\x1b[36m',
  magenta: '\x1b[35m',
  pink: '\x1b[38;2;255;105;180m'
}

const nativeFunctions = {
  echo: (...args: any[]) => console.log(...args),
  len: (value: any) => {
    if (typeof value === 'string') return value.length
    if (Array.isArray(value)) return value.length
    if (value && typeof value === 'object') return Object.keys(value).length
    return 0
  },
  type: (value: any) => toValue(value).type,
  range: (start: number, end: number | null) => {
    if (end === null) {
      end = start
      start = 0
    }
    const result: number[] = []
    for (let i = start; i <= end; i++) {
      result.push(i)
    }
    return result
  },
  join: (arr: any[], sep: string = ',') => arr.join(sep),
  split: (str: string, sep: string = ',') => str.split(sep),
  upper: (str: string) => str.toUpperCase(),
  lower: (str: string) => str.toLowerCase(),
  trim: (str: string) => str.trim(),
  list: (...args: any[]) => args,
  dict: (atNamed = {}) => atNamed
}

async function runFile(filePath: string) {
  try {
    const code = readFileSync(filePath, 'utf-8')
    const compiler = new Compiler(code)
    const vm = new VM(compiler.bytecode, nativeFunctions)
    await vm.run()
    return vm.stack.length ? fromValue(vm.stack[vm.stack.length - 1]) : null
  } catch (error: any) {
    console.error(`${colors.red}Error:${colors.reset} ${error.message}`)
    process.exit(1)
  }
}

async function compileFile(filePath: string) {
  try {
    const code = readFileSync(filePath, 'utf-8')
    const compiler = new Compiler(code)
    return bytecodeToString(compiler.bytecode)
  } catch (error: any) {
    console.error(`${colors.red}Error:${colors.reset} ${error.message}`)
    process.exit(1)
  }
}

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> [...args]

${colors.bright}Commands:${colors.reset}
  ${colors.cyan}run      ${colors.yellow}./my-file.sh${colors.reset}    Execute a file with Shrimp
  ${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}repl${colors.reset}                     Start REPL
  ${colors.cyan}help${colors.reset}                     Print this help message
  ${colors.cyan}version${colors.reset}                  Print version`)
}

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

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

  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)
    }

    try { mkdirSync('/tmp/shrimp') } catch { }
    const path = `/tmp/shrimp/${randomUUID()}.sh`
    writeFileSync(path, code)
    console.log(await runFile(path))
    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(await compileFile(file))
    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)
    }
    await runFile(file)
    return
  }

  await runFile(command)
}

await main()