#!/usr/bin/env nix-shell
#!nix-shell -i ysh -p nettools -p oils-for-unix

# TODO: simplify variant handling by defining `desko-full = desko`, etc, in nix config.

const SELF = $(hostname)
# which hosts to deploy to (in order), when "all hosts" is specified or implied.
const ALL_HOSTS = [ "moby", "flowy", "desko", "servo" ]
# which variants to build (in order), when "all variants" is specified or implied.
const ALL_VARIANTS = ["-min", "-light", "", "-next-min", "-next-light", "-next"]


# mutated by `main`
var DRY_RUN = false
var NIX_ARGS = []  # extra args to pass to nix invocations

proc usage (;; status=1) {
  echo "deploy: deploy a nix config to a remote machine, possibly activating it"
  echo ""
  echo "usage: deploy [options] [host] [host2 ...]"
  echo "options:"
  echo "- --build: only build; don't copy or deploy"
  echo "- --copy: only build + copy files, nothing more"
  echo "- --switch  (default)"
  echo "- --test: switch to the new configuration, but do not make it bootable"
  echo "- --dry-run: show what would be done without actually doing it"
  echo "- --pre: alias for --action copy ..."
  echo '         if otherwise unspecified, implies `--variant all`'
  echo "         if no host is specified, copies to all known hosts"
  echo "- --reboot: reboot the target machine after deploying (if deployed with no errors)"
  echo "- --reboot-force: reboot the target machine after deploying (even on error)"
  echo "- --variant light|min|''|all  (default: '')"
  echo "- --wireguard always|never|opportunistic: deploy over wireguard"
  echo "- --ip <address>: deploy to the specific IP address"
  echo "- --deriv /nix/store/...: prebuilt store path (or .drv to realize) to deploy instead of (re-)building the default target"
  echo ""
  echo "common idioms:"
  echo "- deploy all: deploy all hosts, sequentially"
  echo "- deploy --pre: build and copy all hosts"
  echo "- deploy desko lappy: build and deploy just those hosts"
  echo "- deploy: deploy the local host"
  exit "$status"
}

proc info (...plain; ...escape) {
  for e in (escape) {
    call plain->append(toJson(e))
  }
  echo "[deploy]" @plain >&2
}

var CliArgs = Object(null, {
  action: null,  # str, e.g. build, copy, switch, test
  deriv: null,  # str, e.g. /nix/store/...
  dryRun: false,  # bool
  hosts: [],  # list of str, e.g. [ "desko", "moby" ]
  help: false,
  ip: null,
  pre: false,
  reboot: false,
  rebootForce: false,
  variants: [],  # list of str, e.g. [ "-min" ] or [ "-light", "" ]
  wireguard: "opportunistic",  # str: e.g. "opportunistic", "never", "always"
  unhandled: [],  # unhandled arguments; forward them to nix?
})

# each NixConfig instance represents _one_ nix build target + host-specific deployment strategy,
# e.g. "build lappy-min; then copy it to lappy (over wireguard) but don't activate it"
var NixConfig = Object(null, {
  host: null,     # str, e.g. "desko"
  variant: "",    # str, e.g. "-min", "-light"
  drvPath: null,  # str, e.g. "/nix/store/...-nixos-system.drv"
  storePath: null,  # str, e.g. "/nix/store/...-nixos-system"
  action: "switch",  # str, e.g. "switch", "copy", .... TODO: convert to enum?
  ip: null,  # str, e.g. "10.78.79.51"
  wireguard: "opportunistic",  # str, e.g. "opportunistic", "never", "always".  TODO: convert to enum?
  reboot: false,  # may be false, true, or "force"
})

# pass through normal hosts, but expand aliases like `"all"` into `[ "moby", "lappy", ... ]`.
# @arg(host): str
# @return: list[str]
func expandHostAlias (host) {
  if (host === "all" ) {
    return (ALL_HOSTS)
  } else {
    return ([host])
  }
}
# pass through normal variants, but expand aliases like `"all"` into `[ "-min", "-light", ... ]`.
# @arg(v): str
# @return: list[str]
func expandVariantAlias (v) {
  if (v === "all") {
    return (ALL_VARIANTS)
  } elif (v === "") {
    # "full" variant
    return ([""])
  } else {
    return (["-$v"])
  }
}

