# vim: fileencoding=utf8 ft=nim et sw=2 ts=2 sts=2
import os
import posix
#import asyncfutures
from asyncfutures import
  Future, newFuture, fail, complete, `or`, `and`, `callback=`, failed, read
#import asyncdispatch
from asyncdispatch import
  register, unregister, addRead, addWrite, AsyncFD
#import strformat
#import sequtils
#import strutils

const
  ERR_PERM = 100
  ERR_TEMP = 101
  ERR_EXEC_OTHER = 126
  ERR_EXEC_ENOENT = 127
  DEBUG_OUT {.intdefine.}: int = 0

proc writeHelp(errorcode: int): void =
  write(stdout, "writeHelp()\n")
  quit(errorcode)

proc writeVersion(): void =
  write(stdout, "writeVersion()\n")
  quit(0)

type
  Args = tuple[
    endl: string,
    hideendl: bool,
    color: bool,
    color_escapes: array[4, string],
    exe: seq[string]
  ]
proc posixParseArgs(): Args =
  var args: Args = ("\n", false, true, ["32", "36", "31", "35"], @[])
  let argc = paramCount()
  var n = 1
  while n <= argc:
    let a = paramStr(n)
    if len(a) >= 1 and a[0] == '-':
      case a
      of "-h", "--help": writeHelp(0)
      of "-v", "--version": writeVersion()
      of "-H", "--hide-newlines": args.hideendl = true
      of "-c", "--nocolor": args.color = false
      of "-n": args.endl = "\n"
      of "-r": args.endl = "\r"
      of "-D": args.endl = "\r\n"
      of "-N": args.endl = ""
      of "-E", "--eol":
        inc(n)
        if n > argc:
          write(stderr, "Missing argument to " & a & "\n")
          writeHelp(ERR_PERM)
        args.endl = paramStr(n)
      of "-C", "--colors", "--colours":
        for i in 0..3:
          inc(n)
          if n > argc:
            write(stderr, "Missing argument to " & a & "\n")
            writeHelp(ERR_PERM)
          args.color_escapes[i] = paramStr(n)
      else:
        write(stderr, "Unrecognised parameter: \"" & a & "\"\n")
        writeHelp(ERR_PERM)
    else:
      break
    inc(n)
  while n <= argc:
    args.exe &= paramStr(n)
    inc(n)
  return args


proc libc_fatal(msg: string, exitcode=ERR_TEMP): void =
  var err = posix.errno
  write(stderr, "fatal: " & msg & " ")
  write(stderr, posix.strerror(err))
  write(stderr, "\n")
  quit(exitcode)


