{ config, lib, pkgs, ... }:
with lib;
let
  cfg = config.services.ezwg;

  peerNameReplacement = lib.replaceChars [ "/" "-" " " "+" "=" ] [
    "-"
    "\\x2d"
    "\\x20"
    "\\x2b"
    "\\x3d"
  ];

  ranges = serverIPs:
    let
      generateRangesScript =
        builtins.toFile "exclusionary-wildcard-ranges-generator.py" ''
          import ipaddress
          serverNetworks = [${map (ip: "ip_network('${ip}/32')") serverIPs}]
          ranges = [ipaddress.ip_network('0.0.0.0/0')]
          for serverNetwork in serverNetworks:
            ranges = map(lambda r: list(r.address_exclude(serverNetwork)), ranges)
          print(':'.join(ranges))
        '';
      rangesOutput = pkgs.runCommandNoCC "exclusionary-wildcard-ranges" { } ''
        ${pkgs.python3}/bin/python3 ${generateRangesScript} > $out
      '';
    in lib.splitString ":" (builtins.readFile "${rangesOutput}");

  subnet = vlanIP: vlanSize:
    let
      generateSubnetScript =
        builtins.toFile "subnet-without-host-bits-generator.py" ''
          import ipaddress
          n1 = ipaddress.ip_network('${vlanIP}/${toString vlanSize}', False)
          print(n1, end="")
        '';
      subnetOutput = pkgs.runCommandNoCC "subnet-without-host-bits" { } ''
        ${pkgs.python3}/bin/python3 ${generateSubnetScript} > $out
      '';
    in builtins.readFile "${subnetOutput}";

  serverOpts.options = {
    ip = mkOption {
      type = types.str;
      description = "The IP of the wg server";
    };
    port = mkOption {
      type = types.int;
      default = 51820;
      description = "The port of the wg server";
    };
    publicKey = mkOption {
      type = types.str;
      description = "The public key of the wg server";
    };
  };

  instanceOpts.options = {
    servers = mkOption {
      description = "Configuration of servers to connect to";
      default = { };
      type = with types; listOf (submodule serverOpts);
    };
    autoStart = mkOption {
      type = types.bool;
      default = true;
      description = "Start this wireguard tunnel by default";
    };
    proxy = mkOption {
      type = types.bool;
      default = true;
      description = "Route all your traffic through this connection";
    };
    vlanSize = mkOption {
      type = types.int;
      default = 24;
      description = "Size of your VLAN";
    };
    privateKeyFile = mkOption {
      type = types.str;
      description = "Private wg key";
    };
    vlanIP = mkOption {
      type = types.str;
      description = "The IP to use on the wg VLAN";
    };
  };
in {
  options.services.ezwg = {
    enable = mkEnableOption "Enable simple Wireguard connection";
    instances = mkOption {
      description = "Configuration of instances of Wireguard";
      default = { };
      type = with types; attrsOf (submodule instanceOpts);
    };
  };

  config = mkIf cfg.enable {
    networking.firewall.checkReversePath = false;

    systemd.paths = mapAttrs' (instName: inst: {
      name = "wireguard-${instName}";
      value = if inst.autoStart then { } else { wantedBy = mkForce [ ]; };
    }) cfg.instances;

    systemd.services = lib.listToAttrs (flatten (mapAttrsToList (instName: inst:
      [{
        name = "wireguard-${instName}";
        value = if inst.autoStart then { } else { wantedBy = mkForce [ ]; };
      }] ++ map (server: {
        name =
          "wireguard-${instName}-peer${peerNameReplacement server.publicKey}";
        value = if inst.autoStart then { } else { wantedBy = mkForce [ ]; };
      }) inst.servers) cfg.instances));

    networking.wireguard.interfaces = mapAttrs (instName: inst:
      let
        allowedIPs = if inst.proxy then
          ranges (map (s: s.ip) inst.servers)
        else
          [ (subnet inst.vlanIP inst.vlanSize) ];
      in {
        ips = [ "${inst.vlanIP}/${toString inst.vlanSize}" ];
        privateKeyFile = inst.privateKeyFile;
        peers = map (server: {
          inherit allowedIPs;
          publicKey = server.publicKey;
          endpoint = "${server.ip}:${toString server.port}";
          persistentKeepalive = 25;
        }) inst.servers;
      }) cfg.instances;

  };
}