# internal function for use by `parseArgs`.
# transforms argv into `CliArgs` object. should be side-effect free.
# @arg(cur, next, ...): each one is str; unpacked only for internal use
# @arg(p): intermediate CliArgs object; meant for internal use
# @return: CliArgs instance
func _parseArgs(p,cur=null, next=null, ...rest) {
  if (cur === null) {
    # end
    return (p)
  } else {
    case (cur) {
      --build | --copy | --switch | --test {
        setvar p.action = cur[2:]
        return (_parseArgs (p, next, ...rest))
      }
      --deriv {
        setvar p.deriv = next
        return (_parseArgs (p, ...rest))
      }
      --dry-run {
        setvar p.dryRun = true
        return (_parseArgs (p, next, ...rest))
      }
      --help {
        setvar p.help = true
        return (_parseArgs (p, next, ...rest))
      }
      --ip {
        setvar p.ip = next
        return (_parseArgs (p, ...rest))
      }
      --pre {
        setvar p.pre = true
        return (_parseArgs (p, next, ...rest))
      }
      --reboot {
        setvar p.reboot = true
        return (_parseArgs (p, next, ...rest))
      }
      --reboot-force | --force-reboot {
        setvar p.rebootForce = true
        return (_parseArgs (p, next, ...rest))
      }
      --variant {
        call p.variants->append(next)
        return (_parseArgs (p, ...rest))
      }
      --wireguard {
        setvar p.wireguard = next
        return (_parseArgs (p, ...rest))
      }
      all | crappy | desko | flowy | lappy | moby | servo {
        call p.hosts->append(cur)
        return (_parseArgs (p, next, ...rest))
      }
      (else) {
        call p.unhandled->append(cur)
        return (_parseArgs (p, next, ...rest))
      }
    }
  }
}

# front-facing argument parser.
# transforms argv into `CliArgs` object. should be side-effect free.
# @arg(...): str
# @return: CliArgs instance
func parseArgs(...args) {
  var p=Object(CliArgs, {})
  return (_parseArgs (p, ...args))
}

proc destructive (...cmd; ; echo=false) {
  if (DRY_RUN) {
    info "dry-run:" (...cmd)
  } else {
    if (echo) {
      echo @cmd >&2
    }
    @cmd
  }
}

# return "$ip", "$host" or "$host-hn", based on if ip was specified,
# or if wireguard was requested.
func resolveHost (host, ip, wireguard) {
  if (host === SELF) {
    return (host)
  } elif (ip) {
    return (ip)
  } else {
    case ("$wireguard-$host") {
      opportunistic-moby {
        return ("$host-hn")
      }
      opportunistic-* {
        return (host)
      }
      never-* {
        return (host)
      }
      always-* {
        return ("$host-hn")
      }
      (else) {
        return ("$host-hn")
      }
    }
  }
}

# return the number of seconds to allot to `nix copy` when copying the given variant.
# this is done to avoid stalled copies from blocking the entire process, while hopefully not breaking the copies that are actually important
func timeoutFor (variant) {
  case (variant) {
    -min | -light | -next {
      return (3600)
    }
    -next-min | -next-light {
      return (1800)
    }
    (else) {
      # this catches the normal variant
      return (14400)
    }
  }
}

proc runOnTarget (host, ...cmd) {
  # run the command on the machine we're deploying to.
  # if that's a remote machine, then do it via ssh, else local shell.
  if (host !== SELF) {
    info "running on remote ($host):" @cmd
    ssh "$host" @cmd
  } else {
    info "running locally ($SELF):" @cmd
    @cmd
  }
}

