# Upgrading the packages (nixpkgs: the kernel and everything else) of an
# installed router without a new release of nixos-router.
#
# The branch to follow is the setting `router.update.channel` (unstable, or a
# release like 26.05). nixpkgs is pinned in the configuration flake's
# flake.lock, so an upgrade is a new lock:
#
#   router-upgrade check      what the channel's newest revision is, against the running one
#   router-upgrade preview    lock a copy of the configuration flake to it and evaluate the
#                             system: which packages change version, how much would be
#                             downloaded and built (nothing is downloaded or activated)
#   router-upgrade download   fetch and build it now (a flaky uplink is no problem later)
#   router-upgrade apply      take the lock and run `router-apply` (commit-confirm)
#
# The revision comes from channels.nixos.org: the branch's last revision that
# passed the NixOS tests. A configuration directory that was upgraded keeps
# the choice in nixpkgs-pin.json, which `router-update` locks again after it
# installs a new release of nixos-router. The directory as it was is kept in
# config.pre-update, as for updates: a rollback or a failure puts it back.
{
  config,
  lib,
  pkgs,
  ...
}:
let
  cfg = config.routerSettings;
  up = cfg.upgrade;
  snapshot = "${cfg.stateDirectory}/config.pre-update";
  dir = "${cfg.stateDirectory}/upgrade";

  upgrade = pkgs.writeShellApplication {
    name = "router-upgrade";
    runtimeInputs = [
      config.nix.package
      config.systemd.package
      pkgs.coreutils
      pkgs.curl
      pkgs.jq
      pkgs.util-linux
    ];
    text = ''
      usage() {
        echo "usage: router-upgrade check | preview [--force] | download | apply [timeout]" >&2
        echo "  check     the newest revision of router.update.channel, against the running one" >&2
        echo "  preview   the packages that change and the size of the download (nothing is fetched)" >&2
        echo "  download  fetch and build what was previewed, without activating it" >&2
        echo "  apply     activate what was previewed (fetching it if need be), with commit-confirm" >&2
        exit 2
      }
      die() { echo "router-upgrade: $*" >&2; exit 1; }
      [ $# -gt 0 ] || usage
      cmd=$1
      shift
      [ "$(id -u)" -eq 0 ] || die "must run as root"

      config=${cfg.directory}
      mkdir -p ${dir}

      channel=$(jq -r '.router.update.channel // empty' "$config/settings.json")
      [ -n "$channel" ] || die "no channel: set router.update.channel (unstable, or a release like 26.05)"
      if [ "$channel" = unstable ]; then branch=nixos-unstable; else branch="nixos-$channel"; fi

      running_rev=$(/run/current-system/sw/bin/nixos-version --revision || true)
      running_version=$(cat /run/current-system/nixos-version)
      short() { echo "''${1:0:7}"; }

      newest_rev() {
        local rev
        rev=$(curl -fsS --max-time 30 "${up.channelsUrl}/$branch/git-revision") \
          || die "cannot read ${up.channelsUrl}/$branch/git-revision (no internet, or no channel of that name)"
        [[ $rev =~ ^[0-9a-f]{40}$ ]] || die "unexpected answer for $branch: ''${rev:0:60}"
        echo "$rev"
      }

      check() {
        newest=$(newest_rev)
        if [ "$newest" = "$running_rev" ]; then same=true; else same=false; fi
        jq -n --arg channel "$channel" --arg branch "$branch" --arg time "$(date -u +%FT%TZ)" \
          --arg running "$running_rev" --arg version "$running_version" --arg newest "$newest" \
          --argjson upToDate "$same" \
          '{channel: $channel, branch: $branch, time: $time, running: {rev: $running, version: $version},
            newest: {rev: $newest}, upToDate: $upToDate}' > ${dir}/check.json.new
        mv ${dir}/check.json.new ${dir}/check.json
        echo "channel $channel ($branch)"
        echo "  running  $running_version"
        echo "  newest   $(short "$newest")"
        if [ "$same" = true ]; then echo "up to date"; else echo "a newer revision is available"; fi
      }

      # The names (with versions) of store paths, one per line.
      names() { sed -e 's|^/nix/store/[a-z0-9]\{32\}-||' -e 's/\.drv$//'; }

      # The packages of the running system none of whose versions is among the
      # new system's derivations': {name, from, to}. Approximate: a package that is rebuilt
      # without a new version is not listed, and build-time versions of a package that
      # the system does not run are not either.
      changes() { # file of the running system's paths, file of the new system's
        jq -n --rawfile old "$1" --rawfile new "$2" '
          def lines: split("\n") | map(select(. != "" and (test("\\.(tar|patch|diff|zip|tgz|tbz2?|txz|gz|xz|bz2|sh|py|nix|json|conf|asc|sig|lock)$") | not)));
          def strip: sub("-(bin|dev|lib|out|man|doc|info|devdoc|static|lib32|modules|modules-shrunk|libgcc|getent|debug)$"; "");
          def parse: capture("^(?<n>.*?)-(?<v>[0-9].*)$") // {n: ., v: ""};
          def table: map(strip | parse) | group_by(.n) | map({key: .[0].n, value: (map(.v) | unique)}) | from_entries;
          ($old | lines | table) as $o | ($new | lines | table) as $n
          | [ $o | to_entries[] | select($n[.key] != null and .value != [""])
              | select(((.value - $n[.key]) | length) == (.value | length) and (($n[.key] - .value) | length) > 0)
              | {name: .key, from: (.value | join(", ")), to: ($n[.key] - .value | join(", "))} ]
          | sort_by(.name)'
      }

      preview() {
        local force=0
        [ "''${1:-}" != --force ] || force=1
        check
        rm -rf ${dir}/preview.json ${dir}/flake ${dir}/result
        if [ "$same" = true ] && [ "$force" -eq 0 ]; then return 0; fi

        mkdir ${dir}/flake
        for f in flake.nix settings.json local.nix flake.lock; do
          if [ -e "$config/$f" ]; then cp "$config/$f" ${dir}/flake/; fi
        done
        ref="${up.nixpkgsRef}"
        ref=''${ref//%rev%/$newest}
        echo "router-upgrade: locking nixpkgs to $(short "$newest")"
        nix flake lock --override-input nixos-router/nixpkgs "$ref" "path:${dir}/flake"
        node=$(jq -r '.nodes[.nodes[.root].inputs["nixos-router"]].inputs.nixpkgs' ${dir}/flake/flake.lock)
        hash=$(jq -er --arg n "$node" '.nodes[$n].locked.narHash' ${dir}/flake/flake.lock)
        case $ref in
          *narHash=*) pin=$ref ;;
          *\?*) pin="$ref&narHash=$hash" ;;
          *) pin="$ref?narHash=$hash" ;;
        esac

        attr="path:${dir}/flake#nixosConfigurations.router.config.system"
        echo "router-upgrade: evaluating the new system"
        new_version=$(nix eval --raw "$attr.nixos.version")
        # The release (26.05) of a version like 26.05.20260101.abcdef0 or 26.11pre-git.
        release() { sed -E 's/^([0-9]+\.[0-9]+).*/\1/' <<<"$1"; }
        newer=$(printf '%s\n%s\n' "$(release "$new_version")" "$(release "$running_version")" | sort -V | tail -1)
        [ "$newer" = "$(release "$new_version")" ] \
          || die "$new_version is older than the running $running_version: NixOS cannot be downgraded"
        drv=$(nix eval --raw "$attr.build.toplevel.drvPath")

        nix path-info -r /run/current-system | names > ${dir}/old.names
        nix path-info -r "$drv" | names > ${dir}/new.names
        changes ${dir}/old.names ${dir}/new.names > ${dir}/changes.json
        rm -f ${dir}/old.names ${dir}/new.names
        if [ "$(jq length ${dir}/changes.json)" -gt 0 ]; then
          jq -r '.[] | "  \(.name): \(.from) -> \(.to)"' ${dir}/changes.json
        else
          echo "  no package changes"
        fi
        if jq -e 'any(.[]; .name == "linux")' ${dir}/changes.json > /dev/null; then reboot=true; else reboot=false; fi

        # Nothing is fetched: the substituters only say what they have.
        dry=$(nix build --dry-run --no-link "$attr.build.toplevel" 2>&1 || true)
        download=$(grep -o '([0-9.]* [KMGT]iB download, [0-9.]* [KMGT]iB unpacked)' <<<"$dry" | head -1 || true)
        builds=$(grep -cE '^ +/nix/store/[^ ]*\.drv$' <<<"$dry" || true)
        echo "  to fetch: ''${download:-nothing, or no cache reachable}; to build here: $builds"

        jq -n --arg channel "$channel" --arg rev "$newest" --arg ref "$pin" --arg version "$new_version" \
          --arg time "$(date -u +%FT%TZ)" --arg download "$download" --argjson builds "$builds" \
          --argjson reboot "$reboot" --slurpfile changes ${dir}/changes.json \
          '{channel: $channel, rev: $rev, ref: $ref, version: $version, time: $time, download: $download,
            builds: $builds, rebootNeeded: $reboot, changes: $changes[0]}' > ${dir}/preview.json
        echo "$new_version$([ "$reboot" = false ] || echo ' (a new kernel runs after a reboot)'): router-upgrade apply activates it"
      }

      download() {
        [ -r ${dir}/preview.json ] || die "nothing is previewed (router-upgrade preview)"
        free=$(df --output=avail -B1M /nix/store | tail -1)
        [ "$free" -ge 3072 ] || die "less than 3 GiB free on /nix/store (''${free} MiB): free some space first"
        echo "router-upgrade: fetching and building (nothing is activated)"
        nix build --no-link --out-link ${dir}/result \
          "path:${dir}/flake#nixosConfigurations.router.config.system.build.toplevel"
        echo "downloaded $(jq -r .version ${dir}/preview.json); router-upgrade apply activates it"
      }

      apply() {
        local timeout=''${1:-}
        [ -r ${dir}/preview.json ] || die "nothing is previewed (router-upgrade preview)"
        exec 9>/run/router-update.lock
        flock -n 9 || die "an update or upgrade is running"
        [ ! -d ${snapshot} ] || die "an earlier update or upgrade is not confirmed yet"
        cmp -s "$config/flake.nix" ${dir}/flake/flake.nix \
          || die "the configuration flake changed since it was previewed: preview again"
        cp -a "$config" ${snapshot}
        fail() {
          echo "router-upgrade: $*" >&2
          /run/current-system/sw/bin/router-update-restore
          exit 1
        }
        cp ${dir}/flake/flake.lock "$config/flake.lock"
        jq '{channel, rev, ref}' ${dir}/preview.json > "$config/nixpkgs-pin.json"

        echo "router-upgrade: building and activating $(jq -r .version ${dir}/preview.json)"
        if ! router-apply ''${timeout:+"$timeout"}; then
          if systemctl is-active --quiet router-rollback.timer; then
            # Activation failed half way: the rollback restores the configuration.
            systemctl start router-rollback.service
            die "the upgrade failed; the previous system is back"
          fi
          fail "the upgrade failed; the previous system is back"
        fi
        rm -rf ${dir}/preview.json ${dir}/flake ${dir}/result
        if systemctl is-active --quiet router-rollback.timer; then
          echo "router-upgrade: active; confirm it (router-confirm, or in the web UI) or the previous system comes back"
        else
          rm -rf ${snapshot} # nothing changed: nothing to confirm
        fi
      }

      case $cmd in
        check) check ;;
        preview) preview "$@" ;;
        download) download ;;
        apply) apply "$@" ;;
        *) usage ;;
      esac
    '';
  };
in
{
  options.routerSettings.upgrade = {
    channelsUrl = lib.mkOption {
      type = lib.types.str;
      default = "https://channels.nixos.org";
      description = "Where the newest revision of a nixpkgs branch is read (`<branch>/git-revision`).";
    };
    nixpkgsRef = lib.mkOption {
      type = lib.types.str;
      default = "github:NixOS/nixpkgs/%rev%";
      description = "The flake reference of a nixpkgs revision (`%rev%` stands for it).";
    };
  };

  config.environment.systemPackages = [ upgrade ];
}