{ config, lib, pkgs, ... }:
with lib;
let
  cfg = config.orbekk.router;

  devices = ["eno1" "eno2"];
  vpnMark = 3;
  mullvadMark = 2;
  heMark = 1;

  mullvadPort = config.orbekk.mullvad.listenPort;
  vpnPort = config.orbekk.vpn.listenPort;
  vpnPrefix = "2001:470:8e2e:1000";

  router-netns-up = pkgs.writeScript "router-netns-up" ''
    #!${pkgs.bash}/bin/bash
    if ip netns list | grep -q router; then
      echo "Netns setup is not idempotent. Needs restart."
      exit 0
    fi
    ip netns add router
    ip netns exec router ip link set lo up
    ${lib.concatMapStrings (device: ''
      ip link set ${device} netns router
    '')
    devices}
    ip link add router-vport type veth peer name dragon-vport netns router
    ip link add vpn-vport type veth peer name dragonvpn-vport netns router
    ip link add wg-vpn-vport type veth peer name wg-vpnlan-vport netns router

    ip netns add vpn
    ip netns exec vpn ip link set lo up
    ip link set vpn-vport netns vpn

    ip netns exec router ${pkgs.procps}/bin/sysctl -w net.ipv4.conf.default.forwarding=1
    ip netns exec router ${pkgs.procps}/bin/sysctl -w net.ipv4.conf.all.forwarding=1
    ip netns exec router ${pkgs.procps}/bin/sysctl -w net.ipv6.conf.default.forwarding=1
    ip netns exec router ${pkgs.procps}/bin/sysctl -w net.ipv6.conf.all.forwarding=1
  '';

  router-netns-down = pkgs.writeScript "router-netns-down" ''
    #!${pkgs.bash}/bin/bash
  '';

  router-config = { config, lib, pkgs, ... }: {
    system.stateVersion = "22.05";

    environment.systemPackages = with pkgs; [ tcpdump ];
    virtualisation.vswitch.enable = true;
    virtualisation.vswitch.resetOnStart = false;
    networking.vswitches.kjlan = {
      interfaces.wan-vport = { vlan = 10; type = "internal"; };
      interfaces.lan-vport = { vlan = 100; type = "internal"; };
      interfaces.vpnlan-vport = { vlan = 30; type = "internal"; };
      interfaces.wg-vpnlan-vport = { vlan = 200; type = "internal"; };
      interfaces.servers-vport = { vlan = 20; type = "internal"; };
      interfaces.admin-vport = { vlan = 255; type = "internal"; };
      interfaces.dragon-vport = { vlan = 20; };
      interfaces.dragonvpn-vport = { vlan = 30; };

      extraOvsctlCmds = ''
        add-bond kjlan bond0 eno1 eno2 lacp=active
        set interface wan-vport mac=\"3c:97:0e:19:7e:5c\"
      '';
    };

    networking.interfaces.eno1 = {};
    networking.interfaces.eno2 = {};
    networking.interfaces.wan-vport = {
      useDHCP = true;
    };
    networking.interfaces.lan-vport = {
      ipv4.addresses = [{address = "172.20.100.1"; prefixLength = 24;}];
      ipv6.addresses = [{address = "2001:470:8e2e:100::1"; prefixLength = 64;}];
    };
    networking.interfaces.servers-vport = {
      ipv4.addresses = [{address = "172.20.20.1"; prefixLength = 24;}];
      ipv6.addresses = [{address = "2001:470:8e2e:20::1"; prefixLength = 64;}];
    };
    networking.interfaces.admin-vport = {
      ipv4.addresses = [{address = "10.10.255.18"; prefixLength = 24;}];
      ipv4.routes = [{address = "10.10.255.0"; prefixLength = 24;}];
    };
    networking.interfaces.wg-vpnlan-vport = {
      ipv6.addresses = [{address = "2001:470:8e2e:1001::1"; prefixLength = 64; }];
    };
    networking.interfaces.vpnlan-vport = {
      ipv4.addresses = [{address = "172.20.30.1"; prefixLength = 24;}];
      ipv6.addresses = [{address = "2001:470:8e2e:30::1"; prefixLength = 64;}];
    };
    networking.sits.he0 = {
     dev = "wan-vport";
     remote = "209.51.161.14";
    };
    networking.interfaces.he0.ipv6.addresses = [{
      address = "2001:470:1f06:1194::2";
      prefixLength = 64;
    }];
    systemd.services.he0-netdev.after = ["kjlan-netdev.service"];

    networking.iproute2.enable = true;
    networking.iproute2.rttablesExtraConfig = ''
      ${toString vpnMark} vpn
      ${toString mullvadMark} mullvad
      ${toString heMark} he
    '';
    systemd.services.network-route-setup = {
      description = "HE tunnel route setup";
      requires = [ "network-online.target" ];
      after = [ "network.target" "network-online.target" ];
      wantedBy = [ "multi-user.target" ];
      path = [ pkgs.iproute ];
      script = ''
        ip -6 rule add fwmark ${toString heMark} table he || true
        ip -6 route replace default dev he0 table he
        ip -6 route replace default dev he0 metric 1

        # ip rule add fwmark ${toString vpnMark} table vpn || true
        # ip -6 rule add fwmark ${toString vpnMark} table vpn || true

        ip rule add fwmark ${toString mullvadMark} table mullvad || true
        ip -6 rule add fwmark ${toString mullvadMark} table mullvad || true

        ip route replace default dev mullvad table mullvad
        ip -6 route replace default dev mullvad table mullvad

        ip route replace default dev mullvad table vpn
        ip -6 route replace default dev mullvad table vpn

        ip -6 route flush cache
        ip route flush cache
      '';
    };

    services.dnsmasq = {
      enable = true;
      settings.server = [ "1.1.1.1" "8.8.8.8" "8.8.4.4" ];
      resolveLocalQueries = false;

      extraConfig = ''
        no-resolv
        no-hosts
        log-debug

        dhcp-authoritative
        enable-ra

        address=/localhost/::1
        address=/localhost/127.0.0.1

        dhcp-range=tag:servers-vport,172.20.20.10,172.20.20.254,5m
        dhcp-option=tag:servers-vport,option:router,172.20.20.1
        dhcp-option=tag:servers-vport,option:dns-server,172.20.20.1
        dhcp-range=tag:servers-vport,::,static,constructor:servers-vport,5m
        dhcp-host=id:dragon,tag:servers-vport,172.20.20.2
        dhcp-host=id:00:01:00:01:21:a2:4e:a8:d0:bf:9c:45:a6:ec,tag:servers-vport,[::d]
        # dhcp-host=id:dragon,::d

        dhcp-range=tag:lan-vport,172.20.100.10,172.20.100.254,5m
        dhcp-option=tag:lan-vport,option:router,172.20.100.1
        dhcp-option=tag:lan-vport,option:dns-server,172.20.100.1
        dhcp-range=tag:lan-vport,::2,::1000,constructor:lan-vport,ra-only

        dhcp-range=tag:vpnlan-vport,172.20.30.10,172.20.30.254,5m
        dhcp-option=tag:vpnlan-vport,option:router,172.20.30.1
        dhcp-option=tag:vpnlan-vport,option:dns-server,193.138.218.74
        dhcp-range=tag:vpnlan-vport,::2,::1000,constructor:vpnlan-vport,ra-only,5m
        dhcp-host=id:00:04:33:32:31:37:37:31:58:4d:32:35:31:37:30:30:4a:44,tag:vpnlan-vport,[::2]
        dhcp-host=id:vpn,tag:vpnlan-vport,172.20.30.2

        dhcp-range=tag:wg-vpnlan-vport,::2,::1000,constructor:wg-vpnlan-vport,ra-only,5m
      '';
    };

    networking.dhcpcd = {
      enable = true;
      extraConfig = ''
        noipv6rs
        noipv6
        nohook resolv.conf
        interface wan-vport
        dhcp
      '';
    };

    networking.firewall.enable = false;
    networking.nftables.enable = true;
    networking.nftables.ruleset =
      let
        ports-to-csv = ports: concatStringsSep "," (map toString ports);
      in ''
        define SERVER_WAN_PORTS = {
          ssh, domain, http, https,
          ${toString mullvadPort}
        }
        define SERVER_LAN_PORTS = {
          tftp, 139, 445, 137, 138,
          1080, # socks
        }
        table inet filter {
          chain input {
            type filter hook input priority 0
            iifname lo accept

            ct state {established, related} counter accept
            meta l4proto {tcp, udp} th dport {bootps, bootpc, domain, dhcpv6-client, dhcpv6-server, ${toString vpnPort}} counter accept

            ip protocol ipv6 counter accept comment "sit tunnel"
            ip protocol icmp limit rate 4/second counter accept comment "icmp v4"

            ip6 nexthdr ipv6-icmp counter accept comment "accept all ICMP types"

            iifname wan-vport counter drop
            meta nftrace set 1
            counter drop
          }

          chain output {
            type filter hook output priority 0
            counter accept
          }

          chain forward {
            type filter hook forward priority 0; policy drop

            ip protocol icmp limit rate 4/second counter accept comment "icmp v4"
            ip6 nexthdr ipv6-icmp limit rate 4/second counter accept comment "accept all ICMP types"

            ct state vmap { established : accept, related : accept, invalid : drop }

            # Don't allow accidental vpn forwarding to wan.
            iifname vpnlan-vport oifname wan-vport counter reject
            oifname he0 counter accept
            oifname wan-vport counter accept
            oifname mullvad counter accept
            iifname servers-vport oifname wg-vpn counter accept
            iifname wg-vpn-vport oifname wg-vpn counter accept
            iifname wg-vpn oifname wg-vpn-vport counter accept

            # Transmission
            ip daddr 172.20.20.2 th dport {9091, 56732} counter accept;
            oifname servers-vport meta l4proto {tcp, udp} th dport $SERVER_WAN_PORTS counter accept
            oifname lan-vport meta l4proto {tcp, udp} th dport 34197  counter accept
            iifname lan-vport oifname vpnlan-vport counter accept
            iifname lan-vport oifname servers-vport meta l4proto {tcp, udp} th dport $SERVER_LAN_PORTS counter accept
            iifname vpnlan-vport oifname servers-vport meta l4proto {tcp, udp} th dport $SERVER_LAN_PORTS counter accept
            iifname servers-vport counter accept

            counter drop
          }
        }

        table inet mangle {
          chain prerouting {
            type filter hook prerouting priority -150
            # ip6 saddr 2001:470:8e2e::/48 ip6 daddr != 2001:470:8e2e::/48 ip6 daddr != fe80::/64 meta nftrace set 1
            ip6 saddr 2001:470:8e2e::/48 ip6 daddr != 2001:470:8e2e::/48 ip6 daddr != fe80::/64 meta mark set ${toString heMark}
            meta nfproto ipv4 iifname vpnlan-vport ip daddr != 172.20.0.0/16 meta mark set ${toString mullvadMark}
            meta nfproto ipv6 ip6 daddr != 2001:470:8e2e::/48 ip6 daddr != fe80::/60 iifname vpnlan-vport meta mark set ${toString mullvadMark}
          }
        }

        table inet nat {
          chain prerouting {
            type nat hook prerouting priority -100; policy accept
            meta nfproto ipv4 iifname wan-vport udp dport 34197 counter dnat to 172.20.100.214
            # OpenTTD:
            meta nfproto ipv4 iifname wan-vport meta l4proto {tcp, udp} th dport 3979 counter dnat to 172.20.30.84
            meta nfproto ipv4 iifname wan-vport meta l4proto {tcp, udp} th dport $SERVER_WAN_PORTS counter dnat to 172.20.20.2
            meta nfproto ipv4 iifname mullvad tcp dport 56732 counter dnat to 172.20.20.2
          }
          chain postrouting {
            type nat hook postrouting priority 100; policy accept
            ip saddr 172.16.0.0/12 oifname {"wan-vport"} counter masquerade
            ip saddr 172.16.0.0/12 oifname {"mullvad"} counter masquerade
            # Nat66 on VPN :(
            meta nfproto ipv6 oifname {"mullvad"} counter masquerade
          }
        }
      '';
  };
