sockets.nim

import zmq, json, os, osproc, strutils, base64, streams, tables, options
import ./messages, ./utils, ./display

type
  Heartbeat* = object
    socket*: ZConnection
    alive: bool

  IOPub* = object
    socket*: ZConnection

  #TODO: encapsulate code better, why is stuff from shell exported?
  Shell* = object
    socket*: ZConnection
    pub*: IOPub     # keep a reference to pub so we can send status message
    count: Natural # Execution counter
    code: OrderedTable[string, string] # hold the code for cells that compile as a cellId:code table
    codeserver : Process # the codeserver process, needs to stay alive as long as we need to compile stuff
    servercode: string #code to build codeserver in hcr
    executingCellId: string # id of the executing cell id

  Control* = object
    socket*: ZConnection

  StdIn* = object
    socket*: ZConnection

  Channels* = concept s
    s.socket is ZConnection

# forward decl
proc updateCodeServer(shell: var Shell, firstInit=false): tuple[output: string, exitcode: int]

# Helpers that generally work on all channel objects
proc hasMsgs*(s: Channels): bool = getsockopt[int](s.socket, EVENTS) == 3
proc close*(s: Channels) = s.close()
proc receiveMsg(s: Channels): WireMessage = decode(s.socket.recv_multipart)

proc sendMsg*(s: Channels, msg: WireMessage) =
  let encoded = encode(msg)
  #debug "Encoded ", reply_type
  s.socket.send_multipart(encoded)

## Heartbeat Socket
proc createHB*(ip: string, hbport: BiggestInt): Heartbeat =
  ## Create the heartbeat socket
  result.socket = zmq.listen("tcp://" & ip & ":" & $hbport, zmq.REP)
  result.alive = true

proc pong*(hb: Heartbeat) =
  var msg = hb.socket.receive()
  hb.socket.send(msg)

proc close*(hb: var Heartbeat) =
  hb.alive = false
  hb.socket.close()

## IOPub Socket
proc createIOPub*(ip: string, port: BiggestInt): IOPub =
  ## Create the IOPub socket
# TODO: transport
  result.socket = zmq.listen("tcp://" & ip & ":" & $port, zmq.PUB)

proc receive*(pub: IOPub) =
  ## Receive a message on the IOPub socket
  let recvdmsg: WireMessage = pub.receiveMsg()
  debug "pub received:\n", $recvdmsg


## Shell Socket

# TODO: move stuff related to codeserver somewhere else
const codeserver = staticRead("codeserver.nim") # $1 -> codecell filename
const jnTempDir* = getHomeDir() / ".jupyternim"

# ORDER IS IMPORTANT
let defaultFlags = ["", # for HCR
                      "-d:release",
                      "--verbosity:0",
                      "--hint[Processing]:off"]


var requiredFlags : array[2,string]= ["-d:jupyter", "-o:"]

when defined useHcr:
  const initialCodecell = staticRead("initialCodeCell.nim")
else:
  const initialCodecell = "discard \"Generated by JupyterNim\""

var flags: seq[string] = @defaultFlags

proc writeCodeFile(shell:Shell) =
  ## Write out the file composed by the cells that were run until now.
  ## The last cell is wrapped in a proc so that it gets run by the codeserver
  ## and produces output.
  var res = ""
  #debug shell.code
  for k, cell in shell.code:
    if k == shell.executingCellId:
      #res.add("jnparentmsg = " & ($(%* shell.pub.lastmsg)).escapeJson) #omg the json escaping
      when defined useHcr:
        # wrap the last cell in the hoist macro:
        res.add("hoist:\n")
        res.add(cell.indent(2) & "\n") # indent to avoid compilation errors
      else:
        res.add("\necho \"#>newcodeout\"\n") # so that we can cut out unneeded outputs
        res.add(cell & "\n")
      break # we don't need all the file if just changing one line
    else:
      res.add(cell & "\n")
  writeFile(jnTempDir / JNfile & "codecells.nim", res)

