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

let
  cfg = config.services.caveman;

  blocklistPath = "/etc/caveman.blocklist";
  profanityPath = "/etc/caveman.profanity";

  dbUser = "caveman";
  dbPassword = "c";

  hunterDefaultSettings = {
    redis = "redis://127.0.0.1:${toString cfg.redis.port}/";
    redis_password_file = cfg.redis.passwordFile;
    database = "host=localhost user=${dbUser} password=${dbPassword} dbname=caveman";
    hosts = [ "mastodon.social" ];
    max_workers = 16;
    prometheus_port = 9101;
    blocklist = blocklistPath;
  };

  hunterSettings = lib.recursiveUpdate hunterDefaultSettings cfg.hunter.settings;

  hunterConfigFile = builtins.toFile "hunter.yaml" (
    builtins.toJSON hunterSettings
  );

  butcherDefaultSettings = {
    redis = "redis://127.0.0.1:${toString cfg.redis.port}/";
    redis_password_file = cfg.redis.passwordFile;
    profanity = profanityPath;
  };

  butcherSettings = lib.recursiveUpdate butcherDefaultSettings cfg.butcher.settings;

  butcherConfigFile = builtins.toFile "butcher.yaml" (
    builtins.toJSON butcherSettings
  );

  gathererDefaultSettings = {
    redis = "redis://127.0.0.1:${toString cfg.redis.port}/";
    redis_password_file = cfg.redis.passwordFile;
    database = "host=localhost user=${dbUser} password=${dbPassword} dbname=caveman";
    listen_port = 8000;
  };

  gathererSettings = lib.recursiveUpdate gathererDefaultSettings cfg.gatherer.settings;

  gathererConfigFile = builtins.toFile "gatherer.yaml" (
    builtins.toJSON gathererSettings
  );

  sieveDefaultSettings = {
    redis = "redis://127.0.0.1:${toString cfg.redis.port}/";
    redis_password_file = cfg.redis.passwordFile;
    in_topic = "relay-in";
    prometheus_port = 9102;
    blocklist = blocklistPath;
  };

  sieveSettings = lib.recursiveUpdate sieveDefaultSettings cfg.sieve.settings;

  sieveConfigFile = builtins.toFile "sieve.yaml" (
    builtins.toJSON sieveSettings
  );

  smokestackDefaultSettings = {
    redis = "redis://127.0.0.1:${toString cfg.redis.port}/";
    redis_password_file = cfg.redis.passwordFile;
    listen_port = 23;
  };

  smokestackSettings = lib.recursiveUpdate smokestackDefaultSettings cfg.smokestack.settings;

  smokestackConfigFile = builtins.toFile "smokestack.yaml" (
    builtins.toJSON smokestackSettings
  );

  limitNOFILE = 1000000;

