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


let
  cfg = config.networking.networkmanager.nm-mullvad;

  template = { id, uuid, pubKey, endpoint }:
    ''
      [connection]
      id=${id}
      uuid=${uuid}
      type=wireguard
      ${lib.optionalString (!cfg.autoConnect.enable || cfg.autoConnect.profile != id) "autoconnect=false"}
      interface-name=${id}

      [wireguard]
      ${lib.optionalString (cfg.listenPort != null) "listen-port=${toString cfg.listenPort}"}
      private-key=${if cfg.privateKey != null then cfg.privateKey else "@PRIVKEY@"}

      [wireguard-peer.${pubKey}]
      endpoint=${endpoint}:51820
      allowed-ips=0.0.0.0/0;::/0;

      [ipv4]
      address1=${if cfg.ipv4Address != null then cfg.ipv4Address else "@IPV4ADDRESS@"}
      dns=10.64.0.1
      dns-search=~;
      method=manual

      [ipv6]
      addr-gen-mode=default
      address1=${if cfg.ipv6Address != null then cfg.ipv6Address else "@IPV6ADDRESS@"}
      method=manual

      [proxy]
    '';

  serversList = import ./mullvad-servers-list.nix;

  allServerIds = builtins.map (x: x.id) serversList;

  availableServersList = builtins.filter (x: lib.elem x.id cfg.availableServers) serversList;

  serverEntryToConnFilename = serverEntry:
    "NetworkManager/system-connections/${serverEntry.id}.nmconnection";

  serverEntryToConnFile = serverEntry:
    {
      "${serverEntryToConnFilename serverEntry}" = {
        text = template
          {
            inherit (serverEntry) id uuid pubKey endpoint;
          };
        mode = "0600";
      };
    };

  connFiles = lib.fold (x: acc: lib.recursiveUpdate (serverEntryToConnFile x) acc) { } availableServersList;