proc startCodeServer(shell: var Shell): Process =
  ## Start the nimcodeserver process (the hcr main program)
  debug "HCR: confirm codeserver.exe exists: ", fileExists( jnTempDir / JNoutCodeServerName)
  if not fileExists( jnTempDir / JNoutCodeServerName):
    debug "forcing codeserver to be rebuilt and reinited"
    flags.add("-f") # maybe forcing a rebuild ?
    discard shell.updateCodeServer(firstInit=true)
    discard flags.pop()
  else:
    debug jnTempDir / JNoutCodeServerName & " already exists"

  result = startProcess(jnTempDir / JNoutCodeServerName)

proc updateCodeServer(shell: var Shell, firstInit=false): tuple[output: string, exitcode: int] =
  ## Write out the source code if firstInit==true, then
  ## write the code file (the "logic")
  ## compile it
  ## Returns the compiler output as a string
  if firstInit:
    debug "Write out codeserver"
    # hcr needs at least one extra cell to generate the initial proc
    shell.executingCellId = "hcrFirst"
    shell.code[shell.executingCellId] = "discard \"Generated for Hcr\""
    writeFile(jnTempDir / JNfile & "codeserver.nim", shell.servercode)

  debug "Write out codecells"
  writeCodeFile(shell)

  when defined useHcr:
    if not firstInit:
      debug "Ensuring codeserver is alive"
      if not shell.codeserver.running:
        debug "The codeserver died, trying to restart it..."
        shell.codeserver = shell.startCodeServer()

  debug "Recompile codeserver"
  debug r"nim c " & flatten(flags) & flatten(requiredFlags) & jnTempDir / JNfile & "codecell|server.nim"
  when defined useHcr:
    const file = "codeserver.nim" # compile the codeserver
  else:
    const file = "codecells.nim"
  result = execCmdEx(r"nim c " & flatten(flags) & flatten(requiredFlags) & escape(jnTempDir / JNfile & file)) # compile the codeserver

proc createShell*(ip: string, shellport: BiggestInt, pub: IOPub): Shell =
  ## Create a shell socket
  #debug "shell at ", ip, " ", shellport

  # flags setup
  requiredFlags[1].add(escape(jnTempDir / JNoutCodeServerName)) # complete the -o: flags
  flags.add("-p:" & escape(getCurrentDir())) # can't importc at compile time
  when not defined(release):
    flags[1] = "-d:debug" #switch release to debug for the compiled file too
    flags[2] = ""#"--verbosity:3" # remove verbosity:0 flag
  when defined useHcr:
    flags[0] = "--hotcodereloading:on" # enable hotcodereloading
    result.servercode = codeserver % [JNfile & "codecells"] # fill in codecells filaname
  # flags

  result.socket = zmq.listen("tcp://" & ip & ":" & $shellport, zmq.ROUTER)
  result.pub = pub
  # add the import to the codecells of shell,
  # this way it will be there when generating the code to be run
  result.code = initOrderedTable[string, string]()
  result.code["initialCell"] = initialCodecell
  when defined useHcr:
    let tmp = result.updateCodeServer(firstInit=true)
    #debug tmp.output
    result.codeserver = result.startCodeServer()

proc handleKernelInfo(s: Shell, m: WireMessage) =
  s.sendMsg( kernelInfoMsg(m) )

const MagicsStrings = ["#>flags", "#>clear all"]
#[TODO: find an efficient way to do the following
        (since it's unlikely that a lot of flags are present, looping all lines is not a good idea)
proc handleMagics(codeseq: seq[string])
  for line in codeseq:
    if line.startsWith(MagicsStrings[0]):
      flags = code[MagicsStrings[0].len+1..^1].split()
      debug "with custom flags:", flags.flatten
    elif line.startsWith...
      ]#

