# fsh - Forka shell
# 2021 - Seculum Forka
# This program is licensed under the agpl v3.0 only license

# imports
import streams
import osproc
import strtabs
import os
import strutils

# globals
var shellVariables = newStringTable(modeCaseSensitive)
# Environment variables get defined in a separate var
var envVariables = newStringTable(modeCaseSensitive)
# populate the table
for k, v in envPairs():
  envVariables[k]=v

# list of builtin commands
let builtins: array[6, string] = [
"set",
"setenv",
"echo",
"exit",
"if",
"cd",
]

# Exception types
type
  ShellError= object of CatchableError
    line: int
    col: int
    whileExec: string
  ShellVariableError= object of ShellError
    variable: string
  ShellCommandError = object of ShellError
    command: string
  InterruptCtrlC = object of CatchableError

# procs for making procedures
proc newShellException(msg: string, line=0, col=0, whileExec=""): ref ShellError =
  var e = newException(ShellError, msg)
  e.line=line
  e.col=col
  e.whileExec = whileExec
  return e

proc newShellVariableException(msg: string, vari: string, line=0, col=0, whileExec=""): ref ShellVariableError =
  var e= newException(ShellVariableError, vari & ": " & msg)
  e.variable =vari
  e.line=line
  e.col=col
  e.whileExec=whileExec
  return e

proc newShellCommandException(msg: string, command: string, line=0, col=0, whileExec=""): ref ShellCommandError =
  var e=newException(ShellCommandError, command & ": " & msg)
  e.command=command
  e.line=line
  e.col=col
  e.whileExec=whileExec
  return e

# forward declarations
proc eval(cmd: string): string
proc execute(cmd: string): int
proc substitute(strm: Stream, delim = ";\p"): string

# ctrlc is the ctrl-c handler
proc ctrlc() {.noconv.} =
  raise newException(InterruptCtrlC, "")

# skipSpaces reads the stream until a non-space character is found
proc skipSpaces(strm: Stream): int =
  while not strm.atEnd():
    if strm.peekChar() == ' ':
      discard strm.readChar()
      result.inc
    else: break

# readLiterals parses the string passed in as a literal, preserving all braces except the first and last ones
proc readLiteral(strm: Stream): string =
  var numbraces=1
  while not strm.atEnd() and numbraces > 0:
    let c = strm.readChar()
    case c:
      of '{':
        numbraces.inc
      of '}':
        numbraces = numbraces - 1
      of '\\':
        if strm.peekChar notin {'{', '}'}: result.add(strm.readChar)
        result.add(strm.readChar)
      else:
        result.add(c)

# a forward declaration. Fuck this shit
proc readSubstitution(strm: Stream): string

# readInterpelation reads the interpelation in between [ and  ]
proc readInterpelation(strm: Stream): string =
  while not strm.atEnd():
    let c = strm.readChar()
    case c:
      of '{':
        result.add('{')
        result.add(strm.readLiteral)
        result.add('}')
      of '[':
        result.add('[')
        result.add(strm.readInterpelation)
        result.add(']')
      of '"':
        result.add("\"")
        result.add(strm.readSubstitution)
        result.add("\"")
      of '\\':
        result.add("\\")
        result.add(strm.readChar())
      of ']':
        break
      else: result.add(c)

# readSubstitution reads a substitutionbetween "quotes"
proc readSubstitution(strm: Stream): string =
  while not strm.atEnd():
    let c = strm.readChar()
    case c:
      of '{':
        result.add('{')
        result.add(strm.readLiteral)
        result.add('}')
      of '[':
        result.add('[')
        result.add(strm.readInterpelation)
        result.add(']')
      of '"':
        break
      of '\\':
        result.add("\\")
        result.add(strm.readChar())
      else: result.add(c)

# readCommand reads a command from stream
proc readCommand(strm: Stream): string =
  while not strm.atEnd():
    let c = strm.readChar()
    case c:
      of '{':
        result.add('{')
        result.add(strm.readLiteral)
        result.add('}')
      of '[':
        result.add('[')
        result.add(strm.readInterpelation)
        result.add(']')
      of '"':
        result.add("\"")
        result.add(strm.readSubstitution)
        result.add("\"")
      of '\\':
        result.add("\\")
        result.add(strm.readChar())
      of '\n', '\r', ';':
        break
      else: result.add(c)

# readVariable reads a variable
proc readVariable(strm: Stream): string =
  let c = strm.readChar()
  case c:
    of '{':
      result = strm.readLiteral()
    of '[':
      result.add(strm.readInterpelation)
    of '"':
      result.add(strm.readSubstitution())
    of ']', '}':
      stderr.write("Unexpected " & c & ".\n")
    else:
      result.add(c)
      while not strm.atEnd():
        let d = strm.readChar()
        case d:
          of 'a'..'z', 'A'..'Z', '0'..'9', '_', ':': result.add(d)
          else: break

# evalVariable finds a value of a variable
proc evalVariable(vari: string): string =
  if shellVariables.hasKey(vari):
    result=shellVariables[vari]
  elif envVariables.hasKey(vari):
    result = envVariables[vari]
  else:
    raise newShellVariableException("No such variable", vari)

# substitute does a substitution on the stream passed in
proc substitute(strm: Stream, delim=";\p"): string =
  while not strm.atEnd():
    let c = strm.readChar()
    if c in delim: break
    case c:
      of '{':
        result.add(strm.readLiteral)
      of '}':
        stderr.write("Extra closing brace")
      of '[':
        result.add(strm.readInterpelation.eval)
      of ']':
        stderr.write("Extra closing bracket")
      of '\\':
        result.add(strm.readChar())
      of '$': result.add(strm.readVariable.evalVariable)
      of ';', '\n', '\r': break
      else:
        result.add(c)

