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

with lib;

let
  cfg = config.services.wireguard-server;
in

{
  options.services.wireguard-server = {
    enable = mkEnableOption "WireGuard VPN server";

    address = mkOption {
      type = types.str;
      example = "10.2.2.1/24";
    };

    peers = mkOption {
      type = types.listOf (types.submodule {
        options = {
          publicKey = mkOption { type = types.str; };
          allowedIPs = mkOption { type = types.listOf types.str; };
        };
      });
      default = [];
    };
  };

  config = mkIf cfg.enable {
    environment.systemPackages = [ pkgs.wireguard-tools ];

    systemd.tmpfiles.rules = [
      "d /etc/wireguard 0700 root root - -"
      "f /etc/wireguard/private.key 0600 root root - -"
    ];

    systemd.services.wireguard-keygen = {
      description = "Generate WireGuard private key";
      before = [ "wg-quick-wg0.service" ];
      wantedBy = [ "multi-user.target" ];
      serviceConfig = {
        Type = "oneshot";
        RemainAfterExit = true;
      };
      script = ''
        if [ ! -s /etc/wireguard/private.key ]; then
          ${pkgs.wireguard-tools}/bin/wg genkey > /etc/wireguard/private.key
          chmod 600 /etc/wireguard/private.key
          echo "Public key: $(${pkgs.wireguard-tools}/bin/wg pubkey < /etc/wireguard/private.key)"
        fi
      '';
    };

    networking.wg-quick.interfaces.wg0 = {
      privateKeyFile = "/etc/wireguard/private.key";
      address = [ cfg.address ];
      listenPort = 51820;

      preUp = ''
        ${pkgs.iptables}/bin/iptables -A FORWARD -i wg0 -j ACCEPT
        ${pkgs.iptables}/bin/iptables -A FORWARD -o wg0 -j ACCEPT
        ${pkgs.iptables}/bin/iptables -t nat -A POSTROUTING -s 10.2.2.0/24 -o eth0 -j MASQUERADE
      '';

      postDown = ''
        ${pkgs.iptables}/bin/iptables -D FORWARD -i wg0 -j ACCEPT
        ${pkgs.iptables}/bin/iptables -D FORWARD -o wg0 -j ACCEPT
        ${pkgs.iptables}/bin/iptables -t nat -D POSTROUTING -s 10.2.2.0/24 -o eth0 -j MASQUERADE
      '';

      peers = map (peer: {
        inherit (peer) publicKey allowedIPs;
      }) cfg.peers;
    };

    networking.firewall = {
      allowedUDPPorts = [ 51820 ];
      trustedInterfaces = [ "wg0" ];
    };

    boot.kernel.sysctl = {
      "net.ipv4.ip_forward" = 1;
      "net.ipv6.conf.all.forwarding" = 1;
    };
  };
}