proc styledOutput(s:string):string =
  result = ""
  if s == "": return
  let last = s.countLines-1
  var i = 0
  for line in s.splitLines:
    var tmp = "<span>"
    if line.startsWith("Hint:"):
      tmp.add("<small style=\"color: #00b479\">")
      tmp.add(line[0..5])
      tmp.add("</small>")
      tmp.add(line[6..^1])
    elif line.startsWith("Warning:"):
      tmp.add("<small style=\"color: #f7ff00\">")
      tmp.add(line[0..8])
      tmp.add("</small>")
      tmp.add(line[8..^1])
    elif line.contains("Hint: "):
      tmp.add(" <small style=\"color: #00b479\">Hint: </small>")
      tmp.add(line[line.find("Hint: ")+len("Hint: ")..^1])
    elif line.contains("Warning: "):
      tmp.add(" <small style=\"color: #f7ff00\">Warning: </small>")
      tmp.add(line[line.find("Warning: ")+len("Warning: ")..^1])
    else:
      tmp.add(line)
    tmp.add("</span>")
    if not (i==last):
      tmp.add("<br/>")
    inc i
    result.add(tmp)


proc analyzeError(s:string): tuple[evalue:string, trace: seq[string]] =
  for line in s.splitLines:
    # don't care if empty
    if line.isEmptyOrWhitespace: continue
    if line.contains(") Error: "):
      result.evalue = line.rsplit(") Error: ")[1]
    result.trace.add(line)
  result.trace.add("Error: " & result.evalue) # hide path, useless here no?

