package ssh

import "strings"

// flagsWithArg lists all SSH flags that consume the following argument.
var flagsWithArg = map[byte]bool{
	'b': true, 'c': true, 'D': true, 'E': true, 'e': true,
	'F': true, 'I': true, 'i': true, 'J': true, 'L': true,
	'l': true, 'm': true, 'o': true, 'O': true, 'p': true,
	'Q': true, 'R': true, 'S': true, 'w': true, 'W': true,
}

// ParsedArgs holds the result of parsing SSH arguments.
type ParsedArgs struct {
	Host    string // hostname without the user@ prefix
	User    string // empty if not specified
	DestIdx int    // index in Args where [user@]host sits
	Args    []string
}

// Parse scans SSH arguments and extracts the destination ([user@]host).
// Returns nil if no destination is found.
func Parse(args []string) *ParsedArgs {
	i := 0
	for i < len(args) {
		arg := args[i]

		// "--" ends option processing
		if arg == "--" {
			i++
			break
		}

		if strings.HasPrefix(arg, "-") && len(arg) > 1 {
			flag := arg[1]
			if flagsWithArg[flag] {
				if len(arg) > 2 {
					// argument is attached, e.g. -p2222
					i++
				} else {
					// argument is the next element, e.g. -p 2222
					i += 2
				}
			} else {
				i++
			}
			continue
		}

		// first non-flag argument is the destination
		host, user := splitUserHost(arg)
		return &ParsedArgs{
			Host:    host,
			User:    user,
			DestIdx: i,
			Args:    args,
		}
	}

	// handle arguments after "--"
	if i < len(args) {
		host, user := splitUserHost(args[i])
		return &ParsedArgs{
			Host:    host,
			User:    user,
			DestIdx: i,
			Args:    args,
		}
	}

	return nil
}

// ReplaceHost returns a copy of args with the destination replaced by newHost,
// preserving any user@ prefix.
func ReplaceHost(args []string, destIdx int, newHost string) []string {
	result := make([]string, len(args))
	copy(result, args)

	original := args[destIdx]
	if at := strings.Index(original, "@"); at != -1 {
		result[destIdx] = original[:at+1] + newHost
	} else {
		result[destIdx] = newHost
	}
	return result
}

// HasUserFlag reports whether a user was specified via -l in args.
// Used to avoid overriding an explicit -l with the configured default user.
func HasUserFlag(args []string) bool {
	for i, a := range args {
		if a == "-l" && i+1 < len(args) {
			return true
		}
		// handle attached form: -lroot
		if len(a) > 2 && a[0] == '-' && a[1] == 'l' {
			return true
		}
	}
	return false
}

func splitUserHost(dest string) (host, user string) {
	if at := strings.Index(dest, "@"); at != -1 {
		return dest[at+1:], dest[:at]
	}
	return dest, ""
}