// Package sshx implements the connection, authentication and host key policy
// layer shared by the sshtools command-line clients.
//
// It is non-interactive by design: credentials are supplied up front and host
// keys are either accepted, pinned or verified against a known_hosts file.
// Nothing here ever prompts.
package sshx

import (
	"errors"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"
)

// DefaultPort is used when neither the target nor the port flag sets one.
const DefaultPort = 22

// Target is a parsed [user@]host[:port] destination.
type Target struct {
	User string
	Host string
	Port int
}

// Addr returns the host:port address to dial.
func (t Target) Addr() string {
	return net.JoinHostPort(t.Host, strconv.Itoa(t.Port))
}

// ParseTarget parses raw as [user@]host[:port]. A non-empty user (from -l)
// overrides the user embedded in raw. A port embedded in raw overrides port.
func ParseTarget(raw, user string, port int) (Target, error) {
	t := Target{Port: port}
	if t.Port == 0 {
		t.Port = DefaultPort
	}
	host := strings.TrimSpace(raw)
	if host == "" {
		return t, errors.New("missing target: expected [user@]host")
	}
	if i := strings.LastIndex(host, "@"); i >= 0 {
		if user == "" {
			user = host[:i]
		}
		host = host[i+1:]
	}
	if user == "" {
		return t, fmt.Errorf("target %q has no user: use user@host or -l user", raw)
	}
	t.User = user

	if strings.HasPrefix(host, "[") {
		end := strings.Index(host, "]")
		if end < 0 {
			return t, fmt.Errorf("invalid target %q: missing ]", raw)
		}
		rest := host[end+1:]
		host = host[1:end]
		switch {
		case rest == "":
		case strings.HasPrefix(rest, ":"):
			p, err := parsePort(rest[1:], raw)
			if err != nil {
				return t, err
			}
			t.Port = p
		default:
			return t, fmt.Errorf("invalid target %q: unexpected %q after ]", raw, rest)
		}
	} else if i := strings.LastIndex(host, ":"); i >= 0 {
		if strings.Count(host, ":") > 1 {
			return t, fmt.Errorf("invalid target %q: bracket IPv6 literals like [::1]:22", raw)
		}
		p, err := parsePort(host[i+1:], raw)
		if err != nil {
			return t, err
		}
		t.Port = p
		host = host[:i]
	}
	if host == "" {
		return t, fmt.Errorf("invalid target %q: missing host", raw)
	}
	if t.Port < 1 || t.Port > 65535 {
		return t, fmt.Errorf("invalid port %d", t.Port)
	}
	t.Host = host
	return t, nil
}

func parsePort(s, raw string) (int, error) {
	p, err := strconv.Atoi(s)
	if err != nil || p < 1 || p > 65535 {
		return 0, fmt.Errorf("invalid target %q: bad port %q", raw, s)
	}
	return p, nil
}

// Config describes a connection attempt. Zero values are usable: an empty
// KnownHosts and HostKey means host keys are accepted without verification.
type Config struct {
	Target     Target
	Password   string
	KeyFile    string
	KnownHosts string
	HostKey    string
	Timeout    time.Duration
	Logf       func(format string, args ...any)
}

func (c Config) logf(format string, args ...any) {
	if c.Logf != nil {
		c.Logf(format, args...)
	}
}

// Dial connects to the target, completes the SSH handshake and authenticates.
// The caller owns the returned client and must close it. Dial never prompts.
func Dial(cfg Config) (*ssh.Client, error) {
	auth, err := authMethods(cfg)
	if err != nil {
		return nil, err
	}
	hostKey, err := hostKeyCallback(cfg)
	if err != nil {
		return nil, err
	}
	addr := cfg.Target.Addr()
	clientConfig := &ssh.ClientConfig{
		User:            cfg.Target.User,
		Auth:            auth,
		HostKeyCallback: hostKey,
		Timeout:         cfg.Timeout,
	}
	cfg.logf("connecting to %s", addr)
	dialer := net.Dialer{Timeout: cfg.Timeout}
	conn, err := dialer.Dial("tcp", addr)
	if err != nil {
		return nil, fmt.Errorf("connect to %s: %w", addr, err)
	}
	if cfg.Timeout > 0 {
		// Bound the handshake and authentication too, then clear the deadline:
		// the remote command itself may legitimately run for hours.
		_ = conn.SetDeadline(time.Now().Add(cfg.Timeout))
	}
	sshConn, chans, reqs, err := ssh.NewClientConn(conn, addr, clientConfig)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("ssh handshake with %s: %w", addr, err)
	}
	_ = conn.SetDeadline(time.Time{})
	cfg.logf("authenticated as %s", cfg.Target.User)
	return ssh.NewClient(sshConn, chans, reqs), nil
}