proc handleExecute(shell: var Shell, msg: WireMessage) =
  ## Handle the ``execute_request`` message
  #debug "HANDLEEXECUTE\n", msg
  inc shell.count

  #TODO: in the future outputs will become a seq[string] so that
  # streams can be sent line by line
  let
    code = msg.content["code"].str # The code to be executed
    # this "fixes" cases in which the frontend doesn't expose the cellid, by
    # requiring users to add a #>cellId:<something> to their cell code.
    # If they don't, we use the message id of the execute_req message.
    # Problem: this last way destroys the ability to track a cell since there's no
    # way to map the cell being re run to its old code, causing compilation errors
    # very fast
    # TODO: follow up issues for nteract to expose this in the cell
    cellId =  if msg.metadata.hasKey("cellId"): msg.metadata["cellId"].str
              elif code.startsWith("#>cellId"): code.splitLines()[0]
              else: $shell.count # msg.header.msg_id

  shell.executingCellId = cellId
  # to prevent reordering of cells due to popping them in case they don't compile,
  # keep a record of the last code that was in this spot.
  let lastTimeCellCode = shell.code.getOrDefault(shell.executingCellId)

  # TODO: move the logic that deals with magics and flags somewhere else
  if code.contains("#>flags"):
    let
      flagstart = code.find("#>flags")+"#>flags".len+1
      nwline = code.find("\u000A", flagstart)
      flagend = if nwline != -1: nwline else: code.len
    flags = code[flagstart..flagend].split()

    debug "with custom flags:", flags.flatten

  if code.contains("#>clear all") and dirExists(jnTempDir):
    debug "Cleaning up..."
    flags = @defaultFlags
    # no need to delete files as they will be overwritten anyway
    shell.code.clear
    shell.code["initialCell"] = initialCodecell
    shell.executingCellId = ""
    let tmp = shell.updateCodeServer()
    debug tmp
    # TODO: resets flags properly
    when defined useHcr:
      shell.codeserver = shell.startCodeServer

  # Send via iopub the block about to be executed
  shell.pub.sendMsg(executeInputMsg(shell.count, code, msg.some))

  # Compile and send compilation messages to jupyter's stdout
  shell.code[cellId] = code
  var compilationResult = shell.updateCodeServer()
  # debug "file before:"
  # debug readFile(jnTempDir / "codecells.nim")
  # debug "file end"

  # debug "server has data: ", shell.codeserver.hasData

  var status: Status = Status.ok
  var streamName: string

  if compilationResult.exitcode != 0:
    # look at python code, do they send as many error messages?

    # return early if the code didn't compile
    status = Status.error
    streamName = "stderr"

    # execution not ok, remove last cell

    debug "Reverted cell from:" & shell.code[cellId]
    shell.code[cellId] = lastTimeCellCode # restore the code that was here before this failure

    #debug "file after:"
    #debug readFile(jnTempDir / "codecells.nim")
    #debug "file end"

    #shell.pub.sendMsg( streamMsg(streamName, compilationResult.output, some msg))
    # C:\Users\stisa\.jupyternim\codecells.nim(6, 1) Error: invalid indentation
    # Compile error: Error
    let betterError = analyzeError(compilationResult.output)
    shell.pub.sendMsg( errorMsg("Compilation Error", betterError.evalue,
                                betterError.trace, some msg))
    # Tell the frontend execution failed from shell
    sleep(1000)

    shell.sendMsg(replyErrorMsg(shell.count,"Compilation Error", betterError.evalue,
                                betterError.trace, msg))

    return

  streamName = "stdout"
  if not compilationResult.output.strip().len == 0:
    # avoid sending back empty strings, avoids wasting space
    # Send compiler messages
    shell.pub.sendMsg( displayDataMsg(showHtml(styledOutput(compilationResult.output)), some msg))

  # Since the compilation was fine, run code and send results with iopub
  var exec_out: string
  # run the new code
  when defined useHcr:
    shell.codeserver.inputStream.writeLine("#runNimCodeServer")
    shell.codeserver.inputStream.flush

    #debug "trying to read all...", shell.codeserver.hasData

    var donewriting = false
    while not doneWriting:
      let tmp = shell.codeserver.outputStream.readLine
      if tmp.rfind("#serverReplied") != -1:
        donewriting = true
        break # we dont want the last message anyway
      exec_out &= tmp & "\n"
  else:
    exec_out = execProcess(jnTempDir / JNoutCodeServerName, getCurrentDir())
    #debug exec_out
    # we are only interested in new output
    let execoutsplit = exec_out.rsplit("#>newcodeout\n") #echo added a newline
    #debug "Split length ", len(execoutsplit)
    if execoutsplit.len > 0: exec_out = execoutsplit[^1]

  #debug "done reading, read: ", exec_out

  # TODO: don't assume no errors are possible at runtime,
  #       check for errors there too

  # Handle display_data
  # FIXME: document this!
  # debug "Handling display data"
  # exec_out only has output from the last cell, so it should be safe to
  # add all the jndd blocks here.
  # Maybe just put code to open a socket and send a display_data message in a utils lib
  # and let the user/plotlib deal with this? Much cleaner?
  # https://nim-lang.org/docs/net.html
  while exec_out.contains("#<jndd>"):
    let
      ddstart = exec_out.find("#<jndd>#")
      ddend = exec_out.find("#<outjndd>#")
      dddata = exec_out[ddstart+len("#<jndd>#")..<ddend]
      parseddata = parseJson(dddata)

    #shell.pub.sendMsg(displayDataMsg(parseddata, some msg))
    shell.pub.sendMsg(displayExecResMsg(shell.count, parseddata, some msg))

    #debug "b",exec_out
    exec_out = exec_out.replace("#<jndd>#" & dddata & "#<outjndd>#", "") # clear out the data from the output
    #debug "a",exec_out

  shell.pub.sendMsg(displayExecResMsg(shell.count, %* {"data": %* {"text/plain": exec_out}}, some msg))
  #shell.pub.sendMsg(execResultMsg(shell.count, exec_out, some msg))

  # Tell the frontend execution was ok, or not from shell
  shell.sendMsg( execReplyMsg(shell.count, status, msg))