proc copy_data(
  fd_from:cint,
  fd_to:cint,
  pp_cb:proc(data:string),
): Future[void] =
  var flags: cint
  
  flags = posix.fcntl(fd_from, posix.F_GETFL)
  if flags == -1: libc_fatal("fcntl():")
  if posix.fcntl(fd_from, posix.F_SETFL, flags or posix.O_NONBLOCK) == -1:
    libc_fatal("fcntl():")
  asyncdispatch.register(fd_from.AsyncFD)

  flags = posix.fcntl(fd_to, posix.F_GETFL)
  if flags == -1: libc_fatal("fcntl():")
  if posix.fcntl(fd_to, posix.F_SETFL, flags or posix.O_NONBLOCK) == -1:
    libc_fatal("fcntl():")
  asyncdispatch.register(fd_to.AsyncFD)

  let buf_size = 4096  # we could query posix.PC_PIPE_BUF, but this is common
  var
    future_result = newFuture[void]("copy_data")
    buffer = newString(buf_size)
    buf_written = 0
    reading = true
    writing = false

  when DEBUG_OUT.bool:
    echo fmt"reading[{fd_from}]: {reading}  writing[{fd_to}]: {writing}"
  # Mutually referential function declaration
  proc read_cb(fd: AsyncFD): bool {.gcsafe.}
  proc write_cb(fd: AsyncFD): bool {.gcsafe.}

  proc read_cb(fd: AsyncFD): bool {.gcsafe.} =
    when DEBUG_OUT.bool:
      echo fmt"read_cb(" & $fd.int & ")"
      echo fmt"reading[{fd_from}]: {reading}  writing[{fd_to}]: {writing}"
    assert(reading)
    assert(not writing)
    let res = posix.read(fd_from, addr buffer[0], buf_size)
    if res < 0:
      let lastError = osLastError()
      if lastError.int32 in {EINTR, EWOULDBLOCK, EAGAIN}:
        when DEBUG_OUT.bool:
          echo "return false"
        return false  # try again
      else:
        future_result.fail(newException(OSError, osErrorMsg(lastError)))
    elif res == 0:
      # End of "file"
      pp_cb("")  # Flush buffer if any
      asyncdispatch.unregister(fd_from.AsyncFD)
      asyncdispatch.unregister(fd_to.AsyncFD)
      when DEBUG_OUT.bool:
        echo "close(", fd_to, ")"
      if posix.close(fd_to) != 0:
        let lastError = osLastError()
        future_result.fail(newException(OSError, osErrorMsg(lastError)))
      else:
        future_result.complete()
    else:
      buffer.setLen(res)
      when DEBUG_OUT.bool:
        echo "read[" & $fd.int & "] -> " & repr(buffer)
      pp_cb(buffer)
      writing = true
      when DEBUG_OUT.bool:
        echo fmt"reading[{fd_from}]: {reading}  writing[{fd_to}]: {writing}"
      asyncdispatch.addWrite(fd_to.AsyncFD, write_cb)
    reading = false
    when DEBUG_OUT.bool:
      echo fmt"reading[{fd_from}]: {reading}  writing[{fd_to}]: {writing}"
      echo "return true"
    return true


  proc write_cb(fd: AsyncFD): bool {.gcsafe.} =
    when DEBUG_OUT.bool:
      echo fmt"write_cb(" & $fd.int & ")"
      echo fmt"reading[{fd_from}]: {reading}  writing[{fd_to}]: {writing}"
    assert(writing)
    assert(not reading)
    let remaining = buffer.len - buf_written
    when DEBUG_OUT.bool:
      echo fmt"written: {buf_written}  remaining: {remaining}"
    let res = posix.write(fd_to, addr buffer[buf_written], remaining.cint)
    if res < 0:
      let lastError = osLastError()
      if lastError.int32 in {EINTR, EWOULDBLOCK, EAGAIN}:
        when DEBUG_OUT.bool:
          echo "return false"
        return false  # try again
      else:
        future_result.fail(newException(OSError, osErrorMsg(lastError)))
        writing = false
        when DEBUG_OUT.bool:
          echo fmt"reading[{fd_from}]: {reading}  writing[{fd_to}]: {writing}"
          echo "return true"
        return true
    else:
      buf_written.inc(res)
      if res < remaining:
        return false  # more to write
      else:
        reading = true
        asyncdispatch.addRead(fd_from.AsyncFD, read_cb)
        buf_written = 0
        buffer.setLen(buf_size)
        writing = false
        when DEBUG_OUT.bool:
          echo fmt"reading[{fd_from}]: {reading}  writing[{fd_to}]: {writing}"
          echo "return true"
        return true

  asyncdispatch.addRead(fd_from.AsyncFD, read_cb)
  return future_result


type DisplayMode = enum
  none,
  normal,
  escaped

proc pp_writer(
  normal: string,
  escaped: string,
  options: Args,
): proc(data:string) =
  var
    to_print = ""
    dm: DisplayMode = DisplayMode.none

  proc write_out() =
    ## Write out and empty to_print buffer.
    if to_print.len() == 0:
      return
    to_print.add("\x1b[0m")
    dm = DisplayMode.none
    write(stderr, to_print)
    to_print.setLen(0)

  proc add_char(c: char) =
    ## Format character and add it to to_print buffer.
    case c
    of ' ' .. '~':
      if dm != DisplayMode.normal:
        to_print.add("\x1b[")
        to_print.add(normal)
        to_print.add("m")
        dm = DisplayMode.normal
      to_print.add(c)
    else:
      if dm != DisplayMode.escaped:
        to_print.add("\x1b[")
        to_print.add(escaped)
        to_print.add("m")
        dm = DisplayMode.escaped
      to_print.addEscapedChar(c)

  case options.endl.len
  of 0:
    proc pp(data: string): void =
      for c in data:
        add_char(c)
      write_out()
    return pp
  of 1:
    let endl: char = options.endl[0]
    if options.hideendl:
      proc pp(data: string): void =
        for c in data:
          if c == endl:
            to_print.add('\n')
          else:
            add_char(c)
        write_out()
      return pp
    else:
      proc pp(data: string): void =
        for c in data:
          add_char(c)
          if c == endl:
            to_print.add('\n')
        write_out()
      return pp
  else:
    if options.hideendl:
      var ringbuf = ""
      var ringpos = 0
      let endl_len = options.endl.len
      var endl_chars: set[char] = {}
      for c in options.endl:
        endl_chars.incl(c)

      proc ring_check(): bool =
        ## Returns true when newline was found.
        for n in 0 .. (endl_len - 1):
          let i = if n > ringpos: endl_len + ringpos - n else: ringpos - n
          if options.endl[^(n + 1)] != ringbuf[i]:
            return false
        return true

      proc pp(data: string): void =
        if data.len() == 0:
          # End of stream, flush buffer
          if ringbuf.len > 0:
            let l = ringbuf.len()
            for n in 1 .. l:
              add_char(ringbuf[(ringpos + n) mod l])
          write_out()
          return

        for c in data:
          if c notin endl_chars:
            if ringbuf.len > 0:
              let l = ringbuf.len()
              for n in 1 .. l:
                add_char(ringbuf[(ringpos + n) mod l])
              ringbuf.setLen(0)
              ringpos = 0
            add_char(c)
            continue

          if ringbuf.len < endl_len:
            ringbuf.add(c)
            ringpos = ringbuf.high()
          else:
            ringpos = (ringpos + 1) mod endl_len
            add_char(ringbuf[ringpos])
            ringbuf[ringpos] = c

          if ringbuf.len < endl_len:
            continue

          if ring_check():
            to_print.add('\n')
            ringbuf.setLen(0)
            ringpos = 0
        write_out()
      return pp
      
    else:
      var ringbuf = ""
      var ringpos = 0
      let endl_len = options.endl.len

      proc ring_add(c: char): bool =
        ## Returns true when newline was found.
        if ringbuf.len < endl_len:
          ringbuf.add(c)
          ringpos = ringbuf.high()
        else:
          ringpos = (ringpos + 1) mod endl_len
          ringbuf[ringpos] = c
        if ringbuf.len < endl_len:
          return false
        for n in 0 .. (endl_len - 1):
          let i = if n > ringpos: endl_len + ringpos - n else: ringpos - n
          if options.endl[^(n + 1)] != ringbuf[i]:
            return false
        ringbuf.setLen(0)
        ringpos = 0
        return true

      proc pp(data: string): void =
        for c in data:
          add_char(c)
          if ring_add(c):
            to_print.add('\n')
        write_out()
      return pp


