// SPDX-License-Identifier: MIT
// Copyright (c) 2023 Robin Jarry

package opt

import (
	"errors"
	"strings"
)

// Shell command line with interpreted arguments.
// Allows access to individual arguments and to preserve shell quoting.
type Args struct {
	raw   []rune
	infos []argInfo
}

// Interpret a shell command line into multiple arguments.
func LexArgs(cmd string) *Args {
	raw := []rune(cmd)
	infos := lexCmdline(raw)
	return &Args{raw: raw, infos: infos}
}

// Shortcut for LexArgs(cmd).Args()
func SplitArgs(cmd string) []string {
	args := LexArgs(cmd)
	return args.Args()
}

// Build a shell command from multiple arguments.
func QuoteArgs(args ...string) *Args {
	quoted := make([]string, len(args))
	for i, arg := range args {
		quoted[i] = QuoteArg(arg)
	}
	return LexArgs(strings.Join(quoted, " "))
}

// Wrap a single argument with appropriate quoting so that it can be used
// in a shell command.
func QuoteArg(arg string) string {
	if strings.ContainsAny(arg, " '\"|?&!#$;[](){}<>*\n\t") {
		// "foo bar" --> "'foo bar'"
		// "foo'bar" --> "'foo'"'"'bar'"
		arg = "'" + strings.ReplaceAll(arg, "'", `'"'"'`) + "'"
	}
	return arg
}

// Get the number of arguments after interpreting shell quotes.
func (a *Args) Count() int {
	return len(a.infos)
}

func (a *Args) LeadingSpace() string {
	if len(a.infos) > 0 {
		first := &a.infos[0]
		if first.start < len(a.raw) {
			return string(a.raw[:first.start])
		}
	}
	return ""
}

func (a *Args) TrailingSpace() string {
	if len(a.infos) > 0 {
		last := &a.infos[len(a.infos)-1]
		if last.end < len(a.raw) {
			return string(a.raw[last.end:])
		}
	}
	return ""
}

var ErrArgIndex = errors.New("argument index out of bounds")

// Remove n arguments from the beginning of the command line.
// Same semantics as the `shift` built-in shell command.
// Will fail if shifting an invalid number of arguments.
func (a *Args) ShiftSafe(n int) ([]string, error) {
	var shifted []string
	switch {
	case n == 0:
		shifted = []string{}
	case n > 0 && n < len(a.infos):
		for i := 0; i < n; i++ {
			shifted = append(shifted, a.infos[i].unquoted)
		}
		a.infos = a.infos[n:]
		start := a.infos[0].start
		a.raw = a.raw[start:]
		for i := range a.infos {
			a.infos[i].start -= start
			a.infos[i].end -= start
		}
	case n == len(a.infos):
		for i := 0; i < n; i++ {
			shifted = append(shifted, a.infos[i].unquoted)
		}
		a.raw = nil
		a.infos = nil
	default:
		return nil, ErrArgIndex
	}
	return shifted, nil
}

// Same as ShiftSafe but cannot fail.
func (a *Args) Shift(n int) []string {
	if n < 0 {
		n = 0
	} else if n > len(a.infos) {
		n = len(a.infos)
	}
	shifted, _ := a.ShiftSafe(n)
	return shifted
}

// Remove n arguments from the end of the command line.
// Will fail if cutting an invalid number of arguments.
func (a *Args) CutSafe(n int) ([]string, error) {
	var cut []string
	switch {
	case n == 0:
		cut = []string{}
	case n > 0 && n < len(a.infos):
		for i := len(a.infos) - n; i < len(a.infos); i++ {
			cut = append(cut, a.infos[i].unquoted)
		}
		a.infos = a.infos[:len(a.infos)-n]
		a.raw = a.raw[:a.infos[len(a.infos)-1].end]
	case n == len(a.infos):
		for i := 0; i < n; i++ {
			cut = append(cut, a.infos[i].unquoted)
		}
		a.raw = nil
		a.infos = nil
	default:
		return nil, ErrArgIndex
	}
	return cut, nil
}

// Same as CutSafe but cannot fail.
func (a *Args) Cut(n int) []string {
	if n < 0 {
		n = 0
	} else if n > len(a.infos) {
		n = len(a.infos)
	}
	cut, _ := a.CutSafe(n)
	return cut
}

// Insert the specified prefix at the beginning of the command line.
func (a *Args) Prepend(cmd string) {
	prefix := []rune(cmd)
	infos := lexCmdline(prefix)
	if len(infos) > 0 && infos[len(infos)-1].end == len(prefix) {
		// No trailing white space to the added command line.
		// Add one space manually.
		prefix = append(prefix, ' ')
	}
	for i := range a.infos {
		a.infos[i].start += len(prefix)
		a.infos[i].end += len(prefix)
	}
	a.raw = append(prefix, a.raw...)
	a.infos = append(infos, a.infos...)
}

// Extend the command line with more arguments.
func (a *Args) Extend(cmd string) {
	suffix := []rune(cmd)
	infos := lexCmdline(suffix)
	if len(infos) > 0 && infos[0].start == 0 {
		// No leading white space to the added command line.
		// Add one space manually.
		a.raw = append(a.raw, ' ')
	}
	for i := range infos {
		infos[i].start += len(a.raw)
		infos[i].end += len(a.raw)
	}
	a.raw = append(a.raw, suffix...)
	a.infos = append(a.infos, infos...)
}

// Get the nth argument after interpreting shell quotes.
func (a *Args) ArgSafe(n int) (string, error) {
	if n < 0 || n >= len(a.infos) {
		return "", ErrArgIndex
	}
	return a.infos[n].unquoted, nil
}

// Get the nth argument after interpreting shell quotes.
// Will panic if the argument index does not exist.
func (a *Args) Arg(n int) string {
	return a.infos[n].unquoted
}

// Get all arguments after interpreting shell quotes.
func (a *Args) Args() []string {
	args := make([]string, 0, len(a.infos))
	for i := range a.infos {
		args = append(args, a.infos[i].unquoted)
	}
	return args
}

// Get the raw command line, with uninterpreted shell quotes.
func (a *Args) String() string {
	return string(a.raw)
}

// Make a deep copy of an Args object.
func (a *Args) Clone() *Args {
	infos := make([]argInfo, len(a.infos))
	copy(infos, a.infos)
	raw := make([]rune, len(a.raw))
	copy(raw, a.raw)
	return &Args{raw: raw, infos: infos}
}