proc parseNimsuggest(nims: string): tuple[found: bool, data: JsonNode] =
  # nimsuggest output is \t separated
  # http://nim-lang.org/docs/nimsuggest.html#parsing-nimsuggest-output
  discard

proc handleIntrospection(shell: Shell, msg: WireMessage) =
  #[ reply
    content = {
    # 'ok' if the request succeeded or 'error', with error information as in all other replies.
    'status' : 'ok',
    # found should be true if an object was found, false otherwise
    'found' : bool,
    # data can be empty if nothing is found
    'data' : dict,
    'metadata' : dict,
    }
  ]#
  debug "TODO: Inspect request:\n",msg.content
  let code = msg.content["code"].str
  let cpos = msg.content["cursor_pos"].num.int
  if code[cpos] == '.':
    discard # make a call to sug in nimsuggest sug <file> <line>:<pos>
  elif code[cpos] == '(':
    discard # make a call to con in nimsuggest con <file> <line>:<pos>
  # TODO: ask nimsuggest about the code
  shell.sendMsg(inspectReplyMsg(Status.ok, parent= some msg))

proc handleCompletion(shell: Shell, msg: WireMessage) =

  let code: string = msg.content["code"].str
  let cpos: int = msg.content["cursor_pos"].num.int - 1
  var
    wordpart = code
    sw = cpos

  if code.len>1:
    debug msg.content
    # backtrack to word start (ie whitespace? CHECK: is it true?)
    while sw > 0 and not (code[sw] in Whitespace):
      sw -= 1
    wordpart = code[sw..cpos] # partial word

  var matches: seq[string] = @[] # list of all matches

  # Snippets
  if "proc".startswith(wordpart):
    matches &= ("proc name(arg: type): returnType = \n  #proc")
  elif "if".startswith(wordpart):
    matches &= ("if (expression):\n  #then")
  elif "method".startswith(wordpart):
    matches &= ("method name(arg: type): returnType = \n  #method")
  elif "iterator".startswith(wordpart):
    matches &= ("iterator name(arg: type): returnType = \n  #iterator")
  elif "array".startswith(wordpart):
    matches &= ("array[length, type]")
  elif "seq".startswith(wordpart):
    matches &= ("seq[type]")
  elif "for".startswith(wordpart):
    matches &= ("for index in iterable):\n  #for loop")
  elif "while".startswith(wordpart):
    matches &= ("while(condition):\n  #while loop")
  elif "block".startswith(wordpart):
    matches &= ("block name:\n  #block")
  elif "case".startswith(wordpart):
    matches &= ("case variable:\nof value:\n  #then\nelse:\n  #else")
  elif "try".startswith(wordpart):
    matches &= ("try:\n  #something\nexcept exception:\n  #handle exception")
  elif "template".startswith(wordpart):
    matches &= ("template name (arg: type): returnType =\n  #template")
  elif "macro".startswith(wordpart):
    matches &= ("macro name (arg: type): returnType =\n  #macro")

  # Single word matches
  let single = ["int", "float", "string", "addr", "and", "as", "asm", "atomic", "bind", "break", "cast",
                "concept", "const", "continue", "converter", "defer", "discard",
                "distinct", "div", "do",
                "elif", "else", "end", "enum", "except", "export", "finally",
                "for", "from", "func",
                "generic", "import", "in", "include", "interface", "is",
                "isnot", "let", "mixin", "mod",
                "nil", "not", "notin", "object", "of", "or", "out", "ptr",
                "raise", "ref", "return", "shl",
                "shr", "static", "tuple", "type", "using", "var", "when",
                "with", "without", "xor", "yield"]

  #magics = ['#>loadblock ','#>passflag ']

  # Add all matches to our list
  #TODO: rewrite this and document
  matches = matches & (filter(single) do (x: string) -> bool: x.startsWith(wordpart))

  # TODO completion+nimsuggest

  # debug msg
  shell.sendMsg( completeReplyMsg( Status.ok, matches, sw, cpos+1, some msg))

