home/hosts/odin/system/impermanence.nix

{ config
, inputs
, lib
, ...
}:
{
  imports = [
    inputs.impermanence.nixosModules.impermanence
  ];

  boot.initrd.postDeviceCommands = lib.mkAfter ''
    #!/bin/sh
    DEVICE=${config.disko.devices.disk.main.device}-part2

    # Mount Btrfs root
    mkdir -p /mnt
    if ! mount -o subvol=/ $DEVICE /mnt; then
      echo "Failed to mount $DEVICE at /mnt"
      exit 1
    fi

    # Create directory for old roots
    mkdir -p /mnt/old-roots

    # Move current root to old-roots with current timestamp
    if [[ -e /mnt/@root ]]; then
      timestamp=$(date +%Y-%m-%d_%H:%M:%S)
      if ! btrfs subvolume snapshot -r /mnt/@root "/mnt/old-roots/@root-$timestamp"; then
        echo "Failed to move /mnt/@root to old-roots"
        umount /mnt
        exit 1
      fi
    fi

    # Function to recursively delete subvolumes
    delete_subvolume_recursively() {
      for i in $(btrfs subvolume list -o "$1" | cut -f 9- -d ' '); do
        delete_subvolume_recursively "/mnt/$i"
      done
      if ! btrfs subvolume delete "$1"; then
        echo "Failed to delete subvolume $1"
      fi
    }

    # Delete old roots more than 5
    index=0
    for i in $(btrfs subvolume list /mnt | grep 'old-roots/@root-' | cut -f 9- -d ' ' | sort -r); do
      if [[ $index -ge 5 ]]; then
        delete_subvolume_recursively "/mnt/$i"
      fi
      index=$((index + 1))
    done

    # Create or restore fresh root
    if [[ -e /mnt/@root-blank ]]; then
      delete_subvolume_recursively /mnt/@root
      if ! btrfs subvolume snapshot /mnt/@root-blank /mnt/@root; then
        echo "Failed to snapshot @root-blank to @root"
        umount /mnt
        exit 1
      fi
    else
      if ! btrfs subvolume create /mnt/@root-blank; then
        echo "Failed to create @root-blank"
        umount /mnt
        exit 1
      fi
      if ! btrfs subvolume create /mnt/@root; then
        echo "Failed to create @root"
        umount /mnt
        exit 1
      fi
    fi

    # Unmount
    if ! umount /mnt; then
      echo "Failed to unmount /mnt"
      exit 1
    fi
  '';

  # Persistent directories for impermanence
  fileSystems."/persist".neededForBoot = true;
  fileSystems."/var/lib".neededForBoot = true;

  environment.persistence."/persist" = {
    hideMounts = true;

    directories = [
      "/etc/ssh"
      "/var/cache"
    ];

    files = [
      "/etc/machine-id"
    ];

    users.thomas.directories = [
      ".ssh"
      ".local/share"
    ];
  };
}