proc deployOneHost (; nixcfg) {
  # unpack the cfg
  var timeout = timeoutFor(nixcfg.variant)
  var host = nixcfg.host
  var attrName = "$[nixcfg.host]$[nixcfg.variant]"

  # first, realize the host config derivation if not exists:
  var myStorePath = nixcfg.storePath
  if (not myStorePath) {
    var drvPath = nixcfg.drvPath
    if (not drvPath) {
      # `nix-build -A foo` evals and then realizes foo, but it never unloads the memory used to eval foo.
      # my exprs are heavyweight, we need that memory for building, so do the evals separately from the realizations:
      info "$attrName: evaluating..."
      setvar drvPath = $(nix eval --raw -f . "hosts.$attrName.toplevel.drvPath")
    }

    info "$attrName: building $drvPath"
    setvar myStorePath = $(destructive nix-store --realize "$drvPath" --add-root "./build/$attrName" @NIX_ARGS)
    # if using `nix-store --add-root`, then nix prints the gc root path, instead of the /nix/store path; resolve it
    setvar myStorePath = $(realpath "$myStorePath")
    info "$attrName: built $myStorePath"
  }

  # mimic `nixos-rebuild --target-host`, in effect:
  # - nix-copy-closure ...
  # - nix-env --set ...
  # - switch-to-configuration <boot|dry-activate|switch|test|>
  # avoid the actual `nixos-rebuild` for a few reasons:
  # - fewer nix evals
  # - more introspectability and debuggability
  # - sandbox friendliness (especially: `git` doesn't have to be run as root)

  var netHost = resolveHost(host, nixcfg.ip, nixcfg.wireguard)

  case (nixcfg.action) {
    copy | switch | test {
      if (host !== SELF) {
        if test -e /run/secrets/nix_signing_key {
          info "signing store paths ..."
          destructive sudo nix store sign -r -k /run/secrets/nix_signing_key "$myStorePath"
        } else {
          info "not signing store paths: /run/secrets/nix_signing_key does not exist"
        }
        # add more `-v` for more verbosity (up to 5).
        # builders-use-substitutes false: optimizes so that the remote machine doesn't try to get paths from its substituters.
        #   we already have all paths here, and the remote substitution is slow to check and SERIOUSLY flaky on moby in particular.
        destructive timeout "$timeout" nix copy -vv --option builders-use-substitutes false --to "ssh-ng://$netHost" "$myStorePath" (echo=true)
      }
    }
  }

  case (nixcfg.action) {
    switch | test {
      info "activating profile... "
      destructive runOnTarget "$netHost" sudo nix-env -p /nix/var/nix/profiles/system --set "$myStorePath"

      # at this point, the new profile will be booted into on next boot.
      # try to switch to that profile _now_, but allow this to be fallible
      # so as to support `--force-reboot` CLI option:
      try {
        destructive runOnTarget "$netHost" sudo "$myStorePath/bin/switch-to-configuration" "$[nixcfg.action]"
      }
      var fail = false
      if failed { setvar fail = true }  # XXX: `failed` special variable is only readable via `if failed`

      if (nixcfg.reboot and (not fail or nixcfg.reboot === "force")) {
        info "$attrName: rebooting $host"
        destructive runOnTarget "$netHost" sane-reboot "$host"
      }

      if (fail) {
        false
      }
    }
  }
}

proc main (...args) {
  var cli = parseArgs (...args)

  if (cli.help) {
    usage (status=0)
  }

  setglobal NIX_ARGS = cli.unhandled;
  if (cli.dryRun) {
    setglobal DRY_RUN = true
  }

  var defaultAction = "switch"
  var defaultHosts = [ SELF ]
  var defaultVariants = [ "" ]
  if (cli.pre) {
    setvar defaultAction = "copy"
    setvar defaultHosts = [ "all" ]
    setvar defaultVariants = [ "all" ]
  }

  # expand hosts and handle default
  var hosts = []
  for h in (cli.hosts or defaultHosts) {
    call hosts->extend(expandHostAlias(h))
  }
  # expand variants and handle default
  var variants = []
  for v in (cli.variants or defaultVariants) {
    call variants->extend(expandVariantAlias(v))
  }

  # compute a plan
  var cfgs = []
  for variant in (variants) {
    for host in (hosts) {
      var cfg = Object(NixConfig, {
        host: host,
        variant: variant,
        # drvPath: (set below),
        # storePath: (set below),
        action: cli.action or defaultAction,
        ip: cli.ip,
        wireguard: cli.wireguard,
        # reboot: (set below),
      })
      if (cli.reboot) {
        setvar cfg.reboot = true
      }
      if (cli.rebootForce) {
        setvar cfg.reboot = "force"
      }
      if (cli.deriv) {
        if (cli.deriv => endsWith(".drv")) {
          setvar cfg.drvPath = cli.deriv
        } else {
          setvar cfg.storePath = cli.deriv
        }
      }
      call cfgs->append(cfg)
    }
  }

  info "----- deployment plan -----"
  if (DRY_RUN) {
    info "DRY RUN"
  }
  if (NIX_ARGS) {
    info "NIX ARGS:" (...NIX_ARGS)
  }
  for cfg in (cfgs) {
    info "- $[cfg.host]$[cfg.variant]: $[cfg.action], reboot=$[cfg.reboot], wg=$[cfg.wireguard]"
  }
  info "---------------------------"

  var failedDeploys = []
  for cfg in (cfgs) {
    try {
      deployOneHost (cfg)
    }
    if failed {
      info "❌$[cfg.host]$[cfg.variant]: failed deployment!"
      call failedDeploys->append(cfg)
    } else {
      info "✅$[cfg.host]$[cfg.variant]: deployment complete"
    }
  }

  if (failedDeploys !== []) {
    echo "FAILED DEPLOYMENT:"
    for d in (failedDeploys) {
      echo "- $[d.host]$[d.variant]"
    }
    exit 1
  }

  echo "SUCCESS"
}

if is-main {
  main @ARGV
}