package cobra

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	spfcobra "github.com/spf13/cobra"
)

// CompletionQueryToken is the hidden request understood by generated scripts.
const CompletionQueryToken = spfcobra.ShellCompRequestCmd

// CompletionScript generates a shell script for the installed executable.
func CompletionScript(w io.Writer, executable, shell string) error {
	root := &spfcobra.Command{Use: executable}
	var generated bytes.Buffer
	var err error
	switch shell {
	case "bash":
		err = root.GenBashCompletionV2(&generated, true)
	case "zsh":
		err = root.GenZshCompletion(&generated)
	case "fish":
		err = root.GenFishCompletion(&generated, true)
	default:
		return fmt.Errorf("unsupported completion shell %q", shell)
	}
	if err != nil {
		return err
	}

	script := generated.String()
	if shell == "zsh" {
		script, err = replaceCompletionText(script, `words=("${=words[1,CURRENT]}")`, `words=("${(@)words[1,CURRENT]}")`)
		if err == nil {
			script, err = replaceCompletionText(script, `=~ '-.*='`, `=~ '-[^=]*='`)
		}
	} else if shell == "fish" {
		script, err = adjustFishCompletion(script, root.Name())
	}
	if err != nil {
		return err
	}
	n, err := io.WriteString(w, script)
	if err == nil && n != len(script) {
		return io.ErrShortWrite
	}
	return err
}

func replaceCompletionText(script, old, replacement string) (string, error) {
	if strings.Count(script, old) != 1 {
		return "", fmt.Errorf("unsupported generated completion script")
	}
	return strings.Replace(script, old, replacement, 1), nil
}

func adjustFishCompletion(script, name string) (string, error) {
	var err error
	cobraName := strings.NewReplacer("-", "_", ":", "_").Replace(name)
	variableName := strings.ReplaceAll(cobraName, ".", "_")
	if cobraName != variableName {
		script = strings.ReplaceAll(script, "__"+cobraName+"_", "__"+variableName+"_")
	}
	script, err = replaceCompletionText(script,
		`set -l lastArg (string escape -- (commandline -ct))`,
		"set -l lastArg (string unescape -- (commandline -ct))\n    set -q lastArg[1]; or set lastArg \"\"")
	if err != nil {
		return "", err
	}
	script, err = replaceCompletionText(script,
		`set -l results (eval $requestComp 2> /dev/null)`,
		`set -lx COBRA_ACTIVE_HELP 0`+"\n    "+`set -l results ($args[1] `+spfcobra.ShellCompRequestCmd+` $args[2..-1] "$lastArg" 2> /dev/null)`)
	if err != nil {
		return "", err
	}
	script, err = replaceCompletionText(script, `string match -r -- '-.*=' "$lastArg"`, `string match -r -- '-[^=]*=' "$lastArg"`)
	if err != nil {
		return "", err
	}
	variable := "__" + variableName + "_comp_results"
	directories := fmt.Sprintf(`if test $dirfilter -eq 1
        set -l flagPrefix (string match -r -- '^-[^=]*=' (commandline -ct))
        set --global %[1]s
        for dir in (__fish_complete_directories)
            set -a %[1]s "$flagPrefix$dir"
        end
        return 0
    end
    if test $filefilter -eq 1`, variable)
	return replaceCompletionText(script, `if test $filefilter -eq 1; or test $dirfilter -eq 1`, directories)
}