proc main(): void =
  let args = posixParseArgs()
  when DEBUG_OUT.bool:
    write(stderr, "args: " & repr(args) & "\n")

  if len(args.exe) == 0:
    echo "No arguments"
    writeHelp(ERR_PERM)

  # despite documentation it doesn't search local dir on unix
  # unless executable name contains /
  let exe = os.findExe(args.exe[0])
  if len(exe) == 0:
    echo "Command not found: " & args.exe[0]
    quit(ERR_EXEC_ENOENT)

  var
    pipein: array[0 .. 1, cint]
    pipeout: array[0 .. 1, cint]

  if posix.pipe(pipein) != 0: libc_fatal("pipe():")
  if posix.pipe(pipeout) != 0: libc_fatal("pipe():")

  var pid = fork()
  case pid
  of -1: libc_fatal("fork():")
  of 0:
    # Child process
    if posix.dup2(pipein[0], 0) == -1: libc_fatal("dup2():")
    if posix.close(pipein[1]) != 0: libc_fatal("close():")
    if posix.dup2(pipeout[1], 1) == -1: libc_fatal("dup2():")
    if posix.close(pipeout[0]) != 0: libc_fatal("close():")
    if posix.execv(exe, allocCStringArray(args.exe)) == -1:
      let exitcode =
        if posix.errno == ENOENT: ERR_EXEC_ENOENT
        else: ERR_EXEC_OTHER
      libc_fatal("exec():", exitcode)
    write(stderr, "fatal: exec() did not execute!")
    quit(ERR_TEMP)
  else:
    # Parent process
    when DEBUG_OUT.bool:
      echo "pipein: " & repr(pipein)
      echo "pipeout: " & repr(pipeout)
    if close(pipein[0]) != 0: libc_fatal("close():")
    if close(pipeout[1]) != 0: libc_fatal("close():")

    # input: 0 -> pipein[1]
    # output: pipeout[0] -> 1
    
    var
      rw_stdin: Future[void] = copy_data(0, pipein[1], pp_writer(
        args.color_escapes[0], args.color_escapes[1], args,
        ))
      rw_stdout: Future[void] = copy_data(pipeout[0], 1, pp_writer(
        args.color_escapes[2], args.color_escapes[3], args,
        ))
      
    asyncdispatch.waitFor(rw_stdin or rw_stdout)
    if rw_stdin.failed:
      rw_stdin.read  # re-raise
    if rw_stdout.failed:
      rw_stdout.read  # re-raise

    asyncdispatch.waitFor(rw_stdin and rw_stdout)
    if rw_stdin.failed:
      rw_stdin.read  # re-raise
    if rw_stdout.failed:
      rw_stdout.read  # re-raise

    var
      status: cint
      ret: cint
      err = posix.EINTR

    while err == posix.EINTR:
      ret = waitpid(pid, status, 0)
      if ret == -1:
        err = posix.errno
      elif ret == pid:
       #define wait_estatus(w) (WIFSIGNALED(w) ? 128 + WTERMSIG(w) :
       #                        WEXITSTATUS(w) >= 128 ? 128 : WEXITSTATUS(w))
       quit(if WIFSIGNALED(status): 128 + WTERMSIG(status)
            elif WEXITSTATUS(status) >= 128: 128
            else: WEXITSTATUS(status))
      else:
        write(stderr, "fatal: waitpid(): unexpected return value: " & repr(ret))
        quit(ERR_TEMP)
    libc_fatal("waitpid():")

when isMainModule:
    main()