// Package sshx provides the SSH transport used to drive a target host that has
// nothing installed but sshd and the docker CLI.
package sshx

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

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

// AuthMethod selects how to authenticate against the target host.
type AuthMethod string

const (
	AuthPassword AuthMethod = "password"
	AuthKey      AuthMethod = "key"
	AuthAgent    AuthMethod = "agent"
)

// Config describes one target host.
type Config struct {
	ID   string `json:"id"`
	Name string `json:"name"`
	Host string `json:"host"`
	Port int    `json:"port"`
	User string `json:"user"`

	Auth AuthMethod `json:"auth"`
	// Password is used with AuthPassword, and as the passphrase fallback when
	// a key is encrypted.
	Password string `json:"password,omitempty"`
	// PrivateKey holds PEM key material for AuthKey. PrivateKeyPath is read
	// from disk instead when PrivateKey is empty.
	PrivateKey     string `json:"privateKey,omitempty"`
	PrivateKeyPath string `json:"privateKeyPath,omitempty"`
	Passphrase     string `json:"passphrase,omitempty"`

	// Sudo prefixes every docker command with sudo -n, for hosts where the
	// login user is not in the docker group.
	Sudo bool `json:"sudo"`
	// DockerCmd overrides the docker binary, e.g. "podman" or an absolute path.
	DockerCmd string `json:"dockerCmd,omitempty"`

	// SaveSecrets persists the password and key material to the connection
	// store. When false the secrets live only for the current process.
	SaveSecrets bool `json:"saveSecrets"`

	// Timeout is the TCP/handshake timeout. Zero means 20s.
	Timeout time.Duration `json:"-"`
}

func (c Config) addr() string {
	port := c.Port
	if port == 0 {
		port = 22
	}
	return net.JoinHostPort(c.Host, strconv.Itoa(port))
}

// Client is a live SSH connection to a target host.
type Client struct {
	cfg  Config
	conn *ssh.Client
}

// HostKeyError reports that the target's host key is unknown or has changed.
// The UI shows the fingerprint and asks the operator to confirm before the key
// is written to the known-hosts store.
type HostKeyError struct {
	Host        string
	Fingerprint string
	KeyType     string
	Changed     bool // true when a different key was already trusted
}

func (e *HostKeyError) Error() string {
	if e.Changed {
		return fmt.Sprintf("host key for %s CHANGED (%s %s); refusing to connect", e.Host, e.KeyType, e.Fingerprint)
	}
	return fmt.Sprintf("host key for %s is not trusted yet (%s %s)", e.Host, e.KeyType, e.Fingerprint)
}

// Dial opens a connection, verifying the host key against the known-hosts
// store. It returns a *HostKeyError when the operator has to make a trust
// decision first.
func Dial(ctx context.Context, cfg Config, hk *KnownHosts) (*Client, error) {
	auths, err := authMethods(cfg)
	if err != nil {
		return nil, err
	}
	timeout := cfg.Timeout
	if timeout == 0 {
		timeout = 20 * time.Second
	}

	var hkErr *HostKeyError
	clientCfg := &ssh.ClientConfig{
		User:    cfg.User,
		Auth:    auths,
		Timeout: timeout,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			err := hk.Check(hostname, remote, key)
			var he *HostKeyError
			if errors.As(err, &he) {
				hkErr = he
			}
			return err
		},
	}

	d := net.Dialer{Timeout: timeout}
	rawConn, err := d.DialContext(ctx, "tcp", cfg.addr())
	if err != nil {
		return nil, fmt.Errorf("connect to %s: %w", cfg.addr(), err)
	}
	sshConn, chans, reqs, err := ssh.NewClientConn(rawConn, cfg.addr(), clientCfg)
	if err != nil {
		rawConn.Close()
		if hkErr != nil {
			return nil, hkErr
		}
		return nil, fmt.Errorf("ssh handshake with %s: %w", cfg.addr(), err)
	}
	return &Client{cfg: cfg, conn: ssh.NewClient(sshConn, chans, reqs)}, nil
}

// Close terminates the connection.
func (c *Client) Close() error { return c.conn.Close() }

// Config returns the configuration this client was dialled with.
func (c *Client) Config() Config { return c.cfg }

// Result is the outcome of a remote command.
type Result struct {
	Stdout   string
	Stderr   string
	ExitCode int
}

// Run executes a command line on the remote host and collects its output.
// The command is passed to the remote login shell, so it may contain pipes.
func (c *Client) Run(ctx context.Context, cmdline string) (*Result, error) {
	var stdout, stderr bytes.Buffer
	code, err := c.run(ctx, cmdline, nil, &stdout, &stderr)
	res := &Result{Stdout: stdout.String(), Stderr: stderr.String(), ExitCode: code}
	return res, err
}

