package host

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"strings"
	"time"

	"github.com/tidbcloud/controlplane9/internal/verbose"
)

var sshCommandRunner = verbose.RunCommandStdout

const (
	directSSHTimeout    = 10 * time.Second
	ssmSSHTimeout       = 30 * time.Second
	gcloudIAPSSHTimeout = 60 * time.Second
)

// SSHTimeout returns the command and tunnel startup budget for the selected
// transport. SSM includes the Session Manager handshake before SSH starts.
func SSHTimeout(selected *Session) time.Duration {
	if selected != nil && selected.Remote != nil {
		switch strings.TrimSpace(selected.Remote.SSHTransport) {
		case SSHTransportSSM:
			return ssmSSHTimeout
		case SSHTransportGCloudIAP:
			return gcloudIAPSSHTimeout
		}
	}
	return directSSHTimeout
}

// SSHBaseArgs returns the common ssh CLI arguments for a persisted remote session.
func SSHBaseArgs(selected *Session) ([]string, error) {
	return sshBaseArgs(selected, false)
}

// SSHTransferBaseArgs returns reusable ssh arguments for short-lived file and
// image transfers. Direct sessions share one control connection.
func SSHTransferBaseArgs(selected *Session) ([]string, error) {
	return sshBaseArgs(selected, true)
}

func sshBaseArgs(selected *Session, multiplex bool) ([]string, error) {
	if selected == nil || selected.Remote == nil {
		return nil, fmt.Errorf("remote host metadata is not available")
	}
	if strings.TrimSpace(selected.Remote.SSHTransport) == SSHTransportGCloudIAP {
		return nil, fmt.Errorf("gcloud-IAP transport does not expose raw ssh arguments")
	}
	if strings.TrimSpace(selected.Remote.SSHUser) == "" {
		return nil, fmt.Errorf("remote ssh user is not configured")
	}
	if strings.TrimSpace(selected.Remote.SSHPrivateKeyPath) == "" {
		return nil, fmt.Errorf("remote ssh private key path is not configured")
	}
	host := firstNonEmpty(selected.Remote.PublicDNS, selected.Remote.PublicIP)
	if host == "" {
		return nil, fmt.Errorf("remote ssh host is not configured")
	}
	// #93 deliberately targets ephemeral per-session dev EC2 hosts inside the
	// dev VPC, so bootstrap ergonomics currently win over host-key persistence.
	args := []string{
		"-o", "BatchMode=yes",
		"-o", "StrictHostKeyChecking=no",
		"-o", "UserKnownHostsFile=/dev/null",
		"-o", "LogLevel=ERROR",
		"-o", fmt.Sprintf("ConnectTimeout=%d", int(SSHTimeout(selected)/time.Second)),
		"-o", "ConnectionAttempts=3",
		"-o", "ServerAliveInterval=30",
		"-o", "ServerAliveCountMax=3",
	}
	switch transport := firstNonEmpty(selected.Remote.SSHTransport, SSHTransportDirect); transport {
	case SSHTransportDirect:
		if multiplex {
			sessionName := strings.TrimSpace(selected.ClusterName)
			if sessionName == "" {
				return nil, fmt.Errorf("remote session cluster name is not configured")
			}
			args = append(args,
				"-o", "ControlMaster=auto",
				"-o", "ControlPersist=60s",
				"-o", "ControlPath=/tmp/"+sessionName+"-ssh-%i-%C",
			)
		} else {
			args = append(args, "-o", "ControlMaster=no", "-o", "ControlPath=none")
		}
	case SSHTransportSSM:
		args = append(args, "-o", "ControlMaster=no", "-o", "ControlPath=none")
		profile := firstNonEmpty(selected.Remote.Profile, selected.Profile)
		region := firstNonEmpty(selected.Remote.Region, selected.Region)
		instanceID := strings.TrimSpace(selected.Remote.InstanceID)
		if profile == "" || region == "" || instanceID == "" {
			return nil, fmt.Errorf("SSM ssh transport requires profile, region, and instance ID")
		}
		proxyCommand := fmt.Sprintf(
			"ProxyCommand=aws ssm start-session --profile %s --region %s --target %s --document-name AWS-StartSSHSession --parameters portNumber=%%p",
			shellEscape(profile),
			shellEscape(region),
			shellEscape(instanceID),
		)
		args = append(args, "-o", proxyCommand)
	default:
		return nil, fmt.Errorf("unsupported remote ssh transport %q", transport)
	}
	return append(args,
		"-i", selected.Remote.SSHPrivateKeyPath,
		selected.Remote.SSHUser+"@"+host,
	), nil
}

// SSHRemoteCommand returns the single remote command string that should be
// passed to ssh after the host argument.
func SSHRemoteCommand(argv ...string) string {
	return SSHRemoteShellCommand(shellEscapeArgs(argv...))
}

// SSHRemoteShellCommand returns the single remote shell command string that
// should be passed to ssh after the host argument.
func SSHRemoteShellCommand(script string) string {
	return "sh -lc " + shellEscape(script)
}

// SSHCommandArgs returns the local ssh argv for one remote multi-arg command.
func SSHCommandArgs(selected *Session, argv ...string) ([]string, error) {
	return sshCommandArgs(selected, true, argv...)
}

func sshCommandArgs(selected *Session, multiplex bool, argv ...string) ([]string, error) {
	base, err := sshBaseArgs(selected, multiplex)
	if err != nil {
		return nil, err
	}
	args := append([]string{}, base...)
	args = append(args, SSHRemoteCommand(argv...))
	return args, nil
}

// RunSSH executes a remote command over ssh.
func RunSSH(ctx context.Context, selected *Session, stdin []byte, argv ...string) ([]byte, error) {
	return runSSH(ctx, selected, stdin, true, argv...)
}

