From b87b657770c0d780523b553a32d3aefc4c5012fb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20Chocholat=C3=BD?= <chocholaty.david@protonmail.com>
Date: Sun, 3 Nov 2024 18:20:23 +0100
Subject: [PATCH] zellij: Add additional options for integrating with shells

---
 modules/programs/zellij.nix | 71 ++++++++++++++++++++++++++++++-------
 1 file changed, 59 insertions(+), 12 deletions(-)

diff --git a/modules/programs/zellij.nix b/modules/programs/zellij.nix
index 0906a028d..4145b3714 100644
--- a/modules/programs/zellij.nix
+++ b/modules/programs/zellij.nix
@@ -8,18 +8,39 @@ let
   yamlFormat = pkgs.formats.yaml { };
   zellijCmd = getExe cfg.package;
 
+  autostartOnShellStartModule = types.submodule {
+    options = {
+      enable = mkEnableOption "" // {
+        description = ''
+          Whether to autostart Zellij session on shell creation.
+        '';
+      };
+
+      attachExistingSession = mkEnableOption "" // {
+        description = ''
+          Whether to attach to the default session after being autostarted if a Zellij session already exists.
+        '';
+      };
+
+      exitShellOnExit = mkEnableOption "" // {
+        description = ''
+          Whether to exit the shell when Zellij exits after being autostarted.
+        '';
+      };
+    };
+  };
 in {
   meta.maintainers = [ hm.maintainers.mainrs ];
 
   options.programs.zellij = {
-    enable = mkEnableOption "zellij";
+    enable = mkEnableOption "Zellij";
 
     package = mkOption {
       type = types.package;
       default = pkgs.zellij;
       defaultText = literalExpression "pkgs.zellij";
       description = ''
-        The zellij package to install.
+        The Zellij package to install.
       '';
     };
 
@@ -41,6 +62,17 @@ in {
       '';
     };
 
+    # Additional options.
+    autostartOnShellStart = mkOption {
+      type = types.nullOr autostartOnShellStartModule;
+      default = null;
+      description = ''
+        Options related to autostarting Zellij on shell creation.
+        Requires enable<Shell>Integration to apply to the respective <Shell>.
+      '';
+    };
+
+    # Shell integration.
     enableBashIntegration = mkEnableOption "Bash integration" // {
       default = false;
     };
@@ -69,17 +101,32 @@ in {
         text = lib.hm.generators.toKDL { } cfg.settings;
       };
 
-    programs.bash.initExtra = mkIf cfg.enableBashIntegration (mkOrder 200 ''
-      eval "$(${zellijCmd} setup --generate-auto-start bash)"
-    '');
+    programs.zsh.initExtra = mkIf (cfg.enableZshIntegration)
+      (if cfg.autostartOnShellStart.enable then ''
+        eval "$(${zellijCmd} setup --generate-auto-start zsh)"
+      '' else
+        "");
 
-    programs.zsh.initExtra = mkIf cfg.enableZshIntegration (mkOrder 200 ''
-      eval "$(${zellijCmd} setup --generate-auto-start zsh)"
-    '');
-
-    programs.fish.interactiveShellInit = mkIf cfg.enableFishIntegration
-      (mkOrder 200 ''
+    programs.fish.interactiveShellInit = mkIf (cfg.enableFishIntegration)
+      (if cfg.autostartOnShellStart.enable then ''
         eval (${zellijCmd} setup --generate-auto-start fish | string collect)
-      '');
+      '' else
+        "");
+
+    programs.bash.initExtra = mkIf (cfg.enableBashIntegration)
+      (if cfg.autostartOnShellStart.enable then ''
+        eval "$(${zellijCmd} setup --generate-auto-start bash)"
+      '' else
+        "");
+
+    home.sessionVariables = mkIf cfg.autostartOnShellStart.enable {
+      ZELLIJ_AUTO_ATTACH =
+        if cfg.autostartOnShellStart.attachExistingSession then
+          "true"
+        else
+          "false";
+      ZELLIJ_AUTO_EXIT =
+        if cfg.autostartOnShellStart.exitShellOnExit then "true" else "false";
+    };
   };
 }