import std/[asyncdispatch, deques, json, os, osproc, streams, strutils, tables]
import cbor, dhall/[render, terms], eris

const selfDescribedCbor = 55799

proc writeErisLinks(path: string; node: CborNode) =
  # Inspired by "Package Metadata for Core Files"
  # - https://systemd.io/COREDUMP_PACKAGE_METADATA/
  # - http://netbsd.org/docs/kernel/elf-notes.html
  const name = "Sigil\x00"
  var links = newStringStream()
  links.writeCborTag(selfDescribedCbor)
  links.writeCbor(node)
  var file = openFileStream(path, fmWrite)
  file.write(uint32 name.len)
  file.write(uint32 links.data.len)
  file.write(uint32 1)
  file.write(name)
  while (file.getPosition and 3) != 0: file.write(byte 0)
  file.write(links.data)
  while (file.getPosition and 3) != 0: file.write(byte 0)
  close(file)

proc toCbor(cap: ErisCap): CborNode =
  result = initCborBytes(cap.bytes)
  result.tag = erisCborTag

proc toDhall(js: JsonNode): Value =
  case js.kind
  of JString:
    result = newValue(js.str)
  of JObject:
    result = newRecordLiteral(js.fields.len)
    for key, val in js.fields.pairs:
      result.table[key] = val.toDhall
  else:
    raiseAssert "unhandled JSON value"

proc isElf(path: string): bool =
  var magic: array[4, char]
  let file = open(path)
  discard readChars(file, magic)
  close(file)
  magic == [0x7f.char, 'E', 'L', 'F']

