From 384863fcdfef3bfdc6b58d8ca35c1fc807cba82f Mon Sep 17 00:00:00 2001
From: Evan Fiordeliso <evan.fiordeliso@gmail.com>
Date: Mon, 13 Nov 2023 17:00:49 -0500
Subject: [PATCH] Add zsh implementation

---
 command.go         |  4 ++++
 completions_zsh.go | 59 ++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/command.go b/command.go
index e82b256..c42a169 100644
--- a/command.go
+++ b/command.go
@@ -46,6 +46,10 @@ func (c *Command) Root() *Command {
 	return c.Parent.Root()
 }
 
+func (c *Command) IsRoot() bool {
+	return c.isRoot
+}
+
 func (c *Command) CommandPath() string {
 	if c.Parent == nil {
 		return ""
diff --git a/completions_zsh.go b/completions_zsh.go
index 885ec41..f157e64 100644
--- a/completions_zsh.go
+++ b/completions_zsh.go
@@ -9,11 +9,62 @@ import (
 
 func WriteZshCompletions(out io.Writer, rootCmd *Command) error {
 	return zshTpl.Execute(out, map[string]any{
-		"rootCmd":    rootCmd,
-		"globalOpts": opt.Globals(),
+		"RootCmd":    rootCmd,
+		"GlobalOpts": opt.Globals(),
 	})
 }
 
-var zshTpl = template.Must(template.New("zsh").Parse(`
-	
+// TODO: Fix indentation and other spacing
+var zshTpl = template.Must(template.New("zsh").Funcs(tplFuncs).Parse(`
+{{- $rootCmd := .RootCmd -}}
+{{- $progName := $rootCmd.Name -}}
+{{- $varName := under $rootCmd.Name -}}
+
+{{- define "cmd" }}
+{{- $varPrefix := varPrefix .Cmd.CommandPath -}}
+{{- if not .Cmd.IsRoot}}
+{{ .Cmd.Name }})
+{{ end -}}
+  {{ if gt (len .Cmd.Subcommands) 0 -}}
+  local -a {{ $varPrefix }}commands
+  {{ $varPrefix }}commands=(
+    {{- range .Cmd.Subcommands}}
+    '{{ .Name }}:{{ .ShortDescription }}'
+    {{- end }}
+  )
+  {{ end }}
+
+  {{- if or (gt (len .Cmd.Opts) 0) (gt (len .Cmd.Subcommands) 0) -}}
+  _arguments
+  {{- end -}}
+  {{- range .Cmd.Opts -}}
+    {{ " " }}\
+    '-{{ .ShortName }}[{{ .Description }}]'
+  {{- end }}
+  {{- if gt (len .Cmd.Subcommands) 0 -}}
+    {{ " " }}\
+    '1: :{_describe 'command' commands}' \
+    '*:: :->args'
+  
+  case $state in
+    args)
+      case $words[1] in
+        {{- range .Cmd.Subcommands }}
+          {{ template "cmd" (map "Cmd" .) }}
+        {{- end }}
+      esac
+      ;;
+  esac
+  {{- end }}
+{{- if not .Cmd.IsRoot}}
+  ;;
+{{ end }}
+{{ end -}}
+
+
+function _{{ $varName }} {
+  {{ template "cmd" (map "Cmd" .RootCmd) -}}
+}
+
+compdef _{{ $varName }} {{ $progName }}
 `))