"encoding/json"

	"sort"

	"text/template"

	"github.com/spf13/pflag"

const (

	zshCompArgumentAnnotation   = "cobra_annotations_zsh_completion_argument_annotation"

	zshCompArgumentFilenameComp = "cobra_annotations_zsh_completion_argument_file_completion"

	zshCompArgumentWordComp     = "cobra_annotations_zsh_completion_argument_word_completion"

	zshCompDirname              = "cobra_annotations_zsh_dirname"

)

var (

	zshCompFuncMap = template.FuncMap{

		"genZshFuncName":              zshCompGenFuncName,

		"extractFlags":                zshCompExtractFlag,

		"genFlagEntryForZshArguments": zshCompGenFlagEntryForArguments,

		"extractArgsCompletions":      zshCompExtractArgumentCompletionHintsForRendering,

	}

	zshCompletionText = `

{{/* should accept Command (that contains subcommands) as parameter */}}

{{define "argumentsC" -}}

{{ $cmdPath := genZshFuncName .}}

function {{$cmdPath}} {

  local -a commands

  _arguments -C \{{- range extractFlags .}}

    {{genFlagEntryForZshArguments .}} \{{- end}}

    "1: :->cmnds" \

    "*::arg:->args"

  case $state in

  cmnds)

    commands=({{range .Commands}}{{if not .Hidden}}

      "{{.Name}}:{{.Short}}"{{end}}{{end}}

    )

    _describe "command" commands

    ;;

  esac

  case "$words[1]" in {{- range .Commands}}{{if not .Hidden}}

  {{.Name}})

    {{$cmdPath}}_{{.Name}}

    ;;{{end}}{{end}}

  esac

}

{{range .Commands}}{{if not .Hidden}}

{{template "selectCmdTemplate" .}}

{{- end}}{{end}}

{{- end}}

{{/* should accept Command without subcommands as parameter */}}

{{define "arguments" -}}

function {{genZshFuncName .}} {

{{"  _arguments"}}{{range extractFlags .}} \

    {{genFlagEntryForZshArguments . -}}

{{end}}{{range extractArgsCompletions .}} \

    {{.}}{{end}}

}

{{end}}

{{/* dispatcher for commands with or without subcommands */}}

{{define "selectCmdTemplate" -}}

{{if .Hidden}}{{/* ignore hidden*/}}{{else -}}

{{if .Commands}}{{template "argumentsC" .}}{{else}}{{template "arguments" .}}{{end}}

{{- end}}

{{- end}}

{{/* template entry point */}}

{{define "Main" -}}

#compdef _{{.Name}} {{.Name}}

{{template "selectCmdTemplate" .}}

{{end}}

`

)

// zshCompArgsAnnotation is used to encode/decode zsh completion for

// arguments to/from Command.Annotations.

type zshCompArgsAnnotation map[int]zshCompArgHint

type zshCompArgHint struct {

	// Indicates the type of the completion to use. One of:

	// zshCompArgumentFilenameComp or zshCompArgumentWordComp

	Tipe string `json:"type"`

	// A value for the type above (globs for file completion or words)

	Options []string `json:"options"`

}

// GenZshCompletion generates a zsh completion file and writes to the passed

// writer. The completion always run on the root command regardless of the

// command it was called from.

	tmpl, err := template.New("Main").Funcs(zshCompFuncMap).Parse(zshCompletionText)

	if err != nil {

		return fmt.Errorf("error creating zsh completion template: %v", err)

	}

	return tmpl.Execute(w, c.Root())

// MarkZshCompPositionalArgumentFile marks the specified argument (first

// argument is 1) as completed by file selection. patterns (e.g. "*.txt") are

// optional - if not provided the completion will search for all files.

