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

with lib;

let
  listenPort = 51820;
  udp2rawPort = 51821;
  subnet = "2a03:4000:52:ada:1";
  subnetLength = 80;
  hostLength = subnetLength + 16;

  links = [
    { from = "vidhar";
      to = "surtr";
      endpointHost = "surtr.yggdrasil.li";
      persistentKeepalive = 25;
      dynamicEndpointRefreshSeconds = 86400;
    }
    { from = "sif";
      to = "surtr";
      endpointHost = "surtr.yggdrasil.li";
      persistentKeepalive = 25;
      dynamicEndpointRefreshSeconds = 86400;
    }
  ];
  routes = [
    { from = "sif";
      to = "vidhar";
      via = "surtr";
    }
    { from = "vidhar";
      to = "sif";
      via = "surtr";
    }
  ];
  hostIPs = {
    surtr = ["${subnet}::/${toString hostLength}"];
    vidhar = ["${subnet}:1::/${toString hostLength}"];
    sif = ["${subnet}:2::/${toString hostLength}"];
  };

  mkPublicKeyPath = host: ./hosts + "/${host}.pub";
  mkPrivateKeyPath = host: ./hosts + "/${host}.priv";
  
  publicKeyPath = mkPublicKeyPath hostName;
  privateKeyPath = mkPrivateKeyPath hostName;
  inNetwork = pathExists privateKeyPath && pathExists publicKeyPath;
  hostLinks = filter ({ from, to, ... }: from == hostName || to == hostName) links;
  hostRoutes = filter ({ from, to, ... }: from == hostName || to == hostName) routes;
  isRouter = inNetwork && any ({via, ...}: via == hostName) routes;
  linkToPeer = ix: opts@{from, to, ...}:
    let
      other = if from == hostName then to else from;
    in {
      allowedIPs = hostIPs.${other} ++ concatMap (rArgs: if rArgs.from != hostName || rArgs.via != to then [] else hostIPs.${rArgs.to}) routes;
      publicKey = trim (readFile (mkPublicKeyPath other));
    } // (optionalAttrs (from == hostName) (filterAttrs (n: _v: !(elem n ["from" "to" "endpointHost"])) opts // optionalAttrs (opts ? "endpointHost") { endpoint = "localhost:${toString (udp2rawPort + ix)}"; }));

  trim = str: if hasSuffix "\n" str then trim (removeSuffix "\n" str) else str;
  stripSubnet = addr: let matchRes = builtins.match "^(.*)/[0-9]+$" addr; in if matchRes == null then addr else elemAt matchRes 0;
in {
  config = {
    assertions = [
      { assertion = inNetwork || !(pathExists privateKeyPath || pathExists publicKeyPath);
        message = "yggdrasil-wg: Either both public and private keys must exist or neither.";
      }
      { assertion = !inNetwork || (hostIPs ? "${hostName}");
        message = "yggdrasil-wg: Entry in hostIPs must exist.";
      }
    ] ++ map ({from, to, ...}: let other = if from == hostName then to else from; in { assertion = pathExists (mkPublicKeyPath other); message = "yggdrasil-wg: This host (${hostName}) has a link with ‘${other}’, but no public key is available for ‘${other}’."; }) hostLinks;

    networking.wireguard.interfaces = mkIf inNetwork {
      yggdrasil = {
        allowedIPsAsRoutes = false;
        inherit listenPort;
        ips = hostIPs.${hostName};
        peers = imap0 linkToPeer hostLinks;
        privateKeyFile = config.sops.secrets."yggdrasil-wg.priv".path;
        postSetup = ''
          ${concatMapStringsSep "\n" (linkArgs: let other = if linkArgs.from == hostName then linkArgs.to else linkArgs.from; in concatMapStringsSep "\n" (otherIP: "ip route replace \"${otherIP}\" dev \"yggdrasil\" table \"main\"") hostIPs.${other}) hostLinks}
          ${concatMapStringsSep "\n" (routeArgs: let other = if routeArgs.from == hostName then routeArgs.to else routeArgs.from; in concatMapStringsSep "\n" (otherIP: concatMapStringsSep "\n" (viaIP: "ip route replace \"${otherIP}\" via \"${viaIP}\" dev \"yggdrasil\" table \"main\"") (map stripSubnet hostIPs.${routeArgs.via})) hostIPs.${other}) hostRoutes}
        '';
      };
    };

    systemd.services = listToAttrs (filter ({ value, ...}: value != null) (imap0 (ix: opts@{to, from, ...}: let other = if from == hostName then to else from; in nameValuePair "yggdrasil-udp2raw@${other}" (if opts ? "endpointHost" then {
      requiredBy = ["wireguard-yggdrasil.service"];
      
      serviceConfig = {
        ExecStart = "${pkgs.udp2raw}/bin/udp2raw ${if from == hostName then "-c -l localhost:${toString (udp2rawPort + ix)} -r ${opts.endpointHost}:${toString (udp2rawPort + ix)}" else "-s -l 0.0.0.0:${toString (udp2rawPort + ix)} -r localhost:${toString listenPort}"} -k tmpkey --auth-mode hmac_sha1 --raw-mode faketcp -a";
      };
    } else null)) hostLinks)) // {
      firewall.path = optionals isRouter [pkgs.procps];
    };

    sops.secrets = mkIf (pathExists privateKeyPath) {
      "yggdrasil-wg.priv" = {
        format = "binary";
        sopsFile = privateKeyPath;
      };
    };

    networking.hosts = mkIf inNetwork (listToAttrs (concatMap ({name, value}: map (ip: nameValuePair (stripSubnet ip) ["${name}.yggdrasil"]) value) (mapAttrsToList nameValuePair hostIPs)));

    networking.firewall = mkIf isRouter {
      extraCommands = ''
        ip6tables -A FORWARD -i yggdrasil -o yggdrasil -j nixos-fw-accept
        ip46tables -A FORWARD -j nixos-fw-log-refuse
        sysctl net.ipv6.conf.all.forwarding=1
      '';
      extraStopCommands = ''
        sysctl net.ipv6.conf.all.forwarding=0
        ip46tables -D FORWARD -j nixos-fw-log-refuse || true
        ip6tables -D FORWARD -i yggdrasil -o yggdrasil -j nixos-fw-accept || true
      '';
    };
  };
}