{ config, lib, pkgs, ... }:
let
  cfg = config.orbekk.simple-firewall;
in
{
  options = {
    orbekk.simple-firewall = {
      enable = lib.mkEnableOption "Enable nftables firewall";

      allowedUDPPorts = lib.mkOption {
        type = lib.types.listOf lib.types.anything;
        default = [];
        description = "List of TCP ports to allow";
      };

      allowedTCPPorts = lib.mkOption {
        type = lib.types.listOf lib.types.anything;
        default = [];
        description = "List of TCP ports to allow";
      };
    };
  };

  config = lib.mkIf cfg.enable {
    networking.firewall.enable = lib.mkForce false;
    networking.nftables.enable = true;

    networking.nftables.ruleset =
      let
        csvPorts = ports: lib.concatStringsSep "," (map toString ports);
      in ''
      table inet filter {
        set allowed_tcp_ports {
          type inet_service
          flags interval
          ${lib.optionalString (cfg.allowedTCPPorts != [])''
            elements = {${csvPorts cfg.allowedTCPPorts}}
          ''}
        }
        set allowed_udp_ports {
          type inet_service
          flags interval
          ${lib.optionalString (cfg.allowedUDPPorts != [])''
            elements = {${csvPorts cfg.allowedUDPPorts}}
          ''}
        }

        chain input {
          type filter hook input priority 0

          iif lo accept

          ct state {established, related} accept

          ip protocol icmp limit rate 4/second counter accept
          ip6 nexthdr ipv6-icmp limit rate 4/second counter accept

	        ip6 saddr 2001:470:8e2e:1000::/64 counter accept

          tcp dport @allowed_tcp_ports counter accept
          udp dport @allowed_udp_ports counter accept

          counter drop
        }

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