in
{
  options.networking.networkmanager.nm-mullvad = {
    enable = lib.mkEnableOption "Mullvad VPN profile for NetworkManager";

    listenPort = lib.mkOption {
      type = lib.types.nullOr lib.types.number;
      default = null;

      description = ''
        Set the port to listen on. Useful if you have a firewall active.
      '';
    };

    openFirewall = lib.mkOption {
      type = lib.types.bool;
      default = false;
      description = ''
        Whether to open firewall for given `listenPort`.
      '';
    };

    autoConnect = {
      enable = lib.mkEnableOption "autoconnect for `autoConnect.profile`";

      profile = lib.mkOption {
        type = lib.types.nullOr lib.types.str;
        default = null;
        description = ''
          ID of the profile to autoconnect to.
        '';
      };
    };

    privateKey = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = ''
        WireGuard private key to use for authenticating with Mullvad.
      '';
    };

    privateKeyPath = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = ''
        Path for the file containing WireGuard private key to use for
        authenticating with Mullvad.
      '';
    };

    ipv4Address = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = ''
        IPv4 address to assign to the WireGuard interface.
        This should include the CIDR notation, e.g.) `10.64.0.1/32`
      '';
    };

    ipv4AddressPath = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = ''
        Path for the file containing the IPv4 address to assign to the
        WireGuard interface.
        This should include the CIDR notation, e.g.) `10.64.0.1/32`
      '';
    };

    ipv6Address = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = ''
        IPv6 address to assign to the wireguard interface.
        This should include the CIDR notation, e.g.) `fc00:bbbb:bbbb:bb01::1/128`
      '';
    };

    ipv6AddressPath = lib.mkOption {
      type = lib.types.nullOr lib.types.str;
      default = null;
      description = ''
        Path for the file containing the IPv6 address to assign to the
        WireGuard interface.
        This should include the CIDR notation, e.g.) `fc00:bbbb:bbbb:bb01::1/128`
      '';
    };

    availableServers = lib.mkOption
      {
        type = lib.types.listOf lib.types.str;
        default = allServerIds;
        description = ''
          List of servers that will be made available.
          Defaults to all servers, which might introduce clutter.
        '';
      };
  };

  config = lib.mkIf cfg.enable {
    assertions = [
      {
        assertion = cfg.openFirewall == true -> cfg.listenPort != null;
        message = ''
          A custom `listenPort` must be given when `openFirewall` is set.
        '';
      }
      {
        assertion = cfg.autoConnect.enable == true -> cfg.autoConnect.profile != null;
        message = ''
          Profile to use for autoconnect must be set when enabling autoconnect.
        '';
      }
      {
        assertion = cfg.privateKey != null || cfg.privateKeyPath != null;
        message = ''
          Either `privateKey` or `privateKeyPath` must be set.
        '';
      }
      {
        assertion = cfg.ipv4Address != null || cfg.ipv4AddressPath != null;
        message = ''
          Either `ipv4Address` or `ipv4AddressPath` must be set.
        '';
      }
      {
        assertion = cfg.ipv6Address != null || cfg.ipv6AddressPath != null;
        message = ''
          Either `ipv6Address` or `ipv6AddressPath` must be set.
        '';
      }
    ];

    environment.etc = connFiles;

    system.activationScripts.nm-mullvad-substitutePrivateKey = lib.mkIf (cfg.privateKeyPath != null) (lib.stringAfter [ "etc" "specialfs" "var" ]
        ''
        if [ -f "${cfg.privateKeyPath}" ]; then
          ${pkgs.systemd}/bin/systemd-cat -t nixos echo "<6>loading networking.networkmanager.nm-mullvad.privateKeyPath from ${cfg.privateKeyPath}"
          ${lib.concatMapStringsSep "\n"
            (f: "${pkgs.gnused}/bin/sed -ie \"s_@PRIVKEY@_$(< ${cfg.privateKeyPath})_\" ${f}")
            (builtins.map (s: "/etc/${serverEntryToConnFilename s}") availableServersList)}
        else
            warning_msg="WARNING: networking.networkmanager.nm-mullvad.privateKeyPath (${cfg.privateKeyPath}) does not exist."
            echo "$warning_msg"
            ${pkgs.systemd}/bin/systemd-cat -t nixos echo "<4>$warning_msg"
        fi
      '');

    system.activationScripts.nm-mullvad-substituteIpv4Address = lib.mkIf (cfg.ipv4AddressPath != null) (lib.stringAfter [ "etc" "specialfs" "var" ]
        ''
        if [ -f "${cfg.ipv4AddressPath}" ]; then
          ${pkgs.systemd}/bin/systemd-cat -t nixos echo "<6>loading networking.networkmanager.nm-mullvad.ipv4AddressPath from ${cfg.ipv4AddressPath}"
          ${lib.concatMapStringsSep "\n"
            (f: "${pkgs.gnused}/bin/sed -ie \"s!@IPV4ADDRESS@!$(< ${cfg.ipv4AddressPath})!\" ${f}")
            (builtins.map (s: "/etc/${serverEntryToConnFilename s}") availableServersList)}
        else
            warning_msg="WARNING: networking.networkmanager.nm-mullvad.ipv4AddressPath (${cfg.ipv4AddressPath}) does not exist."
            echo "$warning_msg"
            ${pkgs.systemd}/bin/systemd-cat -t nixos echo "<4>$warning_msg"
        fi
      '');

    system.activationScripts.nm-mullvad-substituteIpv6Address = lib.mkIf (cfg.ipv6AddressPath != null) (lib.stringAfter [ "etc" "specialfs" "var" ]
        ''
        if [ -f "${cfg.ipv6AddressPath}" ]; then
          ${pkgs.systemd}/bin/systemd-cat -t nixos echo "<6>loading networking.networkmanager.nm-mullvad.ipv6AddressPath from ${cfg.ipv6AddressPath}"
          ${lib.concatMapStringsSep "\n"
            (f: "${pkgs.gnused}/bin/sed -ie \"s!@IPV6ADDRESS@!$(< ${cfg.ipv6AddressPath})!\" ${f}")
            (builtins.map (s: "/etc/${serverEntryToConnFilename s}") availableServersList)}
        else
            warning_msg="WARNING: networking.networkmanager.nm-mullvad.ipv6AddressPath (${cfg.ipv6AddressPath}) does not exist."
            echo "$warning_msg"
            ${pkgs.systemd}/bin/systemd-cat -t nixos echo "<4>$warning_msg"
        fi
      '');

    networking.firewall.allowedUDPPorts = lib.mkIf cfg.openFirewall [ cfg.listenPort ];
    networking.firewall.logReversePathDrops = lib.mkIf cfg.openFirewall true;
    networking.firewall.extraCommands = lib.mkIf cfg.openFirewall ''
     ip46tables -t mangle -I nixos-fw-rpfilter -p udp -m udp --sport ${toString cfg.listenPort} -j RETURN
     ip46tables -t mangle -I nixos-fw-rpfilter -p udp -m udp --dport ${toString cfg.listenPort} -j RETURN
    '';
    networking.firewall.extraStopCommands = lib.mkIf cfg.openFirewall ''
     ip46tables -t mangle -D nixos-fw-rpfilter -p udp -m udp --sport ${toString cfg.listenPort} -j RETURN || true
     ip46tables -t mangle -D nixos-fw-rpfilter -p udp -m udp --dport ${toString cfg.listenPort} -j RETURN || true
    '';
  };
}