// authMethods builds the credential list. Public keys come first; the password
// is offered both as the "password" method and as keyboard-interactive, because
// plenty of servers (PAM setups in particular) only offer the latter.
func authMethods(cfg Config) ([]ssh.AuthMethod, error) {
	var methods []ssh.AuthMethod
	if cfg.KeyFile != "" {
		signer, err := loadSigner(cfg.KeyFile)
		if err != nil {
			return nil, err
		}
		methods = append(methods, ssh.PublicKeys(signer))
	}
	if cfg.Password != "" {
		password := cfg.Password
		methods = append(methods, ssh.Password(password))
		methods = append(methods, ssh.KeyboardInteractive(
			func(name, instruction string, questions []string, echos []bool) ([]string, error) {
				answers := make([]string, len(questions))
				for i := range answers {
					answers[i] = password
				}
				return answers, nil
			}))
	}
	if len(methods) == 0 {
		return nil, errors.New("no credentials: pass -p <password> (or set GOLINK_PASSWORD) or -i <keyfile>")
	}
	return methods, nil
}

func loadSigner(path string) (ssh.Signer, error) {
	pemBytes, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read key %s: %w", path, err)
	}
	signer, err := ssh.ParsePrivateKey(pemBytes)
	if err != nil {
		var missing *ssh.PassphraseMissingError
		if errors.As(err, &missing) {
			return nil, fmt.Errorf("key %s is passphrase protected: golink is non-interactive, decrypt it with 'ssh-keygen -p' or use -p", path)
		}
		return nil, fmt.Errorf("parse key %s: %w", path, err)
	}
	return signer, nil
}

// hostKeyCallback selects the host key policy: pinned fingerprint, known_hosts
// file, or accept anything (the default, since a prompt is not an option).
func hostKeyCallback(cfg Config) (ssh.HostKeyCallback, error) {
	var cb ssh.HostKeyCallback
	switch {
	case cfg.HostKey != "":
		cb = pinnedHostKey(cfg.HostKey)
	case cfg.KnownHosts != "":
		verify, err := knownhosts.New(cfg.KnownHosts)
		if err != nil {
			return nil, fmt.Errorf("load known_hosts %s: %w", cfg.KnownHosts, err)
		}
		cb = func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			err := verify(hostname, remote, key)
			if err == nil {
				return nil
			}
			var keyErr *knownhosts.KeyError
			if errors.As(err, &keyErr) {
				if len(keyErr.Want) == 0 {
					return fmt.Errorf("host %s is not in %s: add it with 'ssh-keyscan -p %d %s' or pass -hostkey", hostname, cfg.KnownHosts, cfg.Target.Port, cfg.Target.Host)
				}
				return fmt.Errorf("host key mismatch for %s against %s (possible MITM): %w", hostname, cfg.KnownHosts, err)
			}
			return err
		}
	default:
		cb = ssh.InsecureIgnoreHostKey()
	}
	if cfg.Logf != nil {
		inner := cb
		cb = func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			cfg.logf("host key %s (%s)", ssh.FingerprintSHA256(key), key.Type())
			return inner(hostname, remote, key)
		}
	}
	return cb, nil
}

func pinnedHostKey(want string) ssh.HostKeyCallback {
	want = strings.TrimSpace(want)
	return func(hostname string, remote net.Addr, key ssh.PublicKey) error {
		sha256 := ssh.FingerprintSHA256(key)
		if strings.EqualFold(want, sha256) ||
			strings.EqualFold(want, strings.TrimPrefix(sha256, "SHA256:")) ||
			strings.EqualFold(want, ssh.FingerprintLegacyMD5(key)) {
			return nil
		}
		return fmt.Errorf("host key for %s does not match -hostkey:\n  offered: %s\n  wanted:  %s", hostname, sha256, want)
	}
}