// Package connhelper provides helpers for connecting to a remote daemon host with custom logic.
package connhelper

import (
	"context"
	"net"
	"net/url"
	"os"
	"strings"

	"github.com/docker/cli/cli/connhelper/commandconn"
	"github.com/docker/cli/cli/connhelper/ssh"
	"github.com/pkg/errors"
)

const (
	// DockerSSHRemoteBinaryEnv is the environment variable that can be used to
	// override the default Docker binary called over SSH
	DockerSSHRemoteBinaryEnv = "DOCKER_SSH_REMOTE_BINARY"
)

// ConnectionHelper allows to connect to a remote host with custom stream provider binary.
type ConnectionHelper struct {
	Dialer func(ctx context.Context, network, addr string) (net.Conn, error)
	Host   string // dummy URL used for HTTP requests. e.g. "http://docker"
}

// GetConnectionHelper returns Docker-specific connection helper for the given URL.
// GetConnectionHelper returns nil without error when no helper is registered for the scheme.
//
// ssh://<user>@<host> URL requires Docker 18.09 or later on the remote host.
func GetConnectionHelper(daemonURL string) (*ConnectionHelper, error) {
	return getConnectionHelper(daemonURL, nil)
}

// GetConnectionHelperWithSSHOpts returns Docker-specific connection helper for
// the given URL, and accepts additional options for ssh connections. It returns
// nil without error when no helper is registered for the scheme.
//
// Requires Docker 18.09 or later on the remote host.
func GetConnectionHelperWithSSHOpts(daemonURL string, sshFlags []string) (*ConnectionHelper, error) {
	return getConnectionHelper(daemonURL, sshFlags)
}

func getConnectionHelper(daemonURL string, sshFlags []string) (*ConnectionHelper, error) {
	u, err := url.Parse(daemonURL)
	if err != nil {
		return nil, err
	}
	if u.Scheme == "ssh" {
		sp, err := ssh.ParseURL(daemonURL)
		if err != nil {
			return nil, errors.Wrap(err, "ssh host connection is not valid")
		}
		sshFlags = addSSHTimeout(sshFlags)
		sshFlags = disablePseudoTerminalAllocation(sshFlags)
		remoteDockerBinary := dockerSSHRemoteBinary()
		return &ConnectionHelper{
			Dialer: func(ctx context.Context, network, addr string) (net.Conn, error) {
				args := []string{remoteDockerBinary}
				if sp.Path != "" {
					args = append(args, "--host", "unix://"+sp.Path)
				}
				args = append(args, "system", "dial-stdio")
				return commandconn.New(ctx, "ssh", append(sshFlags, sp.Args(args...)...)...)
			},
			Host: "http://docker.example.com",
		}, nil
	}
	// Future version may support plugins via ~/.docker/config.json. e.g. "dind"
	// See docker/cli#889 for the previous discussion.
	return nil, err
}

// GetCommandConnectionHelper returns Docker-specific connection helper constructed from an arbitrary command.
func GetCommandConnectionHelper(cmd string, flags ...string) (*ConnectionHelper, error) {
	return &ConnectionHelper{
		Dialer: func(ctx context.Context, network, addr string) (net.Conn, error) {
			return commandconn.New(ctx, cmd, flags...)
		},
		Host: "http://docker.example.com",
	}, nil
}

func addSSHTimeout(sshFlags []string) []string {
	if !strings.Contains(strings.Join(sshFlags, ""), "ConnectTimeout") {
		sshFlags = append(sshFlags, "-o ConnectTimeout=30")
	}
	return sshFlags
}

// disablePseudoTerminalAllocation disables pseudo-terminal allocation to
// prevent SSH from executing as a login shell
func disablePseudoTerminalAllocation(sshFlags []string) []string {
	for _, flag := range sshFlags {
		if flag == "-T" {
			return sshFlags
		}
	}
	return append(sshFlags, "-T")
}

// dockerSSHRemoteBinary returns the binary to use when executing Docker
// commands over SSH. It defaults to "docker" if the DOCKER_SSH_REMOTE_BINARY
// environment variable is not set.
func dockerSSHRemoteBinary() string {
	value := os.Getenv(DockerSSHRemoteBinaryEnv)
	if value == "" {
		return "docker"
	}

	return value
}