swap-sentinel.nix

{ config, lib, pkgs, ... }:

let
  cfg = config.services.btrfs-flex-swap;
in {
  options.services.btrfs-flex-swap = {
    enable = lib.mkEnableOption "Dynamic Btrfs flex-swap manager";
    
    minFreeMB = lib.mkOption {
      type = lib.types.int;
      default = 512;
      description = "Threshold of free swap in MB before creating a new file.";
    };
    
    chunkSizeGB = lib.mkOption {
      type = lib.types.int;
      default = 1;
      description = "Size of each dynamic swap chunk in GB.";
    };
  };

  config = lib.mkIf cfg.enable {
    systemd.services.btrfs-flex-swap = {
      description = "Dynamic Btrfs Flex Swap Daemon";
      after = [ "local-fs.target" ];
      wantedBy = [ "multi-user.target" ];

      serviceConfig = {
        Type = "simple";
        Restart = "always";
        RestartSec = "10s";
        # Systemd natively creates and manages /var/lib/btrfs-flex-swap
        StateDirectory = "btrfs-flex-swap"; 
        
        # We still need to ensure +C is set on the directory before files are made
        ExecStartPre = let
          preScript = pkgs.writeShellScript "btrfs-flex-swap-pre" ''
            ${pkgs.e2fsprogs}/bin/chattr +C /var/lib/btrfs-flex-swap || true
          '';
        in "${preScript}";

        ExecStart = let
          script = pkgs.writeShellScript "btrfs-flex-swap-run" ''
            set -euo pipefail
            
            STORAGE_PATH="/var/lib/btrfs-flex-swap"
            CHUNK_SIZE_MB=$(( ${toString cfg.chunkSizeGB} * 1024 ))
            THRESHOLD=${toString cfg.minFreeMB}
            
            echo "Flex-swap monitor started on $STORAGE_PATH..."

            while true; do
              # Get current free swap in MB securely
              FREE_SWAP=$(${pkgs.procps}/bin/free -m | ${pkgs.gawk}/bin/awk '/^Swap:/ {print $4}')
              
              if [ "$FREE_SWAP" -lt "$THRESHOLD" ]; then
                TIMESTAMP=$(${pkgs.coreutils}/bin/date +%s)
                SWAPFILE="$STORAGE_PATH/extra-$TIMESTAMP.swap"
                
                echo "Low swap ($FREE_SWAP MB). Allocating $SWAPFILE..."
                
                # Use the modern, atomic Btrfs-specific swapfile creation tool
                ${pkgs.btrfs-progs}/bin/btrfs filesystem mkswapfile --size ${toString cfg.chunkSizeGB}G "$SWAPFILE"
                ${pkgs.util-linux}/bin/swapon "$SWAPFILE"
              
              elif [ "$FREE_SWAP" -gt "$(( CHUNK_SIZE_MB + THRESHOLD + 512 ))" ]; then
                # Safely find the oldest dynamic swap file
                OLDEST_DYNAMIC=$(${pkgs.findutils}/bin/find "$STORAGE_PATH" -name 'extra-*.swap' -type f -printf '%T+ %p\n' 2>/dev/null | ${pkgs.coreutils}/bin/sort | ${pkgs.coreutils}/bin/head -n 1 | ${pkgs.coreutils}/bin/cut -d' ' -f2 || true)
                
                if [ -n "$OLDEST_DYNAMIC" ]; then
                  echo "High swap headroom ($FREE_SWAP MB). Offloading $OLDEST_DYNAMIC..."
                  
                  # Temporarily disable exit-on-error. If swap is too busy, swapoff will fail.
                  # We want to catch the failure and retry next loop rather than crash the service.
                  set +e
                  ${pkgs.util-linux}/bin/swapoff "$OLDEST_DYNAMIC"
                  SWAPOFF_STATUS=$?
                  set -e
                  
                  if [ $SWAPOFF_STATUS -eq 0 ]; then
                    ${pkgs.coreutils}/bin/rm "$OLDEST_DYNAMIC"
                    echo "Successfully removed $OLDEST_DYNAMIC."
                  else
                    echo "Notice: Failed to swapoff $OLDEST_DYNAMIC (device busy). Will retry next cycle."
                  fi
                fi
              fi
              
              ${pkgs.coreutils}/bin/sleep 30
            done
          '';
        in "${script}";
      };
    };
  };
}