client.go

package client

import (
	"context"
	_ "embed"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/databricks/cli/experimental/ssh/internal/keys"
	"github.com/databricks/cli/experimental/ssh/internal/proxy"
	"github.com/databricks/cli/experimental/ssh/internal/sessions"
	"github.com/databricks/cli/experimental/ssh/internal/sshconfig"
	"github.com/databricks/cli/experimental/ssh/internal/vscode"
	sshWorkspace "github.com/databricks/cli/experimental/ssh/internal/workspace"
	"github.com/databricks/cli/internal/build"
	"github.com/databricks/cli/libs/cmdio"
	"github.com/databricks/cli/libs/log"
	"github.com/databricks/cli/libs/telemetry"
	"github.com/databricks/cli/libs/telemetry/protos"
	"github.com/databricks/databricks-sdk-go"
	"github.com/databricks/databricks-sdk-go/retries"
	"github.com/databricks/databricks-sdk-go/service/compute"
	"github.com/databricks/databricks-sdk-go/service/jobs"
	"github.com/databricks/databricks-sdk-go/service/workspace"
	"github.com/gorilla/websocket"
)

//go:embed ssh-server-bootstrap.py
var sshServerBootstrapScript string

var errServerMetadata = errors.New("server metadata error")

var connectionNameRegex = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_-]*$`)

const (
	sshServerTaskKey         = "start_ssh_server"
	serverlessEnvironmentKey = "ssh_tunnel_serverless"
	minEnvironmentVersion    = 4
)

type ClientOptions struct {
	// Id of the cluster to connect to (for dedicated clusters)
	ClusterID string
	// Connection name (for serverless compute). Used as unique identifier instead of ClusterID.
	ConnectionName string
	// GPU accelerator type (for serverless compute)
	Accelerator string
	// Delay before shutting down the server after the last client disconnects
	ShutdownDelay time.Duration
	// Maximum number of SSH clients
	MaxClients int
	// Indicates that the CLI runs as a ProxyCommand - it should establish ws connection
	// to the cluster and proxy all traffic through stdin/stdout.
	// In the non proxy mode the CLI spawns an ssh client with the ProxyCommand config.
	ProxyMode bool
	// Open remote IDE window with a specific ssh config (empty, 'vscode', or 'cursor')
	IDE string
	// Expected format: "<user_name>,<port>,<cluster_id>".
	// If present, the CLI won't attempt to start the server.
	ServerMetadata string
	// How often the CLI should reconnect to the server with new auth.
	HandoverTimeout time.Duration
	// Max amount of time the server process is allowed to live
	ServerTimeout time.Duration
	// Max amount of time to wait for the SSH server task to reach RUNNING state
	TaskStartupTimeout time.Duration
	// Directory for local SSH tunnel development releases.
	// If not present, the CLI will use github releases with the current version.
	ReleasesDir string
	// Directory for local SSH keys. Defaults to ~/.databricks/ssh-tunnel-keys
	SSHKeysDir string
	// Client public key name located in the ssh-tunnel secrets scope.
	ClientPublicKeyName string
	// Client private key name located in the ssh-tunnel secrets scope.
	ClientPrivateKeyName string
	// If true, the CLI will attempt to start the cluster if it is not running.
	AutoStartCluster bool
	// Optional auth profile name. If present, will be added as --profile flag to the ProxyCommand while spawning ssh client.
	Profile string
	// Additional arguments to pass to the SSH client in the non proxy mode.
	AdditionalArgs []string
	// Optional path to the user known hosts file.
	UserKnownHostsFile string
	// Liteswap header value for traffic routing (dev/test only).
	Liteswap string
	// If true, skip checking and updating IDE settings.
	SkipSettingsCheck bool
	// Environment version for serverless compute.
	EnvironmentVersion int
}

func (o *ClientOptions) Validate() error {
	if !o.ProxyMode && o.ClusterID == "" && o.ConnectionName == "" && o.Accelerator == "" {
		return errors.New("please provide --cluster or --accelerator flag")
	}
	if o.Accelerator != "" && o.ClusterID != "" {
		return errors.New("--accelerator flag can only be used with serverless compute, not with --cluster")
	}
	// TODO: Remove when we add support for serverless CPU
	if o.ConnectionName != "" && o.Accelerator == "" {
		return errors.New("--name flag requires --accelerator to be set (for now we only support serverless GPU compute)")
	}
	if o.ConnectionName != "" && !connectionNameRegex.MatchString(o.ConnectionName) {
		return fmt.Errorf("connection name %q must consist of letters, numbers, dashes, and underscores", o.ConnectionName)
	}
	if o.IDE != "" && o.IDE != vscode.VSCodeOption && o.IDE != vscode.CursorOption {
		return fmt.Errorf("invalid IDE value: %q, expected %q or %q", o.IDE, vscode.VSCodeOption, vscode.CursorOption)
	}
	if o.EnvironmentVersion > 0 && o.EnvironmentVersion < minEnvironmentVersion {
		return fmt.Errorf("environment version must be >= %d, got %d", minEnvironmentVersion, o.EnvironmentVersion)
	}
	return nil
}

func (o *ClientOptions) IsServerlessMode() bool {
	return o.ClusterID == "" && (o.ConnectionName != "" || o.Accelerator != "")
}

// SessionIdentifier returns the unique identifier for the session.
// For dedicated clusters, this is the cluster ID. For serverless, this is the connection name.
func (o *ClientOptions) SessionIdentifier() string {
	if o.IsServerlessMode() {
		return o.ConnectionName
	}
	return o.ClusterID
}

// FormatMetadata formats the server metadata string for use in ProxyCommand.
// Returns empty string if userName is empty or serverPort is zero.
func FormatMetadata(userName string, serverPort int, clusterID string) string {
	if userName == "" || serverPort == 0 {
		return ""
	}
	if clusterID != "" {
		return fmt.Sprintf("%s,%d,%s", userName, serverPort, clusterID)
	}
	return fmt.Sprintf("%s,%d", userName, serverPort)
}

// ToProxyCommand generates the ProxyCommand string for SSH config.
// This method serializes the ClientOptions into a command-line invocation that will
// be parsed back into ClientOptions when the SSH ProxyCommand is executed.
func (o *ClientOptions) ToProxyCommand() (string, error) {
	executablePath, err := os.Executable()
	if err != nil {
		return "", fmt.Errorf("failed to get current executable path: %w", err)
	}

	var proxyCommand string
	if o.IsServerlessMode() {
		proxyCommand = fmt.Sprintf("%q ssh connect --proxy --name=%s --shutdown-delay=%s",
			executablePath, o.ConnectionName, o.ShutdownDelay.String())
		if o.Accelerator != "" {
			proxyCommand += " --accelerator=" + o.Accelerator
		}
	} else {
		proxyCommand = fmt.Sprintf("%q ssh connect --proxy --cluster=%s --auto-start-cluster=%t --shutdown-delay=%s",
			executablePath, o.ClusterID, o.AutoStartCluster, o.ShutdownDelay.String())
	}

	if o.ServerMetadata != "" {
		proxyCommand += " --metadata=" + o.ServerMetadata
	}

	if o.HandoverTimeout > 0 {
		proxyCommand += " --handover-timeout=" + o.HandoverTimeout.String()
	}

	if o.Profile != "" {
		proxyCommand += " --profile=" + o.Profile
	}

	if o.Liteswap != "" {
		proxyCommand += " --liteswap=" + o.Liteswap
	}

	if o.EnvironmentVersion > 0 {
		proxyCommand += " --environment-version=" + strconv.Itoa(o.EnvironmentVersion)
	}

	return proxyCommand, nil
}

func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOptions) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt, syscall.SIGHUP, syscall.SIGTERM)
	go func() {
		<-sigCh
		cmdio.LogString(ctx, "Received termination signal, cleaning up...")
		cancel()
	}()

	event := BuildTelemetryEvent(opts)

	runErr := runConnect(ctx, client, opts, event)
	if runErr != nil {
		event.IsSuccess = false
	} else {
		event.IsSuccess = true
	}

	telemetry.Log(ctx, protos.DatabricksCliLog{
		SshTunnelEvent: event,
	})

	return runErr
}

// BuildTelemetryEvent creates an SshTunnelEvent pre-populated with data from client options.
func BuildTelemetryEvent(opts ClientOptions) *protos.SshTunnelEvent {
	event := &protos.SshTunnelEvent{
		AcceleratorType:  opts.Accelerator,
		IdeType:          opts.IDE,
		AutoStartCluster: opts.AutoStartCluster,
	}

	if opts.IsServerlessMode() {
		event.ComputeType = protos.SshTunnelComputeTypeServerless
	} else {
		event.ComputeType = protos.SshTunnelComputeTypeDedicated
	}

	switch {
	case opts.ProxyMode:
		event.ClientMode = protos.SshTunnelClientModeProxy
	case opts.IDE != "":
		event.ClientMode = protos.SshTunnelClientModeIDE
	default:
		event.ClientMode = protos.SshTunnelClientModeSSH
	}

	// If metadata is provided, the server is already running — this is a reconnect from ProxyCommand.
	event.IsReconnect = opts.ServerMetadata != ""

	return event
}

func runConnect(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOptions, event *protos.SshTunnelEvent) error {
	// For serverless without explicit --name: auto-generate or reconnect to existing session.
	if opts.IsServerlessMode() && opts.ConnectionName == "" && !opts.ProxyMode {
		err := resolveServerlessSession(ctx, client, &opts)
		if err != nil {
			return err
		}
	}

	sessionID := opts.SessionIdentifier()
	if sessionID == "" {
		return errors.New("either --cluster or --accelerator must be provided")
	}

	if !opts.ProxyMode {
		cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s...", sessionID))
	}

	if opts.IDE != "" && !opts.ProxyMode {
		if err := vscode.CheckIDECommand(opts.IDE); err != nil {
			return err
		}
		if err := vscode.CheckIDESSHExtension(ctx, opts.IDE); err != nil {
			return err
		}
	}

	// Check and update IDE settings for serverless mode, where we must set up
	// desired server ports (or socket connection mode) for the connection to go through
	// (as the majority of the localhost ports on the remote side are blocked by iptable rules).
	// Plus the platform (always linux), and extensions (python and jupyter), to make the initial experience smoother.
	if opts.IDE != "" && opts.IsServerlessMode() && !opts.ProxyMode && !opts.SkipSettingsCheck && cmdio.IsPromptSupported(ctx) {
		err := vscode.CheckAndUpdateSettings(ctx, opts.IDE, opts.ConnectionName)
		if err != nil {
			cmdio.LogString(ctx, fmt.Sprintf("Failed to update IDE settings: %v", err))
			cmdio.LogString(ctx, vscode.GetManualInstructions(opts.IDE, opts.ConnectionName))
			cmdio.LogString(ctx, "Use --skip-settings-check to bypass IDE settings verification.")
			shouldProceed, promptErr := cmdio.AskYesOrNo(ctx, "Do you want to proceed with the connection?")
			if promptErr != nil {
				return fmt.Errorf("failed to prompt user: %w", promptErr)
			}
			if !shouldProceed {
				return errors.New("aborted: IDE settings need to be updated manually, user declined to proceed")
			}
		}
	}

	// Only check cluster state for dedicated clusters
	if !opts.IsServerlessMode() {
		cmdio.LogString(ctx, "Checking cluster state...")
		err := checkClusterState(ctx, client, opts.ClusterID, opts.AutoStartCluster)
		if err != nil {
			return err
		}
	}

	secretScopeName, err := keys.CreateKeysSecretScope(ctx, client, sessionID)
	if err != nil {
		return fmt.Errorf("failed to create secret scope: %w", err)
	}

	privateKeyBytes, publicKeyBytes, err := keys.CheckAndGenerateSSHKeyPairFromSecrets(ctx, client, secretScopeName, opts.ClientPrivateKeyName, opts.ClientPublicKeyName)
	if err != nil {
		return fmt.Errorf("failed to get or generate SSH key pair from secrets: %w", err)
	}

	keyPath, err := keys.GetLocalSSHKeyPath(ctx, sessionID, opts.SSHKeysDir)
	if err != nil {
		return fmt.Errorf("failed to get local keys folder: %w", err)
	}

	err = keys.SaveSSHKeyPair(keyPath, privateKeyBytes, publicKeyBytes)
	if err != nil {
		return fmt.Errorf("failed to save SSH key pair locally: %w", err)
	}
	log.Infof(ctx, "Using SSH key: %s", keyPath)
	log.Infof(ctx, "Secrets scope: %s, key name: %s", secretScopeName, opts.ClientPublicKeyName)

	var userName string
	var serverPort int
	var clusterID string

	version := build.GetInfo().Version

	if opts.ServerMetadata == "" {
		cmdio.LogString(ctx, "Uploading binaries...")
		sp := cmdio.NewSpinner(ctx)
		sp.TrackElapsedTime()
		sp.Update("Uploading binaries...")
		if err := UploadTunnelReleases(ctx, client, version, opts.ReleasesDir); err != nil {
			sp.Close()
			return fmt.Errorf("failed to upload ssh-tunnel binaries: %w", err)
		}
		sp.Close()

		serverStartTime := time.Now()
		userName, serverPort, clusterID, err = ensureSSHServerIsRunning(ctx, client, version, secretScopeName, opts)
		if err != nil {
			return fmt.Errorf("failed to ensure that ssh server is running: %w", err)
		}
		event.ServerStartTimeMs = time.Since(serverStartTime).Milliseconds()
	} else {
		// Metadata format: "<user_name>,<port>,<cluster_id>"
		metadata := strings.Split(opts.ServerMetadata, ",")
		if len(metadata) < 2 {
			return fmt.Errorf("invalid metadata: %s, expected format: <user_name>,<port>[,<cluster_id>]", opts.ServerMetadata)
		}
		userName = metadata[0]
		if userName == "" {
			return fmt.Errorf("remote user name is empty in the metadata: %s", opts.ServerMetadata)
		}
		serverPort, err = strconv.Atoi(metadata[1])
		if err != nil {
			return fmt.Errorf("cannot parse port from metadata: %s, %w", opts.ServerMetadata, err)
		}
		if len(metadata) >= 3 {
			clusterID = metadata[2]
		} else {
			clusterID = opts.ClusterID
		}
	}

	// For serverless mode, we need the cluster ID from metadata for Driver Proxy connections
	if opts.IsServerlessMode() && clusterID == "" {
		return errors.New("cluster ID is required for serverless connections but was not found in metadata")
	}

	log.Infof(ctx, "Remote user name: %s", userName)
	log.Infof(ctx, "Server port: %d", serverPort)
	if opts.IsServerlessMode() {
		log.Infof(ctx, "Cluster ID (from serverless job): %s", clusterID)
	}

	if !opts.ProxyMode {
		cmdio.LogString(ctx, "Connected!")
	}

	// Persist the session for future reconnects.
	if opts.IsServerlessMode() && !opts.ProxyMode {
		err = sessions.Add(ctx, sessions.Session{
			Name:          opts.ConnectionName,
			Accelerator:   opts.Accelerator,
			WorkspaceHost: client.Config.Host,
			CreatedAt:     time.Now(),
			ClusterID:     clusterID,
		})
		if err != nil {
			log.Warnf(ctx, "Failed to save session state: %v", err)
		}
	}

	if opts.ProxyMode {
		return runSSHProxy(ctx, client, serverPort, clusterID, opts)
	} else if opts.IDE != "" {
		return runIDE(ctx, client, userName, keyPath, serverPort, clusterID, opts)
	} else {
		log.Infof(ctx, "Additional SSH arguments: %v", opts.AdditionalArgs)
		return spawnSSHClient(ctx, userName, keyPath, serverPort, clusterID, opts)
	}
}

func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, keyPath string, serverPort int, clusterID string, opts ClientOptions) error {
	connectionName := opts.SessionIdentifier()
	if connectionName == "" {
		return errors.New("connection name is required for IDE integration")
	}

	// Get Databricks user name for the workspace path
	currentUser, err := client.CurrentUser.Me(ctx)
	if err != nil {
		return fmt.Errorf("failed to get current user: %w", err)
	}

	// Ensure SSH config entry exists
	configPath, err := sshconfig.GetMainConfigPath(ctx)
	if err != nil {
		return fmt.Errorf("failed to get SSH config path: %w", err)
	}

	err = ensureSSHConfigEntry(ctx, configPath, connectionName, userName, keyPath, serverPort, clusterID, opts)
	if err != nil {
		return fmt.Errorf("failed to ensure SSH config entry: %w", err)
	}

	return vscode.LaunchIDE(ctx, opts.IDE, connectionName, userName, currentUser.UserName)
}

func ensureSSHConfigEntry(ctx context.Context, configPath, hostName, userName, keyPath string, serverPort int, clusterID string, opts ClientOptions) error {
	// Ensure the Include directive exists in the main SSH config
	err := sshconfig.EnsureIncludeDirective(ctx, configPath)
	if err != nil {
		return err
	}

	// Generate ProxyCommand with server metadata
	optsWithMetadata := opts
	optsWithMetadata.ServerMetadata = FormatMetadata(userName, serverPort, clusterID)

	proxyCommand, err := optsWithMetadata.ToProxyCommand()
	if err != nil {
		return fmt.Errorf("failed to generate ProxyCommand: %w", err)
	}

	var hostConfig string
	if opts.IsServerlessMode() {
		hostConfig = sshconfig.GenerateServerlessHostConfig(hostName, userName, keyPath, proxyCommand)
	} else {
		hostConfig = sshconfig.GenerateHostConfig(hostName, userName, keyPath, proxyCommand)
	}

	_, err = sshconfig.CreateOrUpdateHostConfig(ctx, hostName, hostConfig, true)
	if err != nil {
		return err
	}

	log.Infof(ctx, "Updated SSH config entry for '%s'", hostName)
	return nil
}

// getServerMetadata retrieves the server metadata from the workspace and validates it via Driver Proxy.
// sessionID is the unique identifier for the session (cluster ID for dedicated clusters, connection name for serverless).
// For dedicated clusters, clusterID should be the same as sessionID.
// For serverless, clusterID is read from the workspace metadata.
func getServerMetadata(ctx context.Context, client *databricks.WorkspaceClient, sessionID, clusterID, version, liteswap string) (int, string, string, error) {
	wsMetadata, err := sshWorkspace.GetWorkspaceMetadata(ctx, client, version, sessionID)
	if err != nil {
		return 0, "", "", errors.Join(errServerMetadata, err)
	}
	log.Debugf(ctx, "Workspace metadata: %+v", wsMetadata)

	// For serverless mode, the cluster ID comes from the metadata
	effectiveClusterID := clusterID
	if wsMetadata.ClusterID != "" {
		effectiveClusterID = wsMetadata.ClusterID
	}

	if effectiveClusterID == "" {
		return 0, "", "", errors.Join(errServerMetadata, errors.New("cluster ID not available in metadata"))
	}

	workspaceID, err := client.CurrentWorkspaceID(ctx)
	if err != nil {
		return 0, "", "", err
	}
	metadataURL := fmt.Sprintf("%s/driver-proxy-api/o/%d/%s/%d/metadata", client.Config.Host, workspaceID, effectiveClusterID, wsMetadata.Port)
	log.Debugf(ctx, "Metadata URL: %s", metadataURL)
	req, err := http.NewRequestWithContext(ctx, "GET", metadataURL, nil)
	if err != nil {
		return 0, "", "", err
	}
	if liteswap != "" {
		req.Header.Set("x-databricks-traffic-id", "testenv://liteswap/"+liteswap)
	}
	if err := client.Config.Authenticate(req); err != nil {
		return 0, "", "", err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return 0, "", "", err
	}
	defer resp.Body.Close()

	bodyBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return 0, "", "", err
	}
	log.Debugf(ctx, "Metadata response: %s", string(bodyBytes))
	log.Debugf(ctx, "Metadata response status code: %d", resp.StatusCode)

	if resp.StatusCode != http.StatusOK {
		return 0, "", "", errors.Join(errServerMetadata, fmt.Errorf("server is not ok, status code %d", resp.StatusCode))
	}

	return wsMetadata.Port, string(bodyBytes), effectiveClusterID, nil
}

func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient, version, secretScopeName string, opts ClientOptions) error {
	sessionID := opts.SessionIdentifier()
	contentDir, err := sshWorkspace.GetWorkspaceContentDir(ctx, client, version, sessionID)
	if err != nil {
		return fmt.Errorf("failed to get workspace content directory: %w", err)
	}

	err = client.Workspace.MkdirsByPath(ctx, contentDir)
	if err != nil {
		return fmt.Errorf("failed to create directory in the remote workspace: %w", err)
	}

	sshTunnelJobName := "ssh-server-bootstrap-" + sessionID
	jobNotebookPath := filepath.ToSlash(filepath.Join(contentDir, "ssh-server-bootstrap"))
	notebookContent := "# Databricks notebook source\n" + sshServerBootstrapScript
	encodedContent := base64.StdEncoding.EncodeToString([]byte(notebookContent))

	err = client.Workspace.Import(ctx, workspace.Import{
		Path:      jobNotebookPath,
		Format:    workspace.ImportFormatSource,
		Content:   encodedContent,
		Language:  workspace.LanguagePython,
		Overwrite: true,
	})
	if err != nil {
		return fmt.Errorf("failed to create ssh-tunnel notebook: %w", err)
	}

	baseParams := map[string]string{
		"version":                 version,
		"secretScopeName":         secretScopeName,
		"authorizedKeySecretName": opts.ClientPublicKeyName,
		"shutdownDelay":           opts.ShutdownDelay.String(),
		"maxClients":              strconv.Itoa(opts.MaxClients),
		"sessionId":               sessionID,
		"serverless":              strconv.FormatBool(opts.IsServerlessMode()),
	}

	log.Infof(ctx, "Submitting a job to start the ssh server...")

	task := jobs.SubmitTask{
		TaskKey: sshServerTaskKey,
		NotebookTask: &jobs.NotebookTask{
			NotebookPath:   jobNotebookPath,
			BaseParameters: baseParams,
		},
		TimeoutSeconds: int(opts.ServerTimeout.Seconds()),
	}

	if opts.IsServerlessMode() {
		task.EnvironmentKey = serverlessEnvironmentKey
		if opts.Accelerator != "" {
			log.Infof(ctx, "Using accelerator: %s", opts.Accelerator)
			task.Compute = &jobs.Compute{
				HardwareAccelerator: compute.HardwareAcceleratorType(opts.Accelerator),
			}
		}
	} else {
		task.ExistingClusterId = opts.ClusterID
	}

	submitRequest := jobs.SubmitRun{
		RunName:        sshTunnelJobName,
		TimeoutSeconds: int(opts.ServerTimeout.Seconds()),
		Tasks:          []jobs.SubmitTask{task},
	}

	if opts.IsServerlessMode() {
		submitRequest.Environments = []jobs.JobEnvironment{
			{
				EnvironmentKey: serverlessEnvironmentKey,
				Spec: &compute.Environment{
					EnvironmentVersion: strconv.Itoa(max(opts.EnvironmentVersion, minEnvironmentVersion)),
				},
			},
		}
	}

	waiter, err := client.Jobs.Submit(ctx, submitRequest)
	if err != nil {
		return fmt.Errorf("failed to submit job: %w", err)
	}

	log.Infof(ctx, "Job submitted successfully with run ID: %d", waiter.RunId)

	return waitForJobToStart(ctx, client, waiter.RunId, opts.TaskStartupTimeout)
}

func spawnSSHClient(ctx context.Context, userName, privateKeyPath string, serverPort int, clusterID string, opts ClientOptions) error {
	// Create a copy with metadata for the ProxyCommand
	optsWithMetadata := opts
	optsWithMetadata.ServerMetadata = FormatMetadata(userName, serverPort, clusterID)

	proxyCommand, err := optsWithMetadata.ToProxyCommand()
	if err != nil {
		return fmt.Errorf("failed to generate ProxyCommand: %w", err)
	}

	hostName := opts.SessionIdentifier()

	hostKeyChecking := "StrictHostKeyChecking=accept-new"
	if opts.IsServerlessMode() {
		hostKeyChecking = "StrictHostKeyChecking=no"
	}

	sshArgs := []string{
		"-l", userName,
		"-i", privateKeyPath,
		"-o", "IdentitiesOnly=yes",
		"-o", hostKeyChecking,
		"-o", "ConnectTimeout=360",
		"-o", "ProxyCommand=" + proxyCommand,
	}
	if opts.IsServerlessMode() {
		sshArgs = append(sshArgs, "-o", "UserKnownHostsFile=/dev/null")
	} else if opts.UserKnownHostsFile != "" {
		sshArgs = append(sshArgs, "-o", "UserKnownHostsFile="+opts.UserKnownHostsFile)
	}
	sshArgs = append(sshArgs, hostName)
	sshArgs = append(sshArgs, opts.AdditionalArgs...)

	log.Debugf(ctx, "Launching SSH client: ssh %s", strings.Join(sshArgs, " "))
	sshCmd := exec.CommandContext(ctx, "ssh", sshArgs...)

	sshCmd.Stdin = os.Stdin
	sshCmd.Stdout = os.Stdout
	sshCmd.Stderr = os.Stderr

	return sshCmd.Run()
}

func runSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, serverPort int, clusterID string, opts ClientOptions) error {
	createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) {
		return createWebsocketConnection(ctx, client, connID, clusterID, serverPort, opts.Liteswap)
	}
	requestHandoverTick := func() <-chan time.Time {
		return time.After(opts.HandoverTimeout)
	}
	return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, createConn)
}

func checkClusterState(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, autoStart bool) error {
	sp := cmdio.NewSpinner(ctx)
	sp.TrackElapsedTime()
	defer sp.Close()
	if autoStart {
		sp.Update("Ensuring the cluster is running...")
		err := client.Clusters.EnsureClusterIsRunning(ctx, clusterID)
		if err != nil {
			return fmt.Errorf("failed to ensure that the cluster is running: %w", err)
		}
	} else {
		sp.Update("Checking cluster state...")
		cluster, err := client.Clusters.GetByClusterId(ctx, clusterID)
		if err != nil {
			return fmt.Errorf("failed to get cluster info: %w", err)
		}
		if cluster.State != compute.StateRunning {
			return fmt.Errorf("cluster %s is not running, current state: %s. Use --auto-start-cluster to start it automatically", clusterID, cluster.State)
		}
	}
	return nil
}

// waitForJobToStart polls the task status until the SSH server task is in RUNNING state or terminates.
// Returns an error if the task fails to start or if polling times out.
func waitForJobToStart(ctx context.Context, client *databricks.WorkspaceClient, runID int64, taskStartupTimeout time.Duration) error {
	sp := cmdio.NewSpinner(ctx)
	sp.TrackElapsedTime()
	defer sp.Close()
	sp.Update("Starting SSH server...")
	var prevState jobs.RunLifecycleStateV2State

	_, err := retries.Poll(ctx, taskStartupTimeout, func() (*jobs.RunTask, *retries.Err) {
		run, err := client.Jobs.GetRun(ctx, jobs.GetRunRequest{
			RunId: runID,
		})
		if err != nil {
			return nil, retries.Halt(fmt.Errorf("failed to get job run status: %w", err))
		}

		// Find the SSH server task
		var sshTask *jobs.RunTask
		for i := range run.Tasks {
			if run.Tasks[i].TaskKey == sshServerTaskKey {
				sshTask = &run.Tasks[i]
				break
			}
		}

		if sshTask == nil {
			return nil, retries.Halt(fmt.Errorf("SSH server task '%s' not found in job run", sshServerTaskKey))
		}

		if sshTask.Status == nil {
			return nil, retries.Halt(errors.New("task status is nil"))
		}

		currentState := sshTask.Status.State

		// Update spinner if state changed
		if currentState != prevState {
			sp.Update(fmt.Sprintf("Starting SSH server... (task: %s)", currentState))
			prevState = currentState
		}

		// Check if task is running
		if currentState == jobs.RunLifecycleStateV2StateRunning {
			return sshTask, nil
		}

		// Check for terminal failure states
		if currentState == jobs.RunLifecycleStateV2StateTerminated {
			return nil, retries.Halt(errors.New("task terminated before reaching running state"))
		}

		// Continue polling for other states
		return nil, retries.Continues(fmt.Sprintf("waiting for task to start (current state: %s)", currentState))
	})

	return err
}

func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceClient, version, secretScopeName string, opts ClientOptions) (string, int, string, error) {
	sessionID := opts.SessionIdentifier()
	// For dedicated clusters, use clusterID; for serverless, it will be read from metadata
	clusterID := opts.ClusterID

	serverPort, userName, effectiveClusterID, err := getServerMetadata(ctx, client, sessionID, clusterID, version, opts.Liteswap)
	if errors.Is(err, errServerMetadata) {
		cmdio.LogString(ctx, "Starting SSH server...")

		err := submitSSHTunnelJob(ctx, client, version, secretScopeName, opts)
		if err != nil {
			return "", 0, "", fmt.Errorf("failed to submit and start ssh server job: %w", err)
		}

		sp := cmdio.NewSpinner(ctx)
		sp.TrackElapsedTime()
		sp.Update("Waiting for the SSH server to start...")
		maxRetries := 30
		for retries := range maxRetries {
			if ctx.Err() != nil {
				sp.Close()
				return "", 0, "", ctx.Err()
			}
			serverPort, userName, effectiveClusterID, err = getServerMetadata(ctx, client, sessionID, clusterID, version, opts.Liteswap)
			if err == nil {
				sp.Close()
				break
			} else if retries < maxRetries-1 {
				time.Sleep(2 * time.Second)
			} else {
				sp.Close()
				return "", 0, "", fmt.Errorf("failed to start the ssh server: %w", err)
			}
		}
	} else if err != nil {
		return "", 0, "", err
	}

	return userName, serverPort, effectiveClusterID, nil
}

// resolveServerlessSession handles auto-generation and reconnection for serverless sessions.
// It checks local state for existing sessions matching the workspace and accelerator,
// probes them to see if they're still alive, and prompts the user to reconnect or create new.
func resolveServerlessSession(ctx context.Context, client *databricks.WorkspaceClient, opts *ClientOptions) error {
	version := build.GetInfo().Version

	matching, err := sessions.FindMatching(ctx, client.Config.Host, opts.Accelerator)
	if err != nil {
		log.Warnf(ctx, "Failed to load session state: %v", err)
	}

	// Probe sessions to find alive ones (limit to 5 most recent to avoid latency).
	const maxProbe = 5
	if len(matching) > maxProbe {
		matching = matching[len(matching)-maxProbe:]
	}

	var alive []sessions.Session
	for _, s := range matching {
		_, _, _, probeErr := getServerMetadata(ctx, client, s.Name, s.ClusterID, version, opts.Liteswap)
		if probeErr == nil {
			alive = append(alive, s)
		} else {
			cleanupStaleSession(ctx, client, s, version)
		}
	}

	if len(alive) > 0 && cmdio.IsPromptSupported(ctx) {
		choices := make([]string, 0, len(alive)+1)
		for _, s := range alive {
			choices = append(choices, fmt.Sprintf("Reconnect to %s (started %s)", s.Name, s.CreatedAt.Format(time.RFC822)))
		}
		choices = append(choices, "Create new session")

		choice, choiceErr := cmdio.AskSelect(ctx, "Found existing sessions:", choices)
		if choiceErr != nil {
			return fmt.Errorf("failed to prompt user: %w", choiceErr)
		}

		for i, s := range alive {
			if choice == choices[i] {
				opts.ConnectionName = s.Name
				cmdio.LogString(ctx, "Reconnecting to session: "+s.Name)
				return nil
			}
		}
	}

	// No alive session selected — generate a new name.
	opts.ConnectionName = sessions.GenerateSessionName(opts.Accelerator)
	cmdio.LogString(ctx, "Creating new session: "+opts.ConnectionName)
	return nil
}

// cleanupStaleSession removes all local and remote artifacts for a stale session.
func cleanupStaleSession(ctx context.Context, client *databricks.WorkspaceClient, s sessions.Session, version string) {
	// Remove local SSH keys.
	keyPath, err := keys.GetLocalSSHKeyPath(ctx, s.Name, "")
	if err == nil {
		os.RemoveAll(filepath.Dir(keyPath))
	}

	// Remove SSH config entry.
	if err := sshconfig.RemoveHostConfig(ctx, s.Name); err != nil {
		log.Debugf(ctx, "Failed to remove SSH config for %s: %v", s.Name, err)
	}

	// Delete secret scope (best-effort).
	me, err := client.CurrentUser.Me(ctx)
	if err == nil {
		scopeName := fmt.Sprintf("%s-%s-ssh-tunnel-keys", me.UserName, s.Name)
		deleteErr := client.Secrets.DeleteScope(ctx, workspace.DeleteScope{Scope: scopeName})
		if deleteErr != nil {
			log.Debugf(ctx, "Failed to delete secret scope %s: %v", scopeName, deleteErr)
		}
	}

	// Remove workspace content directory (best-effort).
	contentDir, err := sshWorkspace.GetWorkspaceContentDir(ctx, client, version, s.Name)
	if err == nil {
		deleteErr := client.Workspace.Delete(ctx, workspace.Delete{Path: contentDir, Recursive: true})
		if deleteErr != nil {
			log.Debugf(ctx, "Failed to delete workspace content for %s: %v", s.Name, deleteErr)
		}
	}

	// Remove from local state.
	if err := sessions.Remove(ctx, s.Name); err != nil {
		log.Debugf(ctx, "Failed to remove session %s from state: %v", s.Name, err)
	}

	log.Infof(ctx, "Cleaned up stale session: %s", s.Name)
}