server-config/lib.nix

{
  inputs,
  pkgs,
  ...
}:
inputs.nixpkgs.lib.extend (
  final: prev:
  let
    lib = prev;
  in
  {
    # stolen from: https://stackoverflow.com/a/42398526
    optimizeWithFlags =
      pkg: flags:
      lib.overrideDerivation pkg (
        old:
        let
          newflags = lib.foldl' (acc: x: "${acc} ${x}") "" flags;
          oldflags = if (lib.hasAttr "NIX_CFLAGS_COMPILE" old) then "${old.NIX_CFLAGS_COMPILE}" else "";
        in
        {
          NIX_CFLAGS_COMPILE = "${oldflags} ${newflags}";
          # stdenv = pkgs.clang19Stdenv;
        }
      );

    optimizePackage =
      pkg:
      final.optimizeWithFlags pkg [
        "-O3"
        "-march=znver3"
        "-mtune=znver3"
      ];

    vpnNamespaceOpenPort =
      port: service:
      { ... }:
      {
        vpnNamespaces.wg = {
          portMappings = [
            {
              from = port;
              to = port;
            }
          ];

          openVPNPorts = [
            {
              port = port;
              protocol = "both";
            }
          ];
        };
        systemd.services.${service}.vpnConfinement = {
          enable = true;
          vpnNamespace = "wg";
        };
      };

    serviceMountWithZpool =
      serviceName: zpool: dirs:
      { pkgs, config, ... }:
      {
        systemd.services."${serviceName}-mounts" = {
          wants = [ "zfs.target" ] ++ lib.optionals (zpool != "") [ "zfs-import-${zpool}.service" ];
          after = lib.optionals (zpool != "") [ "zfs-import-${zpool}.service" ];
          before = [ "${serviceName}.service" ];

          serviceConfig = {
            Type = "oneshot";
            RemainAfterExit = true;
            ExecStart = lib.getExe (
              pkgs.writeShellApplication {
                name = "ensure-zfs-mounts-with-pool-${serviceName}";
                runtimeInputs = with pkgs; [
                  gawk
                  coreutils
                  config.boot.zfs.package
                ];

                text = ''
                  set -euo pipefail

                  echo "Ensuring ZFS mounts for service: ${serviceName}"
                  echo "Directories: ${lib.strings.concatStringsSep ", " dirs}"

                  # Validate mounts exist (ensureZfsMounts already has proper PATH)
                  ${lib.getExe pkgs.ensureZfsMounts} ${lib.strings.concatStringsSep " " dirs}

                  # Additional runtime check: verify paths are on correct zpool
                  ${lib.optionalString (zpool != "") ''
                    echo "Verifying ZFS mountpoints are on pool '${zpool}'..."

                    if ! zfs_list_output=$(zfs list -H -o name,mountpoint 2>&1); then
                      echo "ERROR: Failed to query ZFS datasets: $zfs_list_output" >&2
                      exit 1
                    fi

                    # This loop handles variable number of directories, shellcheck false positive
                    # shellcheck disable=SC2043
                    for target in ${lib.strings.concatStringsSep " " dirs}; do
                      echo "Checking: $target"

                      # Find dataset that has this mountpoint
                      dataset=$(echo "$zfs_list_output" | awk -v target="$target" '$2 == target {print $1; exit}')

                      if [ -z "$dataset" ]; then
                        echo "ERROR: No ZFS dataset found for mountpoint: $target" >&2
                        exit 1
                      fi

                      # Extract pool name from dataset (first part before /)
                      actual_pool=$(echo "$dataset" | cut -d'/' -f1)

                      if [ "$actual_pool" != "${zpool}" ]; then
                        echo "ERROR: ZFS pool mismatch for $target" >&2
                        echo " Expected pool: ${zpool}" >&2
                        echo " Actual pool:   $actual_pool" >&2
                        echo " Dataset:       $dataset" >&2
                        exit 1
                      fi

                      echo "$target is on $dataset (pool: $actual_pool)"
                    done

                    echo "All paths verified successfully on pool '${zpool}'"
                  ''}

                  echo "Mount validation completed for ${serviceName}"
                '';
              }
            );
          };
        };

        systemd.services.${serviceName} = {
          wants = [
            "${serviceName}-mounts.service"
          ];
          after = [
            "${serviceName}-mounts.service"
          ];
          requires = [
            "${serviceName}-mounts.service"
          ];
        };

        # assert that the pool is even enabled
        #assertions = lib.optionals (zpool != "") [
        #  {
        #    assertion = builtins.elem zpool config.boot.zfs.extraPools;
        #    message = "${zpool} is not enabled in `boot.zfs.extraPools`";
        #  }
        #];
      };
  }
)