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

with lib;

let
  cfg = config.boot.initrd.network.wireguard;
in

{
  options = {
    boot.initrd.network.wireguard = {
      enable = mkEnableOption "Wireguard in initrd";
      privateKeyFile = mkOption {
        type = with types; nullOr str;
        default = null;
        description = "Path to file containing private key for wireguard interface";
      };
      ipv6s = mkOption {
        type = with types; listOf str;
        default = [];
        description = "IPv6 interface addresses";
      };
      ipv4s = mkOption {
        type = with types; listOf str;
        default = [];
        description = "IPv4 interface addresses";
      };
    };
  };

  config = mkIf cfg.enable {
    boot.initrd.kernelModules = [ "wireguard" ];

    boot.initrd.extraUtilsCommands = ''
      for BIN in ${pkgs.iproute2}/bin/*; do
        copy_bin_and_libs $BIN
      done
      [ -f "$out/bin/wg" ] && rm "$out/bin/wg"
      cp -pdv ${pkgs.wireguard-tools}/bin/.wg-wrapped $out/bin/wg
    '';

    boot.initrd.network.postCommands = ''
      ip link add dev "wg-initrd" type wireguard
      ip link set "wg-initrd" mtu 1280

      ${concatMapStringsSep "\n" (ip: ''
        ip address add "${ip}" dev "wg-initrd"
      '') cfg.ipv6s}
      ${concatMapStringsSep "\n" (ip: ''
        ip address add "${ip}" dev "wg-initrd"
      '') cfg.ipv4s}

      wg set "wg-initrd" private-key "/var/src/secrets/wireguard/wg-initrd"

      ip link set up "wg-initrd"

      wg set "wg-initrd" peer "2p1Jqs3bkXbXHFWE6vp1yxHIFoUaZQEARS2nJzbkuBA=" \
        endpoint "$(ip route get 2a01:4f8:c0c:15f1::1 ipproto udp dport 51820 &>/dev/null && echo '[2a01:4f8:c0c:15f1::1]:51820' || echo '78.47.183.82:51820')" \
        persistent-keepalive "25" \
        allowed-ips "0.0.0.0/0,::/0,10.20.30.0/24,2a01:4f8:c0c:15f1::/113"

      ip route add "::/0" dev "wg-initrd" table 1337
      ip route add "0.0.0.0/0" dev "wg-initrd" table 1337

      # Error: Argument "udp" is wrong: Invalid "ipproto" value
      # For some reason ip rule does not recognize "udp" as a value for "ipproto" in initrd,
      # so we pass the numeric value for it manually
      # This is from linux/include/uapi/linux/in.h
      IPPROTO_UDP=17

      ${concatMapStringsSep "\n" (ip: ''
        ip -6 rule add from "${ip}" lookup 1337 prio 19000
        ip -6 rule add from "${ip}" unreachable prio 19001
      '') cfg.ipv6s}
      ip -6 rule add to "2a01:4f8:c0c:15f1::1/128" ipproto $IPPROTO_UDP dport 51820 lookup main prio 20000
      ip -6 rule add to "2a01:4f8:c0c:15f1::1/128" ipproto $IPPROTO_UDP dport 51820 unreachable prio 20001
      ip -6 rule add lookup main prio 21000
      ip -6 rule add lookup 1337 prio 21001
      ip -6 rule add unreachable prio 21000


      ${concatMapStringsSep "\n" (ip: ''
        ip -4 rule add from "${ip}" lookup 1337 prio 19000
        ip -4 rule add from "${ip}" unreachable prio 19001
      '') cfg.ipv4s}
      ip -4 rule add to "78.47.183.82/32" ipproto $IPPROTO_UDP dport 51820 lookup main prio 20000
      ip -4 rule add to "78.47.183.82/32" ipproto $IPPROTO_UDP dport 51820 unreachable prio 20001
      ip -4 rule add lookup main prio 21000
      ip -4 rule add lookup 1337 prio 21001
      ip -4 rule add unreachable prio 21000
    '';

    boot.initrd.postMountCommands = ''
      ip link del dev "wg-initrd"

      ip -6 rule flush
      ip -4 rule flush
    '';

    boot.initrd.secrets = {
      "/var/src/secrets/wireguard/wg-initrd" = cfg.privateKeyFile;
    };
  };
}