proc main =
  var generateManifests, generateNotes, addNotes, runPatchelf: bool
  for arg in commandLineParams():
    case arg
    of "--add-notes": addNotes = true
    of "--run-patchelf": runPatchelf = true
    of "--generate-manifests": generateManifests = true
    of "--generate-notes": generateNotes = true

  if getEnv("dontErisPatch") != "": quit 0
  let
    patchelf = getEnv("ERIS_PATCHELF", "patchelf")
    objcopy = getEnv("ERIS_OBJCOPY", "objcopy")
    nixStore = getEnv("NIX_STORE", "/nix/store")
    jsonManifestSubPath = "nix-support" / "eris-manifest.json"
    dhallManifestSubPath = "nix-support" / "eris-manifest.dhall"

  type PendingFile = ref object
    outputRoot, filePath: string
    replacements: Table[string, string]
    erisLinks: CborNode

  var
    outputManifests = initTable[string, JsonNode]()
    pendingFiles = initDeque[PendingFile]()
    failed = false
  if getEnv("outputs") == "":
    quit """package $outputs not set in environment"""
  for outputName in getEnv("outputs").splitWhitespace:
    if getEnv(outputName) == "":
      quit("package $" & outputName & " not set in environment")
    let outputRoot = getEnv(outputName)
    if fileExists(outputRoot / jsonManifestSubPath):
      stderr.writeLine "Not running ERIS patch hook again"
      quit 0
    outputManifests[outputRoot] = newJObject()

  let buildInputs = getEnv("buildInputs").splitWhitespace

  proc resolveNeed(rpath: seq[string]; need: string): string =
    if need.isAbsolute:
      return need
    for libdir in rpath:
      let absNeed = libdir / need
      if fileExists(absNeed):
        return absNeed
    for outputRoot in outputManifests.keys:
      for relPath in [need, "lib" / need]:
        let absNeed = outputRoot / relPath
        if fileExists(absNeed):
          return absNeed
    for buildInput in buildInputs:
      for relPath in [need, "lib" / need]:
        let absNeed = buildInput / relPath
        if fileExists(absNeed):
          return absNeed

  proc resolveFile(outputRoot, filePath: string): PendingFile =
    result = PendingFile(
        outputRoot: outputRoot,
        filePath: filePath,
        replacements: initTable[string, string](8),
        erisLinks: initCborMap())
    let needs = splitWhitespace(execProcess(
        patchelf, args = ["--print-needed", filePath], options = {poUsePath}))
    let rpath = splitWhitespace(execProcess(
        patchelf, args = ["--print-rpath", filePath], options = {poUsePath}))
    for need in needs:
      if need == "ld.lib.so" or need.startsWith("urn:"): continue
      result.replacements[need] = resolveNeed(rpath, need)

  var capCache = initTable[string, ErisCap]()

  proc fileCap(filePath: string): ErisCap =
    ## Determine the ERIS read capabililty for ``filePath``.
    if capCache.hasKey(filePath):
      result = capCache[filePath]
    else:
      try:
        let str = openFileStream(filePath)
        result = waitFor encode(newDiscardStore(), str, convergentMode)
        capCache["filePath"] = result
        close(str)
      except:
        stderr.writeLine("failed to read \"", filePath, "\"")
        quit 1

  var closureCache = initTable[string, TableRef[string, ErisCap]]()

  proc fileClosure(filePath: string): TableRef[string, ErisCap] =
    ## Recusively find the dependency closure of  ``filePath``.
    let filePath = expandFilename filePath
    if closureCache.hasKey(filePath):
      result = closureCache[filePath]
    else:
      result = newTable[string, ErisCap]()
      var storePath = filePath
      for p in parentDirs(filePath):
        # find the top directory of the ``filePath`` derivation
        if p == nixStore: break
        storePath = p
      if storePath.startsWith nixStore:
        # read the closure manifest of the dependency
        let manifestPath = storePath / jsonManifestSubPath
        if fileExists(manifestPath):
          let
            manifest = parseFile(manifestPath)
            entry = manifest[filePath.extractFilename]
          for path, urn in entry["closure"].pairs:
            result[path] = parseErisUrn urn.getStr
            let otherClosure = fileClosure(path)
            for otherPath, otherCap in otherClosure.pairs:
              # merge the closure of the dependency
              result[otherPath] = otherCap
      closureCache[filePath] = result

  for outputRoot in outputManifests.keys:
    let manifestPath = outputRoot / jsonManifestSubPath
    if fileExists manifestPath: continue
    for filePath in walkDirRec(outputRoot, relative = false):
      # Populate the queue of files to patch
      if filePath.isElf:
        pendingFiles.addLast(resolveFile(outputRoot, filePath))

  var
    prevLen = pendingFiles.len
    prevPrevLen = prevLen.succ
      # used to detect reference cycles
  while pendingFiles.len != 0:
    block selfReferenceCheck:
      # process the files that have been collected
      # taking care not to take a the URN of an
      # unprocessed file
      let
        pendingFile = pendingFiles.popFirst()
        filePath = pendingFile.filePath
      for need, replacementPath in pendingFile.replacements.pairs:
        # search for self-references
        if replacementPath == "":
          stderr.writeLine need, " not found for ", filePath
          failed = true
          continue
        for outputRoot in outputManifests.keys:
          if replacementPath.startsWith(outputRoot):
            for other in pendingFiles.items:
              stderr.writeLine "compare for self-reference:"
              stderr.writeLine '\t', replacementPath
              stderr.writeLine '\t', other.filePath
              if replacementPath == other.filePath:
                stderr.writeLine "defering patch of ", filePath, " with reference to ", other.filePath
                pendingFiles.addLast(pendingFile)
                break selfReferenceCheck
      var
        closure = newJObject()
        replaceCmd = patchelf & " --set-rpath '' " & filePath
      for need, replacementPath in pendingFile.replacements.pairs:
        if replacementPath == "": continue
        let
          cap = fileCap(replacementPath)
          urn = $cap
        stderr.writeLine "replace reference to ", need, " with ", urn
        closure[replacementPath] = %urn
        pendingFile.erisLinks[ replacementPath.toCbor] = cap.toCbor
        replaceCmd.add(" --replace-needed $# $#" % [need, urn])
        for path, cap in fileClosure(replacementPath).pairs:
          closure[path] = %($cap)
          pendingFile.erisLinks[path.toCbor] = cap.toCbor

      if runPatchelf and pendingFile.replacements.len != 0:
        let exitCode = execCmd(replaceCmd)
        if exitCode != 0:
          stderr.writeLine "Patchelf failed - ", replaceCmd
          quit exitCode

      let notePath = getEnv("TMPDIR", ".") / "eris-links.note"
      if generateNotes or addNotes:
        if generateNotes:
          stderr.writeLine("writing .notes.eris-links section to ", notePath)
        sort(pendingFile.erisLinks)
        writeErisLinks(notePath, pendingFile.erisLinks)
      if addNotes:
        # Modify the file with objcopy last because binutils surely
        # has a more conventional interpretation of the ELF standard
        # than patchelf.
        let tmpFile = pendingFile.filePath & ".with-note"
        let objcopyCommand = [
          objcopy,
          "--add-section .note.eris-links=" & notePath,
          "--set-section-flags .note.eris-links=noload,readonly",
          pendingFile.filePath, tmpfile]
        let exitCode = execCmd(objcopyCommand.join(" "))
        if exitCode != 0:
          quit("Adding note to $1 failed" % pendingFile.filePath)
        moveFile(tmpFile, pendingFile.filePath)

      outputManifests[pendingFile.outputRoot][filePath.extractFilename] = %* {
        "cap": $fileCap(filePath),
        "closure": closure,
        "path": filePath
      }

    if pendingFiles.len == prevPrevLen:
      failed = true
      stderr.writeLine "reference cycle detected in the following:"
      for remain in pendingFiles.items:
        stderr.writeLine '\t', " ", remain.filePath
      break
    prevPrevLen = prevLen
    prevLen = pendingFiles.len

  if failed:
    quit -1

  if generateManifests:
    for outputRoot, manifest in outputManifests:
      let supportDir = outputRoot / "nix-support"
      createDir(supportDir)
      writeFile(outputRoot / jsonManifestSubPath, $manifest)
      writeFile(outputRoot / dhallManifestSubPath, $(manifest.toDhall))

main()