# parseCommand parses the string given into a list representing a command
proc parseCommand(cmd: string): seq[string] =
  var newitem = ""
  let strm = newStringStream(cmd)
  while true:
    if strm.atEnd():
      result.add(newitem)
      break
    let c = strm.readChar()
    case c:
      of '{':
        newitem.add(strm.readLiteral)
      of '[':
        newitem.add(strm.readInterpelation.eval)
      of '"':
        newitem.add(strm.readSubstitution.newStringStream.substitute)
      of '$': result.add(strm.readVariable.evalVariable)
      of ' ':
        if (result == @[]) and (newitem == ""): continue
        result.add(newitem)
        newitem = ""
        discard strm.skipSpaces()
      of ';', '\n', '\r':
        result.add(newitem)
        break
      else: newitem.add(c)

# runBuiltin runs a shell builtin
proc runBuiltin(builtin: string, args: openArray[string]): (string, int) =
  case builtin:
    of "set":
      if args.len < 2:
        raise newShellCommandException("Not enough arguments provided", "set")
      shellVariables[args[0]] = args[1..^1].join(" ")
      return ("", 0)
    of "setenv":
      if args.len < 2:
        raise newShellCommandException("Not enough arguments provided", "setenv")
      try:
        envVariables[args[0]] = args[1..^1].join(":")
        putEnv(args[0], args[1..^1].join(":"))
      except OSError as e:
        raise newShellCommandException(e.msg, "setenv")
      return ("", 0)
    of "echo":
      if args[0] == "-n":
        return (args[1..^1].join(" "), 0)
      return (args.join(" ") & "\p", 0)
    of "exit":
      if args.len == 0:
        quit(0)
      elif args.len == 1:
        try:
          quit(args[0].parseInt)
        except ValueError:
          quit(0)
      else:
        try:
          let exitcode=args[0].parseInt
          if exitcode > 0:
            stderr.write(args[1..^1].join(" "), "\p")
          else:
            stdout.write(args[1..^1].join(" "), "\p")
          quit(exitcode)
        except ValueError:
          stdout.write(args.join(" "), "\p")
          quit(0)
    of "if":
      var i=0
      while true:
        # if i is greater than args.high, we exit out. user probs doesn't want an else branch, we aren't going to stop him
        if i > args.high:
          break
    
        # Check if this is the last item, if so just execute it
        # It's an else branch basically
        if i == args.high:
          return ("", execute(args[i]))
        # Get the exitcode of the current arg (the condition)
        let exitcode = args[i].execute
        if exitcode == 0:
          # if the condition is true
          return ("", args[i+1].execute)
        else:
          # if the condition is false
          i.inc
          i.inc
          continue
    of "cd":
      try:
        if args.len == 0:
          shellVariables["LASTPWD"] = getCurrentDir()
          setCurrentDir(getHomeDir())
          return ("", 0)
        else:
          if args[0] == "-":
            let LASTPWD = getCurrentDir()
            setCurrentDir(shellVariables["LASTPWD"])
            shellVariables["LASTPWD"] = LASTPWD
            return ("", 0)
          shellVariables["LASTPWD"] = getCurrentDir()
          setCurrentDir(args[0])
          return ("", 0)
      except OsError as e:
        raise newShellCommandException("No such directory", "cd")
    else:
      raise newShellCommandException("No such builtin implemented", builtin)

# executes executes the command. For now it involves catting the command and printing it
proc execute(cmd: string): int =
  let strm=cmd.newStringStream
  while not strm.atEnd:
    let parsed = strm.readCommand.parseCommand
    if parsed == @[] or parsed == @[""]:
      continue
    let progname = parsed[0]
    let args = parsed[1..parsed.high]
    if progname in builtins:
      let res=runBuiltin(progname, args)
      stdout.write(res[0])
      result=res[1]
      continue
    try:
      let ps = startProcess(progname, args=args, options={poUsePath, poParentStreams})
      result =  ps.waitForExit
      ps.close
    except OSError as e:
      raise newShellCommandException("No such command", progname)
  strm.close

# eval evaluates the given string
proc eval(cmd: string): string =
  let strm=cmd.newStringStream
  while not strm.atEnd():
    let parsed = strm.readCommand.parseCommand
    if parsed == @[] or parsed == @[""]:
      continue
    let progname=parsed[0]
    let args = parsed[1..parsed.high]
    if progname in builtins:
      result = runBuiltin(progname, args)[0]
      continue
    try:
      let ps = startProcess(progname, args=args, options={poUsePath})
      result=ps.outputStream.readAll
      ps.close
    except OSError as e:
      raise newShellCommandException("No such command", progname)
  strm.close

when isMainModule:
  # ctrlc handler
  setControlCHook(ctrlc)
  let stdinstrm = stdin.newFileStream
  # Set the default prompt
  shellVariables["PROMPT"] = "$PWD \\$"
  stdout.write(shellVariables["PROMPT"].newStringStream.substitute)
  while not stdinstrm.atEnd():
    try:
      discard stdinstrm.readCommand.execute
      stdout.write(shellVariables["PROMPT"].newStringStream.substitute)
    except ShellError, ShellCommandError, ShellVariableError:
      let e = getCurrentException()
      stderr.write(e.msg, "\n")
    except InterruptCtrlC:
      continue