// RunSSHIsolated executes a remote command outside the shared transfer connection.
func RunSSHIsolated(ctx context.Context, selected *Session, stdin []byte, argv ...string) ([]byte, error) {
	return runSSH(ctx, selected, stdin, false, argv...)
}

func runSSH(ctx context.Context, selected *Session, stdin []byte, multiplex bool, argv ...string) ([]byte, error) {
	if selected != nil && selected.Remote != nil &&
		strings.TrimSpace(selected.Remote.SSHTransport) == SSHTransportGCloudIAP {
		return runGCloudSSH(ctx, selected, stdin, argv...)
	}
	args, err := sshCommandArgs(selected, multiplex, argv...)
	if err != nil {
		return nil, err
	}
	cmd := exec.CommandContext(ctx, "ssh", args...)
	if len(stdin) > 0 {
		cmd.Stdin = bytes.NewReader(stdin)
	}
	output, err := sshCommandRunner(cmd)
	if err != nil {
		return nil, fmt.Errorf("ssh %s: %w\n%s", strings.Join(argv, " "), err, string(output))
	}
	return output, nil
}

func runGCloudSSH(ctx context.Context, selected *Session, stdin []byte, argv ...string) ([]byte, error) {
	instance, project, zone, err := gcloudIAPTarget(selected)
	if err != nil {
		return nil, err
	}
	args := []string{
		"compute", "ssh", instance,
		"--project=" + project,
		"--zone=" + zone,
		"--tunnel-through-iap",
		"--quiet",
		"--command=" + SSHRemoteCommand(argv...),
	}
	cmd := exec.CommandContext(ctx, "gcloud", args...)
	if len(stdin) > 0 {
		cmd.Stdin = bytes.NewReader(stdin)
	}
	output, err := sshCommandRunner(cmd)
	if err != nil {
		return nil, fmt.Errorf("gcloud compute ssh %s: %w\n%s", strings.Join(argv, " "), err, string(output))
	}
	return output, nil
}

// CopyFileToGCloudIAP copies one local file to the selected GCE host.
func CopyFileToGCloudIAP(ctx context.Context, selected *Session, localPath, remotePath string) error {
	if strings.TrimSpace(localPath) == "" || strings.TrimSpace(remotePath) == "" {
		return fmt.Errorf("gcloud-IAP file copy requires local and remote paths")
	}
	if selected == nil || selected.Remote == nil ||
		strings.TrimSpace(selected.Remote.SSHTransport) != SSHTransportGCloudIAP {
		return fmt.Errorf("gcloud-IAP file copy requires a GCE host")
	}
	instance, project, zone, err := gcloudIAPTarget(selected)
	if err != nil {
		return err
	}
	args := []string{
		"compute", "scp", localPath, instance + ":" + remotePath,
		"--project=" + project,
		"--zone=" + zone,
		"--tunnel-through-iap",
		"--quiet",
	}
	output, err := sshCommandRunner(exec.CommandContext(ctx, "gcloud", args...))
	if err != nil {
		return fmt.Errorf("gcloud compute scp: %w\n%s", err, string(output))
	}
	return nil
}

func gcloudIAPTarget(selected *Session) (instance, project, zone string, err error) {
	if selected == nil || selected.Remote == nil || selected.GCP == nil {
		return "", "", "", fmt.Errorf("gcloud-IAP host metadata is not available")
	}
	instance = strings.TrimSpace(selected.Remote.InstanceName)
	project = strings.TrimSpace(selected.GCP.ProjectID)
	zone = strings.TrimSpace(selected.GCP.Zone)
	if instance == "" || project == "" || zone == "" {
		return "", "", "", fmt.Errorf("gcloud-IAP transport requires instance name, project, and zone")
	}
	if user := strings.TrimSpace(selected.Remote.SSHUser); user != "" {
		instance = user + "@" + instance
	}
	return instance, project, zone, nil
}

// ResetSSHControlMaster forces the next direct SSH command to pick up remote
// group membership changes made during bootstrap.
func ResetSSHControlMaster(ctx context.Context, selected *Session) error {
	if selected == nil || selected.Remote == nil ||
		firstNonEmpty(selected.Remote.SSHTransport, SSHTransportDirect) != SSHTransportDirect {
		return nil
	}
	args, err := sshBaseArgs(selected, true)
	if err != nil {
		return err
	}
	target := args[len(args)-1]
	args = append(args[:len(args)-1], "-O", "exit", target)
	output, err := sshCommandRunner(exec.CommandContext(ctx, "ssh", args...))
	if err != nil && !strings.Contains(string(output), "No such file or directory") {
		return fmt.Errorf("reset SSH control master: %w\n%s", err, string(output))
	}
	return nil
}

func shellEscapeArgs(argv ...string) string {
	escaped := make([]string, 0, len(argv))
	for _, arg := range argv {
		escaped = append(escaped, shellEscape(arg))
	}
	return strings.Join(escaped, " ")
}

func shellEscape(value string) string {
	if value == "" {
		return "''"
	}
	return "'" + strings.ReplaceAll(value, "'", `'"'"'`) + "'"
}

// ShellEscapeArgs exposes the package ssh/shell escaping for other host-aware callers.
func ShellEscapeArgs(argv ...string) string {
	return shellEscapeArgs(argv...)
}

// ShellEscape exposes single-argument shell escaping for other host-aware callers.
func ShellEscape(value string) string {
	return shellEscape(value)
}

func firstNonEmpty(values ...string) string {
	for _, value := range values {
		value = strings.TrimSpace(value)
		if value != "" {
			return value
		}
	}
	return ""
}