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}-${zpool}";
                  runtimeInputs = with pkgs; [
                    gawk
                    coreutils
                    config.boot.zfs.package
                  ];

                  text = ''
                    set -euo pipefail

                    echo "Ensuring ZFS mounts for service: ${serviceName} (pool: ${zpool})"
                    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

                      # 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} (pool: ${zpool})"
                  '';
                }
              ))
            ];
          };
        };

        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`";
        #  }
        #];
      };
  }
)