## Ingest a file's content into Sonic
import os
import strutils
import sonic

var
  verbose: bool = false

const
  USAGE_DOC = staticRead("static/usage.txt")

proc showUsage(msg = "") =
  ## Display a usage message and quit
  if msg != "":
    stderr.writeLine "ERROR: $#\n" % msg
  quit USAGE_DOC.replace("${app}", getAppFilename()),
       (if msg == "": 0 else: 10)

proc envParam(name: string; default = ""): string =
  ## Return an environment parameter or return an error
  let eVal = getEnv(name, default)
  if eVal == "":
    quit("Expected a value: $" & name, 1)
  return eVal

proc getChannel(mode = SonicChannel.Ingest): Sonic =
  ## Return a channel based on command line params
  let
    host = envParam("SONIC_HOST")
    port = envParam("SONIC_PORT")
    secret = envParam("SONIC_SECRET")
  try:
    return open(host, port.parseInt(), secret, mode)
  except ValueError:
    let err = getCurrentExceptionMsg()
    quit "Not a valid number: $#\n$#" % [port, err], 1

proc consolidate(channel: Sonic) =
  ## Trigger a consolidation
  discard channel.execCommand("TRIGGER", @["consolidate"])

proc close(chn: Sonic) =
  ## Close a sonic channel
  let outp = chn.quit()
  if verbose:
    stderr.writeLine "<closing $# channel: #$>" % [$(chn.channel), outp]

proc intAt(args: openArray[string]; pos: int; default: int): int =
  ## Parse a positional parameter if it exists, if not use default
  if args.len-1 < pos or args[pos] == "":
    return default
  try:
    return args[pos].parseInt()
  except:
    let err = getCurrentExceptionMsg()
    quit("Not a valid number: $#\n$#" % [args[pos], err], 1)

################################################################
## Execute the user-facing commands
##
proc cmdPing() =
  ## Execute the "ping" command
  let
    chn = getChannel(SonicChannel.Control)
    response = chn.execCommand("PING")
  chn.close()
  quit response, if response == "PONG": 0 else: 1

proc cmdCount(collection, bucket, objId: string) =
  ## Return indexed search data count for collection/bucket/objId
  let
    chn = getChannel(SonicChannel.Search)
    response = chn.count(collection, bucket, objId)
  chn.close()
  quit $response, 0

proc cmdPush(collection, bucket, obj, data: string) =
  ## Ingest a file's content into a Sonic instance
  var
    justOne = false
    stream = stdin

  if data.len == 0:
    quit "Data is an empty string", 1
  elif data == "-":
    stderr.writeLine "Reading from <stdin>"
  elif data[0] == '@':
    let filename = data[1 ..< data.len]
    try:
      stream = open(filename, bufSize=8000)
    except:
      let err = getCurrentExceptionMsg()
      quit "Can't open \"$#\":\n$#" % [filename, err], 10
  else:
    justOne = true

  var
    ingCh = getChannel(SonicChannel.Ingest)
    ctlCh = getChannel(SonicChannel.Control)
    rMsg = ""
    rCode = 0

  if justOne:
    let pushedOk = ingCh.push(collection, bucket, obj, data)
    ctlCh.consolidate()
    rMsg = if pushedOk: "" else: "push command returned a warning"
    rCode = if pushedOk: 0 else: 1
  else:
    let
      objIdPrefix = if data[0] == '@': "$#/$#:" % [obj, data[1 ..< data.len]]
                    else: obj & ":"
    var
      line = newStringOfCap(256)
      count = 0
    stderr.write("push: ")
    while stream.readLine(line):
      let
        objId = objIdPrefix & $count
        pushed = ingCh.push(collection, bucket, objId, line)
      stderr.write(if pushed: "." else: "x")
      inc count
      if (count mod 31) == 0:
        ctlCh.consolidate()
        stderr.write("#")
    stderr.write("\n")
    if data[0] == '@':
      close(stream)
  ctlCh.close()
  ingCh.close()
  quit rMsg, rCode