proc handleHistory(shell: Shell, msg: WireMessage) =
  debug "Unhandled: history"
  shell.sendMsg(historyReplyMsg( %*{}, some msg))


proc handleCommInfo(s: Shell, msg: WireMessage) =
  debug "Unhandled: CommInfoReq"
  debug msg.content
  if msg.content.hasKey("target_name"):
    debug "CommInfo about ", msg.content["target_name"].getStr
    # A dictionary of the comms, indexed by uuids (comm_id).
    #[content = {  'comms': { comm_id: { 'target_name': str,  },    }, }]#
    s.sendMsg( commReplyMsg( some msg))
  else:
    s.sendMsg( commReplyMsg( some msg))

proc handle(s: var Shell, m: WireMessage) =
  debug "shell: handle ", m.kind
  case m.kind
  of kernelInfoRequest:
    debug "Sending Kernel info"
    handleKernelInfo(s, m)
  of executeRequest:
    handleExecute(s, m)
  of shutdownRequest:
    debug "kernel wants to shutdown"
    quit()
  of inspectRequest: handleIntrospection(s, m)
  of completeRequest: handleCompletion(s, m)
  of historyRequest: handleHistory(s, m)
  of commInfoRequest: handleCommInfo(s, m)
  else:
    debug "unhandled message: ", m.kind

proc receive*(shell: var Shell) =
  ## Receive a message on the shell socket, decode it and handle operations
  let recvdmsg: WireMessage = shell.receiveMsg()
  debug "shell: ", $recvdmsg.kind
  #debug recvdmsg.content
  #debug "end shell"
  shell.pub.sendMsg( statusMsg(State.busy, recvdmsg.some) )
  shell.handle(recvdmsg)
  shell.pub.sendMsg( statusMsg(State.idle, recvdmsg.some) )

proc close*(sl: var Shell) =
  sl.socket.close()
  when defined useHcr:
    sl.codeserver.terminate()


## Control socket
proc createControl*(ip: string, port: BiggestInt): Control =
  ## Create the control socket
  result.socket = zmq.listen("tcp://" & ip & ":" & $port, zmq.ROUTER)


proc handle(c: Control, m: WireMessage) =
  if m.kind in {shutdown_request, interrupt_request}:
    #var content : JsonNode
    debug "shutdown requested"
    #content = %* { "restart": false }
    c.sendMsg( shutdownMsg(m.kind, some m))
    quit()
  else:
    debug "Control: unhandled message ", m.kind

proc receive*(cont: Control) =
  ## Receive a message on the control socket and handle operations
  let recvdmsg: WireMessage = cont.receiveMsg()
  debug "received: ", $recvdmsg.kind
  cont.handle(recvdmsg)


## Stdin socket
proc createStdIn*(ip: string, port: BiggestInt): StdIn =
  ## Create the stdin socket
  result.socket = zmq.listen("tcp://" & ip & ":" & $port, zmq.ROUTER)

proc requestInput(c: StdIn, inputmsg: string, m: WireMessage) {.used.}=
  doAssert m.kind == input_request
  # TODO: inject a proc readLine(stdin:File):string that prints
  # jnstdin to stdin, make execute watch for it in the stdin of startProcess(nimcodeserver)
  # and then call requestInput and wait for stdin.receive which will be then print the
  # `value` to stdin
  debug "input requested"
  c.sendMsg( inputReqMsg(inputmsg))

proc handle(c: StdIn, m: WireMessage) =
  debug "Handle input reply"
  # need to write this to the stdin of the executing compiled code
  # means breaking a bunch of encapsulation? Maybe just add this
  # logic to execute, by calling a requestInput
  echo m.content["value"]


proc receive*(sin: StdIn) =
  ## Receive a message on the control socket and handle operations
  let recvdmsg: WireMessage = sin.receiveMsg()
  debug "received: ", $recvdmsg.kind
  sin.handle(recvdmsg)