// RunCheck executes a command and turns a non-zero exit into an error that
// carries the remote stderr, which is what the operator needs to see.
func (c *Client) RunCheck(ctx context.Context, cmdline string) (string, error) {
	res, err := c.Run(ctx, cmdline)
	if err != nil {
		return res.Stdout, err
	}
	if res.ExitCode != 0 {
		msg := strings.TrimSpace(res.Stderr)
		if msg == "" {
			msg = strings.TrimSpace(res.Stdout)
		}
		return res.Stdout, fmt.Errorf("remote command failed (exit %d): %s", res.ExitCode, msg)
	}
	return res.Stdout, nil
}

// Stream executes a command, feeding it stdin and writing its stdout to out.
// This is how bulk data crosses the wire: the tar stream produced locally is
// piped straight into a remote `docker cp` without ever touching disk.
func (c *Client) Stream(ctx context.Context, cmdline string, stdin io.Reader, stdout io.Writer) (*Result, error) {
	var stderr bytes.Buffer
	if stdout == nil {
		stdout = io.Discard
	}
	code, err := c.run(ctx, cmdline, stdin, stdout, &stderr)
	res := &Result{Stderr: stderr.String(), ExitCode: code}
	if err != nil {
		return res, err
	}
	if code != 0 {
		return res, fmt.Errorf("remote command failed (exit %d): %s", code, strings.TrimSpace(stderr.String()))
	}
	return res, nil
}

func (c *Client) run(ctx context.Context, cmdline string, stdin io.Reader, stdout, stderr io.Writer) (int, error) {
	sess, err := c.conn.NewSession()
	if err != nil {
		return -1, fmt.Errorf("open ssh session: %w", err)
	}
	defer sess.Close()

	sess.Stdout = stdout
	sess.Stderr = stderr
	if stdin != nil {
		sess.Stdin = stdin
	}

	done := make(chan error, 1)
	go func() { done <- sess.Run(cmdline) }()

	select {
	case <-ctx.Done():
		_ = sess.Signal(ssh.SIGTERM)
		_ = sess.Close()
		return -1, ctx.Err()
	case err := <-done:
		if err == nil {
			return 0, nil
		}
		var ee *ssh.ExitError
		if errors.As(err, &ee) {
			return ee.ExitStatus(), nil
		}
		return -1, err
	}
}

func authMethods(cfg Config) ([]ssh.AuthMethod, error) {
	var methods []ssh.AuthMethod
	switch cfg.Auth {
	case AuthPassword:
		if cfg.Password == "" {
			return nil, errors.New("password authentication selected but no password supplied")
		}
		methods = append(methods,
			ssh.Password(cfg.Password),
			// Many sshd setups answer with keyboard-interactive instead of the
			// plain password method.
			ssh.KeyboardInteractive(func(_, _ string, questions []string, _ []bool) ([]string, error) {
				answers := make([]string, len(questions))
				for i := range answers {
					answers[i] = cfg.Password
				}
				return answers, nil
			}),
		)
	case AuthKey:
		pem := []byte(cfg.PrivateKey)
		if len(pem) == 0 {
			if cfg.PrivateKeyPath == "" {
				return nil, errors.New("key authentication selected but no key supplied")
			}
			b, err := os.ReadFile(cfg.PrivateKeyPath)
			if err != nil {
				return nil, fmt.Errorf("read private key: %w", err)
			}
			pem = b
		}
		var signer ssh.Signer
		var err error
		passphrase := cfg.Passphrase
		if passphrase == "" {
			passphrase = cfg.Password
		}
		if passphrase != "" {
			signer, err = ssh.ParsePrivateKeyWithPassphrase(pem, []byte(passphrase))
		} else {
			signer, err = ssh.ParsePrivateKey(pem)
		}
		if err != nil {
			var pm *ssh.PassphraseMissingError
			if errors.As(err, &pm) {
				return nil, errors.New("private key is encrypted; supply the passphrase")
			}
			return nil, fmt.Errorf("parse private key: %w", err)
		}
		methods = append(methods, ssh.PublicKeys(signer))
	case AuthAgent:
		sock := os.Getenv("SSH_AUTH_SOCK")
		if sock == "" {
			return nil, errors.New("agent authentication selected but SSH_AUTH_SOCK is not set")
		}
		conn, err := net.Dial("unix", sock)
		if err != nil {
			return nil, fmt.Errorf("connect to ssh agent: %w", err)
		}
		methods = append(methods, ssh.PublicKeysCallback(agent.NewClient(conn).Signers))
	default:
		return nil, fmt.Errorf("unknown auth method %q", cfg.Auth)
	}
	return methods, nil
}

// Fingerprint renders a public key the way OpenSSH shows it.
func Fingerprint(key ssh.PublicKey) string { return ssh.FingerprintSHA256(key) }

var _ = knownhosts.Normalize