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

  mainScript =
    let
      colmena = "${pkgs.colmena}/bin/colmena";
      date = "${pkgs.coreutils}/bin/date";
      readlink = "${pkgs.coreutils}/bin/readlink";
      shutdown = "${config.systemd.package}/bin/shutdown";
    in
    if cfg.allowReboot then
      ''
        ${colmena} apply-local boot
        booted="$(${readlink} /run/booted-system/{initrd,kernel,kernel-modules})"
        built="$(${readlink} /nix/var/nix/profiles/system/{initrd,kernel,kernel-modules})"

        ${lib.optionalString (cfg.rebootWindow != null) ''
          current_time="$(${date} +%H:%M)"

          lower="${cfg.rebootWindow.lower}"
          upper="${cfg.rebootWindow.upper}"

          if [[ "''${lower}" < "''${upper}" ]]; then
            if [[ "''${current_time}" > "''${lower}" ]] && \
                [[ "''${current_time}" < "''${upper}" ]]; then
              do_reboot="true"
            else
              do_reboot="false"
            fi
          else
            # lower > upper, so we are crossing midnight (e.g. lower=23h, upper=6h)
            # we want to reboot if cur > 23h or cur < 6h
            if [[ "''${current_time}" < "''${upper}" ]] || \
            [[ "''${current_time}" > "''${lower}" ]]; then
              do_reboot="true"
            else
              do_reboot="false"
            fi
          fi
        ''}

        if [ "''${booted}" = "''${built}" ]; then
          ${colmena} apply-local switch
          ${lib.optionalString (cfg.rebootWindow != null) ''
        elif [ "''${do_reboot}" != true ]; then
          echo "Outside of configured reboot window, skipping."
        ''}
        else
          ${shutdown} -r +1
        fi
      ''
    else
      ''
        ${colmena} apply-local switch
      ''
  ;
in
{
  options.services.colmenaAutoUpgrade = {
    enable = lib.mkEnableOption {
      default = false;
      description = "Enable automatic upgrades for Colmena";
    };

    git = lib.mkOption {
      type = lib.types.submodule {
        options = {
          enable = lib.mkEnableOption "Whether to pull the latest changes from the Git repository before upgrading.";

          branch = lib.mkOption {
            type = lib.types.str;
            default = "origin/main";
            description = "Git branch to checkout after fetching";
          };
        };
      };
    };

    useNixShell = lib.mkOption {
      default = false;
      type = lib.types.bool;
      description = ''
        Whether to run colmena in a nix-shell.
      '';
    };

    dates = lib.mkOption {
      type = lib.types.str;
      default = "04:40";
      example = "daily";
      description = ''
        How often or when upgrade occurs. For most desktop and server systems
        a sufficient upgrade frequency is once a day.

        The format is described in
        {manpage}`systemd.time(7)`.
      '';
    };

    allowReboot = lib.mkOption {
      default = false;
      type = lib.types.bool;
      description = ''
        Reboot the system into the new generation instead of a switch
        if the new generation uses a different kernel, kernel modules
        or initrd than the booted system.
        See {option}`rebootWindow` for configuring the times at which a reboot is allowed.
      '';
    };

    randomizedDelaySec = lib.mkOption {
      default = "0";
      type = lib.types.str;
      example = "45min";
      description = ''
        Add a randomized delay before each automatic upgrade.
        The delay will be chosen between zero and this value.
        This value must be a time span in the format specified by
        {manpage}`systemd.time(7)`
      '';
    };

    fixedRandomDelay = lib.mkOption {
      default = false;
      type = lib.types.bool;
      example = true;
      description = ''
        Make the randomized delay consistent between runs.
        This reduces the jitter between automatic upgrades.
        See {option}`randomizedDelaySec` for configuring the randomized delay.
      '';
    };

    rebootWindow = lib.mkOption {
      description = ''
        Define a lower and upper time value (in HH:MM format) which
        constitute a time window during which reboots are allowed after an upgrade.
        This option only has an effect when {option}`allowReboot` is enabled.
        The default value of `null` means that reboots are allowed at any time.
      '';
      default = null;
      example = {
        lower = "01:00";
        upper = "05:00";
      };
      type =
        with lib.types;
        nullOr (submodule {
          options = {
            lower = lib.mkOption {
              description = "Lower limit of the reboot window";
              type = lib.types.strMatching "[[:digit:]]{2}:[[:digit:]]{2}";
              example = "01:00";
            };

            upper = lib.mkOption {
              description = "Upper limit of the reboot window";
              type = lib.types.strMatching "[[:digit:]]{2}:[[:digit:]]{2}";
              example = "05:00";
            };
          };
        });
    };

    persistent = lib.mkOption {
      default = true;
      type = lib.types.bool;
      example = false;
      description = ''
        Takes a boolean argument. If true, the time when the service
        unit was last triggered is stored on disk. When the timer is
        activated, the service unit is triggered immediately if it
        would have been triggered at least once during the time when
        the timer was inactive. Such triggering is nonetheless
        subject to the delay imposed by RandomizedDelaySec=. This is
        useful to catch up on missed runs of the service when the
        system was powered down.
      '';
    };
  };

  config = lib.mkIf cfg.enable {
    systemd.services.colmena-auto-upgrade = {
      description = "Upgrade nixos with colmena";

      restartIfChanged = false;
      unitConfig.X-StopOnRemoval = false;

      serviceConfig.Type = "oneshot";

      environment =
        config.nix.envVars
        // {
          inherit (config.environment.sessionVariables) NIX_PATH;
          HOME = "/root";
        }
        // config.networking.proxy.envVars;

      path = with pkgs; [
        coreutils
        gnutar
        xz.bin
        gzip
        gitMinimal
        colmena
        config.nix.package.out
        config.programs.ssh.package
      ];

      serviceConfig.WorkingDirectory = "/etc/nixos";
      script =
        let
          git = "${pkgs.gitMinimal}/bin/git";
          nix-shell = "${pkgs.nix}/bin/nix-shell";
        in
        ''
          ${lib.optionalString cfg.git.enable
            ''
            ${git} fetch --prune
            ${git} checkout ${cfg.git.branch}
            ''
          }
          ${if cfg.useNixShell then ''
            ${nix-shell} --run "${pkgs.writeShellScript "colmena-auto-upgrade" mainScript}"
          ''
          else mainScript
          }
        '';
      startAt = cfg.dates;
      after = [ "network-online.target" ];
      wants = [ "network-online.target" ];
    };

    systemd.timers.colmena-auto-upgrade = {
      timerConfig = {
        RandomizedDelaySec = cfg.randomizedDelaySec;
        FixedRandomDelay = cfg.fixedRandomDelay;
        Persistent = cfg.persistent;
      };
    };
  };
}