package cobra

import (

	"bytes"

	"fmt"

	"io"

	"os"

	"strings"

)

// GenZshCompletionFile generates zsh completion file.

func (c *Command) GenZshCompletionFile(filename string) error {

	outFile, err := os.Create(filename)

	if err != nil {

		return err

	}

	defer outFile.Close()

	return c.GenZshCompletion(outFile)

}

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

func (c *Command) GenZshCompletion(w io.Writer) error {

	buf := new(bytes.Buffer)

	writeHeader(buf, c)

	maxDepth := maxDepth(c)

	writeLevelMapping(buf, maxDepth)

	writeLevelCases(buf, maxDepth, c)

	_, err := buf.WriteTo(w)

	return err

}

func writeHeader(w io.Writer, cmd *Command) {

	fmt.Fprintf(w, "#compdef %s\n\n", cmd.Name())

}

func maxDepth(c *Command) int {

	if len(c.Commands()) == 0 {

		return 0

	}

	maxDepthSub := 0

	for _, s := range c.Commands() {

		subDepth := maxDepth(s)

		if subDepth > maxDepthSub {

			maxDepthSub = subDepth

		}

	}

	return 1 + maxDepthSub

}

func writeLevelMapping(w io.Writer, numLevels int) {

	fmt.Fprintln(w, `_arguments \`)

	for i := 1; i <= numLevels; i++ {

		fmt.Fprintf(w, `  '%d: :->level%d' \`, i, i)

		fmt.Fprintln(w)

	}

	fmt.Fprintf(w, `  '%d: :%s'`, numLevels+1, "_files")

	fmt.Fprintln(w)

}

func writeLevelCases(w io.Writer, maxDepth int, root *Command) {

	fmt.Fprintln(w, "case $state in")

	defer fmt.Fprintln(w, "esac")

	for i := 1; i <= maxDepth; i++ {

		fmt.Fprintf(w, "  level%d)\n", i)

		writeLevel(w, root, i)

		fmt.Fprintln(w, "  ;;")

	}

	fmt.Fprintln(w, "  *)")

	fmt.Fprintln(w, "    _arguments '*: :_files'")

	fmt.Fprintln(w, "  ;;")

}

func writeLevel(w io.Writer, root *Command, i int) {

	fmt.Fprintf(w, "    case $words[%d] in\n", i)

	defer fmt.Fprintln(w, "    esac")

	commands := filterByLevel(root, i)

	byParent := groupByParent(commands)

	for p, c := range byParent {

		names := names(c)

		fmt.Fprintf(w, "      %s)\n", p)

		fmt.Fprintf(w, "        _arguments '%d: :(%s)'\n", i, strings.Join(names, " "))

		fmt.Fprintln(w, "      ;;")

	}

	fmt.Fprintln(w, "      *)")

	fmt.Fprintln(w, "        _arguments '*: :_files'")

	fmt.Fprintln(w, "      ;;")

}

func filterByLevel(c *Command, l int) []*Command {

	cs := make([]*Command, 0)

	if l == 0 {

		cs = append(cs, c)

		return cs

	}

	for _, s := range c.Commands() {

		cs = append(cs, filterByLevel(s, l-1)...)

	}

	return cs

}

func groupByParent(commands []*Command) map[string][]*Command {

	m := make(map[string][]*Command)

	for _, c := range commands {

		parent := c.Parent()

		if parent == nil {

			continue

		}

		m[parent.Name()] = append(m[parent.Name()], c)

	}

	return m

}

func names(commands []*Command) []string {

	ns := make([]string, len(commands))

	for i, c := range commands {

		ns[i] = c.Name()

	}

	return ns

}