proc cmdPop(collection, bucket, obj, data: string) =
  ## Pop search data from the given collection/bucket/obj
  if data.len > 0:
    let
      ingCh = getChannel(SonicChannel.Ingest)
      ctlCh = getChannel(SonicChannel.Control)
      popOut = ingCh.pop(collection, bucket, obj, data)
    ctlCh.consolidate()
    ctlCh.close()
    ingCh.close()
    quit $popOut, 0
  quit "Data is an empty string", 1

proc cmdQuery(collection, bucket, terms: string; limit, offset: int) =
  ## Query the indexes, echo the results to stdout
  let
    srChn = getChannel(SonicChannel.Search)
    results = srChn.query(collection, bucket, terms, limit, offset)
  srChn.close()
  quit(results.join("\n"), 0)

proc cmdSuggest(collection, bucket, word: string; limit: int) =
  ## Query suggestions based on the word
  let
    srChn = getChannel(SonicChannel.Search)
    results = srChn.suggest(collection, bucket, word, limit)
  srChn.close()
  quit(results.join("\n"), 0)

proc cmdFlush(collection: string; bucket=""; objId="") =
  ## Flushes all indexed data for the given collection, bucket or object
  let
    cnChn = getChannel(SonicChannel.Control)
    results = cnChn.flush(collection, bucket, objId)
  cnChn.close()
  quit($results, 0)

################################################################
## Parse the command line and dispatch appropriate actions
##
proc main() =
  ## Parse the command line, dispatch the appropriate actions
  var args = commandLineParams()
  
  if args.len == 0 or "-h" in args or "--help" in args:
    showUsage()

  verbose = ("-v" in args) or ("--verbose" in args)
  if verbose:
    while "-v" in args:
      args.delete(args.find("-v"))
    while "--verbose" in args:
      args.delete(args.find("--verbose"))

  case args[0]:
    of "help":
      showUsage()

    of "ping":        # ping
      if args.len > 1:
        showUsage("'ping' takes no arguments")
      cmdPing()

    of "count":     # count <collection> [bucket [object]]
      let
        bucket = (if args.len >= 3: args[2] else: "")
        objId = (if args.len >= 4: args[3] else: "")
      if args.len > 4:
        showUsage("Too many arguments for 'count'")
      cmdCount(args[1], bucket, objId)

    of "push":      # push <collection> <bucket> <object> "data|@filename|-"
      if args.len != 5:
        let pre = if args.len < 5: "Missing" else: "Too many"
        showUsage(pre & " arguments for 'push'")
      cmdPush(args[1], args[2], args[3], args[4])

    of "pop":       # pop <collection> <bucket> <object> "data"
      if args.len != 5:
        let pre = if args.len < 5: "Missing" else: "Too many"
        showUsage(pre & " arguments for 'pop'")
      cmdPop(args[1], args[2], args[3], args[4])

    of "query":     # query <collection> <bucket> "terms" [limit=10] [offset=0]
      if args.len < 4 or args.len > 6:
        let pre = if args.len < 4: "Missing" else: "Too many"
        showUsage(pre & " arguments for 'query'")
      let
        limit = args.intAt(4, 10)
        offset = args.intAt(5, 0)
      cmdQuery(args[1], args[2], args[3], limit, offset)

    of "suggest":     # query <collection> <bucket> "word" [limit=10]
      if args.len < 4 or args.len > 5:
        let pre = if args.len < 4: "Missing" else: "Too many"
        showUsage(pre & " arguments for 'suggest'")
      let
        limit = args.intAt(4, 10)
      cmdSuggest(args[1], args[2], args[3], limit)

    of "flush":       # flush <collection> [bucket [object]]
      if args.len < 3 or args.len > 5:
        let pre = if args.len < 3: "Missing" else: "Too many"
        showUsage(pre & " arguments for 'flush'")
      let
        bucket = (if args.len >= 3: args[3] else: "")
        objId = (if args.len >= 4: args[4] else: "")
      cmdFlush(args[2], bucket, objId)

    else:
      showUsage("Unknown command: " & args[0])

main()
# Fin