From 0394c71f2bc00be1e8e76d2931549721c710904c 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                   | 69 +++++++++++++++----
 .../modules/programs/zellij/enable-shells.nix | 17 ++++-
 2 files changed, 73 insertions(+), 13 deletions(-)

diff --git a/modules/programs/zellij.nix b/modules/programs/zellij.nix
index 3d5c0dd52..6aec07200 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.
       '';
     };
 
@@ -44,6 +65,15 @@ in {
       '';
     };
 
+    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>.
+      '';
+    };
+
     enableBashIntegration =
       lib.hm.shell.mkBashIntegrationOption { inherit config; };
 
@@ -69,17 +99,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 (mkOrder 200 ''
+        eval "$(${zellijCmd} setup --generate-auto-start zsh)"
+      '') else
+        "");
 
-    programs.zsh.initContent = 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";
+    };
   };
 }
diff --git a/tests/modules/programs/zellij/enable-shells.nix b/tests/modules/programs/zellij/enable-shells.nix
index 63e13facc..470075b28 100644
--- a/tests/modules/programs/zellij/enable-shells.nix
+++ b/tests/modules/programs/zellij/enable-shells.nix
@@ -4,9 +4,16 @@
   programs = {
     zellij = {
       enable = true;
-      enableBashIntegration = true;
+
+      autostartOnShellStart = {
+        enable = true;
+        attachExistingSession = true;
+        exitShellOnExit = true;
+      };
+
       enableZshIntegration = true;
       enableFishIntegration = true;
+      enableBashIntegration = true;
     };
     bash.enable = true;
     zsh.enable = true;
@@ -32,5 +39,13 @@
     assertFileContains \
       home-files/.config/fish/config.fish \
       'eval (@zellij@/bin/zellij setup --generate-auto-start fish | string collect)'
+
+    assertFileExists home-path/etc/profile.d/hm-session-vars.sh
+    assertFileContains \
+      home-path/etc/profile.d/hm-session-vars.sh \
+      'export ZELLIJ_AUTO_ATTACH="true"'
+    assertFileContains \
+      home-path/etc/profile.d/hm-session-vars.sh \
+      'export ZELLIJ_AUTO_EXIT="true"'
   '';
 }