in {
  options = {
    orbekk.router = {
      enable = mkEnableOption "Enable router config";
    };
  };

  config = mkIf cfg.enable {
    boot.kernel.sysctl = {
      "net.ipv4.conf.all.forwarding" = true;
      "net.ipv4.conf.default.forwarding" = true;
      "net.ipv6.conf.all.forwarding" = true;
      "net.ipv6.conf.default.forwarding" = true;
    };

    systemd.services."router-netns" = {
      description = "router network namespace";
      before = ["network.target"];
      after = ["network-interfaces.target"];
      path = with pkgs; [bash iproute];
      serviceConfig = {
        Type = "oneshot";
        RemainAfterExit = "yes";
        ExecStart = "${router-netns-up}";
        ExecStop = "${router-netns-down}";
      };
    };

    orbekk.mullvad.enable = true;
    networking.wireguard.enable = true;
    networking.wireguard.interfaces.mullvad.interfaceNamespace = "router";

    systemd.services."container@router" = {
      after = ["router-netns.service"];
      requires = ["router-netns.service"];
      wantedBy = ["network.target"];
    };
    systemd.services.dhcpcd.partOf = ["container@router.service"];

    containers.router = {
      autoStart = true;
      extraFlags = ["--network-namespace-path" "/var/run/netns/router"];
      privateNetwork = false;
      config = router-config;
      additionalCapabilities = ["CAP_SYS_ADMIN" "CAP_NET_ADMIN"];
    };

    containers.vpn = {
      autoStart = true;
      extraFlags = ["--network-namespace-path" "/var/run/netns/vpn"];
      privateNetwork = false;
      config = { config, lib, pkgs, ... }: {
        system.stateVersion = "22.11";
        networking.firewall.enable = false;
        networking.interfaces.vpn-vport.useDHCP = true;
        networking.dhcpcd.extraConfig = ''
          clientid vpn
        '';
      };
      additionalCapabilities = ["CAP_NET_ADMIN"];
    };

    age.secrets.dragon-wireguard-key.file = ./.
        + "/../secrets/dragon-wireguard-key.age";
    networking.wireguard.interfaces.wg-vpn = {
      # fwMark = "${toString vpnMark}";
      socketNamespace = "router";
      interfaceNamespace = "router";
      ips = [ "${vpnPrefix}::1/128" ];
      privateKeyFile = config.age.secrets.dragon-wireguard-key.path;
      listenPort = vpnPort;
      peers = let mkPeer = host: ip: {
        name = host;
        allowedIPs = [ "${vpnPrefix}::${ip}/128" ];
        publicKey = (builtins.readFile (../secrets + "/${host}-wireguard-key.pub"));
        endpoint = null;
      }; in [
        (mkPeer "tiny1" "1001")
        (mkPeer "firelink" "2001")
        (mkPeer "pincer" "2002")
        (mkPeer "steamdeck" "2003")
        (mkPeer "trygve" "2004")
        {
          name = "fedora-pincer";
          publicKey = "tvAFrds/N7JRRA/OH9rTDYMqaxnbSxLCcRm3SatcEmc=";
          allowedIPs = [ "${vpnPrefix}::2005/128" ];
        }
      ];
    };
    systemd.services.wg-vpn.after = ["container@router.target"];

    services.ddclient = {
      enable = true;
      configFile = "/opt/secret/he-ddclient.conf";
    };
    systemd.services.ddclient.after = ["container@router.target"];

    systemd.timers.update-dynamic-dns = {
      wantedBy = ["multi-user.target"];
      timerConfig = {
        Persistent = true;
        OnBootSec = "5m";
        OnUnitActiveSec = "5m";
      };
    };
    systemd.services.update-dynamic-dns = {
      description = "Update dynamic dns records";
      after = ["container@router.target"];
      path = with pkgs; [ bash dnsutils nettools gawk iproute curl ];
      startLimitIntervalSec = 5;
      script = toString ../tools/update-dns.sh;
    };
  };
}