From e757e1b8a9b4bd26bade4fe6552fe4e9d9382137 Mon Sep 17 00:00:00 2001
From: Mohammad Alsaleh <CE.Mohammad.AlSaleh@gmail.com>
Date: Wed, 16 Sep 2015 19:02:35 +0300
Subject: [PATCH] Add zsh support in save/restore functions

Signed-off-by: Mohammad Alsaleh <CE.Mohammad.AlSaleh@gmail.com>
---
 scripts/helpers.sh |  3 ++-
 scripts/restore.sh | 12 +++++++++---
 scripts/save.sh    | 27 ++++++++++++++++++++++++---
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/scripts/helpers.sh b/scripts/helpers.sh
index 555b779..921902b 100644
--- a/scripts/helpers.sh
+++ b/scripts/helpers.sh
@@ -141,5 +141,6 @@ pane_contents_archive_file() {
 
 resurrect_history_file() {
 	local pane_id="$1"
-	echo "$(resurrect_dir)/bash_history-${pane_id}"
+	local shell_name="$2"
+	echo "$(resurrect_dir)/${shell_name}_history-${pane_id}"
 }
diff --git a/scripts/restore.sh b/scripts/restore.sh
index 8e9c550..458d71e 100755
--- a/scripts/restore.sh
+++ b/scripts/restore.sh
@@ -274,10 +274,16 @@ restore_shell_history() {
 	awk 'BEGIN { FS="\t"; OFS="\t" } /^pane/ { print $2, $3, $7, $10; }' $(last_resurrect_file) |
 		while IFS=$d read session_name window_number pane_index pane_command; do
 			if ! is_pane_registered_as_existing "$session_name" "$window_number" "$pane_index"; then
-				if [ "$pane_command" == "bash" ]; then
-					local pane_id="$session_name:$window_number.$pane_index"
-					local read_command="history -r '$(resurrect_history_file "$pane_id")'"
+				local pane_id="$session_name:$window_number.$pane_index"
+				local history_file="$(resurrect_history_file "$pane_id" "$pane_command")"
+
+				if [ "$pane_command" = "bash" ]; then
+					local read_command="history -r '$history_file'"
 					tmux send-keys -t "$pane_id" "$read_command" C-m
+				elif [ "$pane_command" = "zsh" ]; then
+					local accept_line="$(expr "$(zsh -i -c bindkey | grep -m1 '\saccept-line$')" : '^"\(.*\)".*')"
+					local read_command="fc -R '$history_file'; clear"
+					tmux send-keys -t "$pane_id" "$read_command" "$accept_line"
 				fi
 			fi
 		done
diff --git a/scripts/save.sh b/scripts/save.sh
index 97ccd8b..f6615b8 100755
--- a/scripts/save.sh
+++ b/scripts/save.sh
@@ -143,16 +143,37 @@ capture_pane_contents() {
 }
 
 save_shell_history() {
+	if [ "$pane_command" = "bash" ]; then
+		local history_w='history -w'
+		local accept_line='C-m'
+		local end_of_line='C-e'
+		local backward_kill_line='C-u'
+	elif [ "$pane_command" = "zsh" ]; then
+		# fc -W does not work with -L
+		# fc -l format is different from what's written by fc -W
+		# fc -R either reads the format produced by fc -W or considers
+		# the entire line to be a command. That's why we need -n.
+		# fc -l only list the last 16 items by default, I think 64 is more reasonable.
+		local history_w='fc -lLn -64 >'
+
+		local zsh_bindkey="$(zsh -i -c bindkey)"
+		local accept_line="$(expr "$(echo "$zsh_bindkey" | grep -m1 '\saccept-line$')" : '^"\(.*\)".*')"
+		local end_of_line="$(expr "$(echo "$zsh_bindkey" | grep -m1 '\send-of-line$')" : '^"\(.*\)".*')"
+		local backward_kill_line="$(expr "$(echo "$zsh_bindkey" | grep -m1 '\sbackward-kill-line$')" : '^"\(.*\)".*')"
+	else
+		return
+	fi
+
 	local pane_id="$1"
 	local pane_command="$2"
 	local full_command="$3"
-	if [ "$pane_command" == "bash" ] && [ "$full_command" == ":" ]; then
+	if [ "$full_command" = ":" ]; then
 		# leading space prevents the command from being saved to history
 		# (assuming default HISTCONTROL settings)
-		local write_command=" history -w '$(resurrect_history_file "$pane_id")'"
+		local write_command=" $history_w '$(resurrect_history_file "$pane_id" "$pane_command")'"
 		# C-e C-u is a Bash shortcut sequence to clear whole line. It is necessary to
 		# delete any pending input so it does not interfere with our history command.
-		tmux send-keys -t "$pane_id" C-e C-u "$write_command" C-m
+		tmux send-keys -t "$pane_id" "$end_of_line" "$backward_kill_line" "$write_command" "$accept_line"
 	fi
 }