completions.go

// See: https://zsh.sourceforge.io/Doc/Release/Completion-System.html
// See: https://github.com/zsh-users/zsh-completions

package cli

import (
	"io"
)

type CompletionGenerator interface {
	CompletionGenerate(cmd *Command, w io.Writer) error
}

type CompletionDynamicValuer interface {
}

var _ CompletionGenerator = (*ZSHCompletionGenerator)(nil)

type ZSHCompletionGenerator struct{}

func (g *ZSHCompletionGenerator) CompletionGenerate(cmd *Command, w io.Writer) error {
	ew := &easyWriter{w: w}

	path := cmd.Path()

	// Header.
	ew.Writef("#compdef _%[1]s %[1]s\n", cmd.Name)
	ew.Writef("\n")
	ew.Writef("#\n")
	ew.Writef("# Code generated by %s; DO NOT EDIT.\n", cmd.Name)
	ew.Writef("#\n")
	ew.Writef("\n")
	ew.Writef("# zsh completion for %s\n", cmd.Name)
	ew.Writef("\n")

	if err := ew.Err(); err != nil {
		return err
	}

	// Main.
	ew.Writef("_%s() {\n", cmd.Name)

	// Main command.
	if err := g.generateCommand(cmd, ew); err != nil {
		return err
	}

	// Create buffer for path to prevent multiple allocations.
	const minPathBuffer = 8
	pathBuffer := make([]string, max(len(path), minPathBuffer))
	copy(pathBuffer, path)

	// Generate subcommands.
	err := walkCommands(cmd, &pathBuffer, func(cmd *Command) error {
		ew.Writef("\n")

		if err := g.generateCommand(cmd, ew); err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return err
	}

	// Call main command.
	ew.Writef("\n")
	ew.Writef("    {\n")
	ew.Writef("        ")
	if err := g.functionName(path, ew); err != nil {
		return err
	}
	ew.Writef("\n")
	ew.Writef("    } always {\n")
	ew.Writef("        unfunction ")
	if err := g.functionName(cmd.Path(), ew); err != nil {
		return err
	}
	ew.Writef("\n")

	err = walkCommands(cmd, &pathBuffer, func(cmd *Command) error {
		ew.Writef("        unfunction ")
		if err := g.functionName(cmd.Path(), ew); err != nil {
			return err
		}
		ew.Writef("\n")

		return nil
	})
	if err != nil {
		return err
	}
	ew.Writef("    }\n")

	ew.Writef("}\n")

	if err := ew.Err(); err != nil {
		return err
	}

	return nil
}

func (g *ZSHCompletionGenerator) generateCommand(cmd *Command, ew *easyWriter) error {
	app := cmd.App()
	path := cmd.Path()
	flags := cmd.Flags()
	args := cmd.Args()
	rest := cmd.Rest()

	ew.Writef("    #")
	for _, p := range path {
		ew.Writef(" %s", p)
	}
	ew.Writef("\n")

	ew.Writef("    ")
	if err := g.functionName(path, ew); err != nil {
		return err
	}
	ew.Writef("() {\n")

	// Empty command.
	if len(flags) == 0 && len(args) == 0 && rest == nil && len(cmd.Commands) == 0 {
		ew.Writef("        :")
		ew.Writef("    }\n")

		if err := ew.Err(); err != nil {
			return err
		}

		return nil
	}

	// Init.
	// ew.Writef("        local -a")
	//
	// if len(flags) > 0 {
	// 	ew.Writef(" options")
	// }
	//
	// if len(args) > 0 {
	// 	ew.Writef(" args")
	// }
	//
	// if rest != nil {
	// 	ew.Writef(" rest")
	// }
	//
	// if len(cmd.Commands) > 0 {
	// 	ew.Writef(" commands")
	// }
	//
	// ew.Writef("\n")

	ew.Writef("        local -a options args rest commands\n")

	if len(cmd.Commands) > 0 {
		ew.Writef("        local subcmd\n")
	}

	if err := ew.Err(); err != nil {
		return err
	}

	// Options.
	if len(flags) > 0 {
		ew.Writef("\n")
		ew.Writef("        options=(\n")
		for i := range flags {
			f := &flags[i]

			if f.Short != "" || f.Long != "" {
				ew.Writef("            ")
				if err := g.generateFlagDef(cmd, f, len(args), ew); err != nil {
					return err
				}
				ew.Writef("\n")
			}
		}
		ew.Writef("        )\n")

		if err := ew.Err(); err != nil {
			return err
		}
	}

	// Arguments.
	if len(args) > 0 {
		ew.Writef("\n")

		ew.Writef("        args=(\n")
		for i := range args {
			arg := &args[i]

			ew.Writef("            ")
			if err := g.generateArgDef(i, arg, ew); err != nil {
				return err
			}
			ew.Writef("\n")
		}
		ew.Writef("        )\n")

		if err := ew.Err(); err != nil {
			return err
		}
	}

	// Rest.
	if rest != nil {
		ew.Writef("\n")
		ew.Writef("        rest=(\n")
		ew.Writef("            ")
		if err := g.generateRestDef(rest, ew); err != nil {
			return err
		}
		ew.Writef("\n")
		ew.Writef("        )\n")
	}

	// Commands.
	if len(cmd.Commands) > 0 {
		ew.Writef("\n")
		ew.Writef("        commands=(\n")
		for i := range cmd.Commands {
			// TODO(SuperPaintman): optimize path.
			newPath := append(append([]string{}, path...), cmd.Commands[i].Name)

			subCmd, err := app.Command(newPath...)
			if err != nil {
				return err
			}

			ew.Writef("            ")
			if err := g.generateCommandDef(subCmd, ew); err != nil {
				return err
			}
			ew.Writef("\n")
		}
		ew.Writef("        )\n")

		if err := ew.Err(); err != nil {
			return err
		}
	}

	// Arguments and subcommands.
	ew.Writef("\n")
	if len(cmd.Commands) == 0 {
		ew.Writef("        _arguments $options $rest $args\n")
	} else {
		// Detect mode.
		ew.Writef("        if [ \"$CURRENT\" = 2 ]; then\n")
		ew.Writef("            _arguments \\\n")
		ew.Writef("                '1:: :{_describe command commands}'\n")
		if len(flags) > 0 || len(args) > 0 || rest != nil {
			ew.Writef("\n")
			ew.Writef("            _arguments $options $args $rest\n")
		}
		ew.Writef("        else\n")
		ew.Writef("            case \"${words[2]}\" in\n")
		ew.Writef("                '%s'", cmd.Commands[0].Name)

		for i := range cmd.Commands[1:] {
			ew.Writef(" | '%s'", cmd.Commands[i+1].Name)
		}
		ew.Writef(")\n")
		ew.Writef("                    subcmd=\"${words[2]}\"\n")
		ew.Writef("\n")
		ew.Writef("                    _arguments \\\n")
		ew.Writef("                        '1: :{_describe command commands}' \\\n")
		ew.Writef("                        '*:: :->args'\n")
		ew.Writef("                    ;;\n")
		if len(flags) > 0 || len(args) > 0 || rest != nil {
			ew.Writef("\n")
			ew.Writef("                *)\n")
			ew.Writef("                    _arguments $options $args $rest\n")
			ew.Writef("                    ;;\n")
		}
		ew.Writef("            esac\n")
		ew.Writef("        fi\n")

		ew.Writef("\n")
		ew.Writef("        case \"$subcmd\" in")
		for i := range cmd.Commands {
			subCmd := &cmd.Commands[i]

			ew.Writef("\n")
			ew.Writef("            '%s')\n", subCmd.Name)
			ew.Writef("                ")
			if err := g.functionName(subCmd.Path(), ew); err != nil {
				return err
			}
			ew.Writef("\n")
			ew.Writef("                ;;\n")
		}
		ew.Writef("        esac\n")
	}

	ew.Writef("    }\n")

	if err := ew.Err(); err != nil {
		return err
	}

	return nil
}

func (g *ZSHCompletionGenerator) generateFlagDef(cmd *Command, f *Flag, argsCount int, ew *easyWriter) error {
	// Omit flags and args.
	if f.commandFlag { // TODO
		ew.Writef("'(* -")

		for i := 0; i < argsCount; i++ {
			ew.Writef(" %d", i+1)
		}
		ew.Writef(")'")
	} else if fv, ok := f.Value.(boolFlag); ok && fv.IsBoolFlag() {
		ew.Writef("'(")
		if f.Short != "" {
			ew.Writef(cmd.Parser().FormatShortFlag(f.Short))
		}
		if f.Long != "" {
			if f.Long != "" {
				ew.Writef(" ")
			}
			ew.Writef(cmd.Parser().FormatLongFlag(f.Long))
		}
		ew.Writef(")'")
	}

	// Name.
	if f.Short != "" {
		if f.Long != "" {
			ew.Writef("{")
			ew.Writef(cmd.Parser().FormatShortFlag(f.Short))
			ew.Writef(",")
			ew.Writef(cmd.Parser().FormatLongFlag(f.Long))
			ew.Writef("}")
		} else {
			ew.Writef(cmd.Parser().FormatShortFlag(f.Short))
		}
	} else {
		ew.Writef(cmd.Parser().FormatLongFlag(f.Long))
	}

	// Value.
	if _, ok := f.Value.(boolFlag); !ok {
		ew.Writef("'='")
	}

	// Usage.
	if f.Usage != nil {
		ew.Writef("'[")
		if err := f.Usage.Usage(cmd, &zshFlagSanitizer{ew}); err != nil {
			return err
		}
		ew.Writef("]'")
	}

	// Value.
	if _, ok := f.Value.(boolFlag); !ok {
		ew.Writef("': :()'")
	} else {
		// TODO
	}

	if err := ew.Err(); err != nil {
		return err
	}

	return nil
}

func (g *ZSHCompletionGenerator) generateArgDef(i int, a *Arg, ew *easyWriter) error {
	ew.Writef("'%d:", i+1)

	if !a.Required() {
		ew.Writef(":")
	}

	if a.Name != "" {
		ew.Writef(a.Name)
	} else {
		ew.Writef(" ")
	}

	// Value.
	ew.Writef(":()") // TODO

	ew.Writef("'")

	if err := ew.Err(); err != nil {
		return err
	}

	return nil
}

func (g *ZSHCompletionGenerator) generateRestDef(ra *RestArgs, ew *easyWriter) error {
	ew.Writef("'*::")

	if ra.Name != "" {
		ew.Writef(ra.Name)
	} else {
		ew.Writef(" ")
	}

	// Value.
	ew.Writef(":()") // TODO

	ew.Writef("'")

	if err := ew.Err(); err != nil {
		return err
	}

	return nil
}

func (g *ZSHCompletionGenerator) generateCommandDef(subCmd *Command, ew *easyWriter) error {
	ew.Writef("'")

	ew.Writef(subCmd.Name)

	if subCmd.Usage != nil {
		ew.Writef(":")

		if err := subCmd.Usage.Usage(subCmd, &zshFlagSanitizer{ew}); err != nil {
			return err
		}
	}

	ew.Writef("'")

	if err := ew.Err(); err != nil {
		return err
	}

	return nil
}

func (*ZSHCompletionGenerator) functionName(path []string, ew *easyWriter) error {
	if len(path) == 0 {
		return nil
	}

	for _, p := range path {
		ew.Writef("__%s", p)
	}

	if err := ew.Err(); err != nil {
		return err
	}

	return nil
}

type zshFlagSanitizer struct {
	ew *easyWriter
}

func (s *zshFlagSanitizer) Write(p []byte) (n int, err error) {
	var extra int
	for _, b := range p {
		switch b {
		case '[', ']':
			extra += len(`\`)
		case '\'':
			extra += len(`'"'"`)
		}
	}

	if extra == 0 {
		return s.ew.w.Write(p)
	}

	data := make([]byte, 0, len(p)+extra)
	for _, b := range p {
		switch b {
		case '[', ']':
			data = append(data, '\\')
		case '\'':
			data = append(data, '\'', '"', '\'', '"')
		}

		data = append(data, b)
	}

	return s.ew.w.Write(data)
}

func walkCommands(cmd *Command, pathBuffer *[]string, fn func(*Command) error) error {
	app := cmd.App()
	path := cmd.Path()

	for i := range cmd.Commands {
		*pathBuffer = (*pathBuffer)[:len(path)]
		*pathBuffer = append(*pathBuffer, cmd.Commands[i].Name)

		// We can't use a command directly because it might be not initialized and
		// setuped.
		//
		// NOTE(SuperPaintman):
		//     Maybe we need a method for getting subcommands without calling the
		//     app.
		subCmd, err := app.Command(*pathBuffer...)
		if err != nil {
			return err
		}

		if err := fn(subCmd); err != nil {
			return err
		}

		if len(subCmd.Commands) != 0 {
			if err := walkCommands(subCmd, pathBuffer, fn); err != nil {
				return err
			}
		}
	}

	return nil
}

func max(a, b int) int {
	if a > b {
		return a
	}

	return b
}