func (c *Command) MarkZshCompPositionalArgumentFile(argPosition int, patterns ...string) error {

	if argPosition < 1 {

		return fmt.Errorf("Invalid argument position (%d)", argPosition)

	}

	annotation, err := c.zshCompGetArgsAnnotations()

	if err != nil {

		return err

	}

	if c.zshcompArgsAnnotationnIsDuplicatePosition(annotation, argPosition) {

		return fmt.Errorf("Duplicate annotation for positional argument at index %d", argPosition)

	}

	annotation[argPosition] = zshCompArgHint{

		Tipe:    zshCompArgumentFilenameComp,

		Options: patterns,

	}

	return c.zshCompSetArgsAnnotations(annotation)

// MarkZshCompPositionalArgumentWords marks the specified positional argument

// (first argument is 1) as completed by the provided words. At east one word

// must be provided, spaces within words will be offered completion with

// "word\ word".

func (c *Command) MarkZshCompPositionalArgumentWords(argPosition int, words ...string) error {

	if argPosition < 1 {

		return fmt.Errorf("Invalid argument position (%d)", argPosition)

	}

	if len(words) == 0 {

		return fmt.Errorf("Trying to set empty word list for positional argument %d", argPosition)

	}

	annotation, err := c.zshCompGetArgsAnnotations()

	if err != nil {

		return err

	}

	if c.zshcompArgsAnnotationnIsDuplicatePosition(annotation, argPosition) {

		return fmt.Errorf("Duplicate annotation for positional argument at index %d", argPosition)

	}

	annotation[argPosition] = zshCompArgHint{

		Tipe:    zshCompArgumentWordComp,

		Options: words,

	return c.zshCompSetArgsAnnotations(annotation)

}

func zshCompExtractArgumentCompletionHintsForRendering(c *Command) ([]string, error) {

	var result []string

	annotation, err := c.zshCompGetArgsAnnotations()

	if err != nil {

		return nil, err

	}

	for k, v := range annotation {

		s, err := zshCompRenderZshCompArgHint(k, v)

		if err != nil {

			return nil, err

		}

		result = append(result, s)

	}

	if len(c.ValidArgs) > 0 {

		if _, positionOneExists := annotation[1]; !positionOneExists {

			s, err := zshCompRenderZshCompArgHint(1, zshCompArgHint{

				Tipe:    zshCompArgumentWordComp,

				Options: c.ValidArgs,

			})

			if err != nil {

				return nil, err

			}

			result = append(result, s)

	sort.Strings(result)

	return result, nil

func zshCompRenderZshCompArgHint(i int, z zshCompArgHint) (string, error) {

	switch t := z.Tipe; t {

	case zshCompArgumentFilenameComp:

		var globs []string

		for _, g := range z.Options {

			globs = append(globs, fmt.Sprintf(`-g "%s"`, g))

		}

		return fmt.Sprintf(`'%d: :_files %s'`, i, strings.Join(globs, " ")), nil

	case zshCompArgumentWordComp:

		var words []string

		for _, w := range z.Options {

			words = append(words, fmt.Sprintf("%q", w))

		}

		return fmt.Sprintf(`'%d: :(%s)'`, i, strings.Join(words, " ")), nil

	default:

		return "", fmt.Errorf("Invalid zsh argument completion annotation: %s", t)

}

func (c *Command) zshcompArgsAnnotationnIsDuplicatePosition(annotation zshCompArgsAnnotation, position int) bool {

	_, dup := annotation[position]

	return dup

func (c *Command) zshCompGetArgsAnnotations() (zshCompArgsAnnotation, error) {

	annotation := make(zshCompArgsAnnotation)

	annotationString, ok := c.Annotations[zshCompArgumentAnnotation]

	if !ok {

		return annotation, nil

	}

	err := json.Unmarshal([]byte(annotationString), &annotation)

	if err != nil {

		return annotation, fmt.Errorf("Error unmarshaling zsh argument annotation: %v", err)

	}

	return annotation, nil

}

func (c *Command) zshCompSetArgsAnnotations(annotation zshCompArgsAnnotation) error {

	jsn, err := json.Marshal(annotation)

	if err != nil {

		return fmt.Errorf("Error marshaling zsh argument annotation: %v", err)

	}

	if c.Annotations == nil {

		c.Annotations = make(map[string]string)

	}

	c.Annotations[zshCompArgumentAnnotation] = string(jsn)

	return nil

}

func zshCompGenFuncName(c *Command) string {

	if c.HasParent() {

		return zshCompGenFuncName(c.Parent()) + "_" + c.Name()

	return "_" + c.Name()

}

func zshCompExtractFlag(c *Command) []*pflag.Flag {

	var flags []*pflag.Flag

	c.LocalFlags().VisitAll(func(f *pflag.Flag) {

		if !f.Hidden {

			flags = append(flags, f)

		}

	})

	c.InheritedFlags().VisitAll(func(f *pflag.Flag) {

		if !f.Hidden {

			flags = append(flags, f)

		}

	})

	return flags

}

// zshCompGenFlagEntryForArguments returns an entry that matches _arguments

// zsh-completion parameters. It's too complicated to generate in a template.

func zshCompGenFlagEntryForArguments(f *pflag.Flag) string {

	if f.Name == "" || f.Shorthand == "" {

		return zshCompGenFlagEntryForSingleOptionFlag(f)

	}

	return zshCompGenFlagEntryForMultiOptionFlag(f)

func zshCompGenFlagEntryForSingleOptionFlag(f *pflag.Flag) string {

	var option, multiMark, extras string

	if zshCompFlagCouldBeSpecifiedMoreThenOnce(f) {

		multiMark = "*"

	}

	option = "--" + f.Name

	if option == "--" {

		option = "-" + f.Shorthand

	extras = zshCompGenFlagEntryExtras(f)

	return fmt.Sprintf(`'%s%s[%s]%s'`, multiMark, option, zshCompQuoteFlagDescription(f.Usage), extras)

func zshCompGenFlagEntryForMultiOptionFlag(f *pflag.Flag) string {

	var options, parenMultiMark, curlyMultiMark, extras string

	if zshCompFlagCouldBeSpecifiedMoreThenOnce(f) {

		parenMultiMark = "*"

		curlyMultiMark = "\\*"

	options = fmt.Sprintf(`'(%s-%s %s--%s)'{%s-%s,%s--%s}`,

		parenMultiMark, f.Shorthand, parenMultiMark, f.Name, curlyMultiMark, f.Shorthand, curlyMultiMark, f.Name)

	extras = zshCompGenFlagEntryExtras(f)

	return fmt.Sprintf(`%s'[%s]%s'`, options, zshCompQuoteFlagDescription(f.Usage), extras)

func zshCompGenFlagEntryExtras(f *pflag.Flag) string {

	if f.NoOptDefVal != "" {

		return ""

	}

	extras := ":" // allow options for flag (even without assistance)

	for key, values := range f.Annotations {

		switch key {

		case zshCompDirname:

			extras = fmt.Sprintf(":filename:_files -g %q", values[0])

		case BashCompFilenameExt:

			extras = ":filename:_files"

			for _, pattern := range values {

				extras = extras + fmt.Sprintf(` -g "%s"`, pattern)

			}

	return extras

func zshCompFlagCouldBeSpecifiedMoreThenOnce(f *pflag.Flag) bool {

	return strings.Contains(f.Value.Type(), "Slice") ||

		strings.Contains(f.Value.Type(), "Array")

func zshCompQuoteFlagDescription(s string) string {

	return strings.Replace(s, "'", `'\''`, -1)

}