package ssh

import "strings"

// flagsWithArg lists all SSH flags that consume the following argument.
var flagsWithArg = map[byte]bool{
	'B': true, 'P': true, '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 {
			// OpenSSH accepts clusters such as -vp 2222 and -vladmin.
			for j := 1; j < len(arg); j++ {
				if flagsWithArg[arg[j]] {
					if j == len(arg)-1 {
						i++
					}
					break
				}
			}
			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
}

// HasPortFlag reports whether a port was explicitly supplied.
func HasPortFlag(args []string) bool { return hasOption(args, 'p', "Port") }

// HasUserFlag reports whether a login user was explicitly supplied.
func HasUserFlag(args []string) bool { return hasOption(args, 'l', "User") }

func hasOption(args []string, wanted byte, keyword string) bool {
	for i := 0; i < len(args); i++ {
		arg := args[i]
		if arg == "--" || len(arg) < 2 || arg[0] != '-' {
			break
		}
		for j := 1; j < len(arg); j++ {
			flag := arg[j]
			if !flagsWithArg[flag] {
				continue
			}
			value := arg[j+1:]
			if value == "" {
				i++
				if i >= len(args) {
					return false
				}
				value = args[i]
			}
			if flag == wanted {
				return true
			}
			if flag == 'o' {
				fields := strings.FieldsFunc(value, func(r rune) bool { return r == '=' || r == ' ' || r == '\t' })
				if len(fields) >= 2 && strings.EqualFold(fields[0], keyword) {
					return true
				}
			}
			break
		}
	}
	return false
}

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