in
{
  options.services.caveman = with lib; {
    redis.port = mkOption {
      type = types.int;
      default = 6379;
    };
    redis.maxmemory = mkOption {
      type = types.int;
      default = 1024 * 1024 * 1024;
    };
    redis.maxmemory-samples = mkOption {
      type = types.int;
      default = 8;
    };
    redis.passwordFile = mkOption {
      type = types.path;
    };

    hunter.enable = mkEnableOption "caveman hunter";

    hunter.settings = mkOption {
      type = types.anything;
      default = hunterDefaultSettings;
    };

    hunter.logLevel = mkOption {
      type = types.enum [ "ERROR" "WARN" "INFO" "DEBUG" "TRACE" ];
      default = "DEBUG";
    };

    butcher.enable = mkEnableOption "caveman butcher";

    butcher.settings = mkOption {
      type = types.anything;
      default = butcherDefaultSettings;
    };

    butcher.logLevel = mkOption {
      type = types.enum [ "ERROR" "WARN" "INFO" "DEBUG" "TRACE" ];
      default = "DEBUG";
    };

    gatherer.enable = mkEnableOption "caveman gatherer";

    gatherer.settings = mkOption {
      type = types.anything;
      default = gathererDefaultSettings;
    };

    gatherer.logLevel = mkOption {
      type = types.enum [ "ERROR" "WARN" "INFO" "DEBUG" "TRACE" ];
      default = "DEBUG";
    };

    sieve.enable = mkEnableOption "caveman sieve";

    sieve.settings = mkOption {
      type = types.anything;
      default = sieveDefaultSettings;
    };

    sieve.logLevel = mkOption {
      type = types.enum [ "ERROR" "WARN" "INFO" "DEBUG" "TRACE" ];
      default = "DEBUG";
    };

    smokestack.enable = mkEnableOption "caveman smokestack";

    smokestack.settings = mkOption {
      type = types.anything;
      default = smokestackDefaultSettings;
    };

    smokestack.logLevel = mkOption {
      type = types.enum [ "ERROR" "WARN" "INFO" "DEBUG" "TRACE" ];
      default = "DEBUG";
    };
  };

  config = {
    systemd.extraConfig = ''
      DefaultLimitNOFILE=${toString limitNOFILE}
    '';

    networking.firewall.allowedTCPPorts = [
      hunterSettings.prometheus_port
      sieveSettings.prometheus_port
    ];

    systemd.tmpfiles.rules = [
      "L ${profanityPath} - - - - ${./profanity.txt}"
    ];

    # redis restore can be slow
    systemd.services.redis-caveman.serviceConfig.TimeoutStartSec = "infinity";

    services.redis.servers.caveman = {
      enable = true;
      port = cfg.redis.port;
      requirePassFile = cfg.redis.passwordFile;
      settings = {
        inherit (cfg.redis) maxmemory maxmemory-samples;
        maxmemory-policy = "allkeys-lru";
      };
    };
    services.postgresql = {
      enable = true;
      initialScript = pkgs.writeScript "initScript" ''
        CREATE ROLE ${dbUser} LOGIN PASSWORD '${dbPassword}';
        CREATE DATABASE caveman TEMPLATE template0 ENCODING UTF8;
        GRANT ALL PRIVILEGES ON DATABASE caveman TO ${dbUser};
      '';
      ensureUsers = [ {
        name = "collectd";
        ensurePermissions."DATABASE caveman" = "ALL PRIVILEGES";
      } ];
    };
    services.collectd.plugins.postgresql = ''
      <Query hosts_with_tokens>
        Statement "select count(distinct host) from instance_tokens;"
        <Result>
          Type gauge
          InstancePrefix "unique_tokens"
          ValuesFrom "count"
        </Result>
      </Query>
      <Query total_tokens>
        Statement "select count(*) from instance_tokens;"
        <Result>
          Type gauge
          InstancePrefix "total_tokens"
          ValuesFrom "count"
        </Result>
      </Query>

      <Database ${config.networking.hostName}>
        Param database "caveman"
        Query total_tokens
        Query hosts_with_tokens
      </Database>
    '';

    systemd.services.caveman-hunter = lib.mkIf cfg.hunter.enable {
      wantedBy = [ "multi-user.target" ];
      requires = [
        "redis-caveman.service"
        "blocklist-update.service"
      ];
      after = [
        "redis-caveman.service"
        "postgresql.service"
        "network-online.target"
      ];
      environment.RUST_LOG = "caveman=${cfg.hunter.logLevel}";
      serviceConfig = {
        ExecStart = "${pkgs.caveman-hunter}/bin/caveman-hunter ${hunterConfigFile}";
        Type = "notify";
        WatchdogSec = 600;
        Restart = "always";
        RestartSec = 30;
        DynamicUser = true;
        User = "caveman-hunter";
        ProtectSystem = "strict";
        ProtectHome = true;
        ProtectHostname = true;
        ProtectKernelLogs = true;
        ProtectKernelModules = true;
        ProtectKernelTunables = true;
        RestrictNamespaces = true;
        RestrictRealtime = true;
        LockPersonality = true;
        MemoryDenyWriteExecute = true;
        LimitNOFILE = limitNOFILE;
        LimitRSS = "4G";
        MemoryMax = "16G";
      };
    };

    systemd.services.caveman-butcher = lib.mkIf cfg.butcher.enable {
      wantedBy = [ "multi-user.target" ];
      requires = [ "redis-caveman.service" ];
      after = [ "redis-caveman.service" "network-online.target" ];
      environment.RUST_LOG = "caveman=${cfg.butcher.logLevel}";
      serviceConfig = {
        ExecStart = "${pkgs.caveman-butcher}/bin/caveman-butcher ${butcherConfigFile}";
        Type = "notify";
        WatchdogSec = 600;
        Restart = "always";
        RestartSec = 30;
        DynamicUser = true;
        User = "caveman-butcher";
        ProtectSystem = "strict";
        ProtectHome = true;
        ProtectHostname = true;
        ProtectKernelLogs = true;
        ProtectKernelModules = true;
        ProtectKernelTunables = true;
        RestrictNamespaces = true;
        RestrictRealtime = true;
        LockPersonality = true;
        MemoryDenyWriteExecute = true;
        LimitNOFILE = limitNOFILE;
        MemoryMax = "2G";
      };
    };

    systemd.services.caveman-gatherer = lib.mkIf cfg.gatherer.enable {
      wantedBy = [ "multi-user.target" ];
      requires = [ "redis-caveman.service" ];
      after = [ "redis-caveman.service" "postgresql.service" "network-online.target" ];
      environment.RUST_LOG = "caveman=${cfg.gatherer.logLevel}";
      serviceConfig = {
        ExecStart = "${pkgs.caveman-gatherer}/bin/caveman-gatherer ${gathererConfigFile}";
        Type = "notify";
        WatchdogSec = 90;
        Restart = "always";
        RestartSec = 1;
        DynamicUser = true;
        User = "caveman-gatherer";
        ProtectSystem = "strict";
        ProtectHome = true;
        ProtectHostname = true;
        ProtectKernelLogs = true;
        ProtectKernelModules = true;
        ProtectKernelTunables = true;
        RestrictNamespaces = true;
        RestrictRealtime = true;
        LockPersonality = true;
        MemoryDenyWriteExecute = true;
        LimitNOFILE = limitNOFILE;
        WorkingDirectory = "${pkgs.caveman-gatherer}/share/caveman/gatherer";
        MemoryMax = "1G";
      };
    };

    systemd.timers.caveman-gatherer-probe = lib.mkIf cfg.gatherer.enable {
      wantedBy = [ "timers.target" ];
      timerConfig.OnCalendar = "minutely";
    };
    systemd.services.caveman-gatherer-probe = lib.mkIf cfg.gatherer.enable {
      requires = [ "caveman-gatherer.service" ];
      serviceConfig = {
        Type = "oneshot";
        User = "caveman-gatherer-probe";
        DynamicUser = true;
        ProtectSystem = "full";
      };
      path = with pkgs; [ wget ];
      script = ''
        wget -O /dev/null --user-agent=caveman-gatherer-probe 127.0.0.1:${toString gathererSettings.listen_port}/
      '';
    };

    systemd.services.caveman-sieve = lib.mkIf cfg.sieve.enable {
      wantedBy = [ "multi-user.target" ];
      requires = [ "redis-caveman.service" ];
      after = [ "redis-caveman.service" "network-online.target" ];
      environment.RUST_LOG = "caveman=${cfg.sieve.logLevel}";
      serviceConfig = {
        ExecStart = "${pkgs.caveman-sieve}/bin/caveman-sieve ${sieveConfigFile}";
        Type = "notify";
        WatchdogSec = 300;
        Restart = "always";
        RestartSec = 1;
        DynamicUser = true;
        User = "caveman-sieve";
        ProtectSystem = "strict";
        ProtectHome = true;
        ProtectHostname = true;
        ProtectKernelLogs = true;
        ProtectKernelModules = true;
        ProtectKernelTunables = true;
        RestrictNamespaces = true;
        RestrictRealtime = true;
        LockPersonality = true;
        MemoryDenyWriteExecute = true;
        LimitNOFILE = limitNOFILE;
        LimitRSS = "128M:256M";
      };
    };

    systemd.services.caveman-smokestack = lib.mkIf cfg.smokestack.enable {
      wantedBy = [ "multi-user.target" ];
      requires = [ "redis-caveman.service" "caveman-hunter.service" ];
      after = [ "redis-caveman.service" "caveman-hunter.service" "network-online.target" ];
      environment.RUST_LOG = "caveman=${cfg.smokestack.logLevel}";
      serviceConfig = {
        ExecStart = "${pkgs.caveman-smokestack}/bin/caveman-smokestack ${smokestackConfigFile}";
        Type = "notify";
        WatchdogSec = 10;
        Restart = "always";
        RestartSec = 10;
        DynamicUser = true;
        User = "caveman-smokestack";
        ProtectSystem = "strict";
        ProtectHome = true;
        ProtectHostname = true;
        ProtectKernelLogs = true;
        ProtectKernelModules = true;
        ProtectKernelTunables = true;
        RestrictNamespaces = true;
        RestrictRealtime = true;
        LockPersonality = true;
        MemoryDenyWriteExecute = true;
        LimitNOFILE = limitNOFILE;
        LimitRSS = "64M:256M";
        # Allow listening on ports <1024
        AmbientCapabilities = "CAP_NET_BIND_SERVICE";
      };
    };

    systemd.services.blocklist-update = lib.mkIf cfg.hunter.enable {
      after = [ "network.target" "network-online.target" ];
      path = with pkgs; [ coreutils wget ];
      script = ''
        T=$(mktemp blocklistXXXX)
        wget -O $T https://raw.githubusercontent.com/gardenfence/blocklist/main/gardenfence.txt
        chmod a+r $T
        mv $T ${lib.escapeShellArg blocklistPath}
      '';
      serviceConfig = {
        Type = "oneshot";
        Restart = "on-failure";
        RestartSec = 600;
      };
    };
    systemd.timers.blocklist-update = lib.mkIf cfg.hunter.enable {
      wantedBy = [ "timers.target" ];
      timerConfig.OnCalendar = "hourly";
    };
  };
}