// Package k3d owns the provider-neutral remote SSH/tool/k3d runtime shared by
// the EC2 and GCE provisioning adapters.
package k3d

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

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

const (
	remoteGoVersion   = "1.26.1"
	remoteTiltVersion = "0.37.3"
	remoteHelmVersion = "4.2.3"
	remoteTTLPrefix   = "cp9-expire-"
)

var (
	sshReadyTimeout = 5 * time.Minute
	sshRetryDelay   = 5 * time.Second
)

// RemoteRunner executes one command on the selected remote VM.
type RemoteRunner func(context.Context, *host.Session, []byte, ...string) ([]byte, error)

// Options contains only provider-neutral remote runtime knobs.
type Options struct {
	Agents           int
	Memory           string
	K3sImage         string
	TTL              time.Duration
	SSHUser          string
	SSHPrivateKey    string
	PublicHTTP       bool
	PublicHTTPS      bool
	O11YHostPort     int
	AllowSSMFallback bool
}

// Runtime prepares one already-provisioned remote VM as a k3d host.
type Runtime struct {
	Selected  *host.Session
	Options   Options
	RunRemote RemoteRunner
}

// Prepare is the shared remote-host interface used after the provider adapter
// has persisted an exact VM identity and address.
func (r *Runtime) Prepare(ctx context.Context) error {
	if r == nil || r.Selected == nil || !r.Selected.IsRemote() {
		return fmt.Errorf("remote k3d selected session is required")
	}
	r.ensureRemoteSession()
	if err := r.persist(); err != nil {
		return err
	}
	if err := r.waitForSSH(ctx); err != nil {
		return err
	}
	if err := r.resolveArchitecture(ctx); err != nil {
		return err
	}
	if err := r.resolveCPUCount(ctx); err != nil {
		return err
	}
	if err := r.resolvePaths(ctx); err != nil {
		return err
	}
	if err := r.persist(); err != nil {
		return err
	}
	if err := r.bootstrapTools(ctx); err != nil {
		return err
	}
	if err := host.ResetSSHControlMaster(ctx, r.Selected); err != nil {
		return err
	}
	if err := r.preflight(ctx); err != nil {
		return err
	}
	if err := r.ensureCluster(ctx); err != nil {
		return err
	}
	if err := r.exportKubeconfig(ctx); err != nil {
		return err
	}
	return r.persist()
}

func (r *Runtime) run(ctx context.Context, argv ...string) ([]byte, error) {
	if r.RunRemote != nil {
		return r.RunRemote(ctx, r.Selected, nil, argv...)
	}
	return host.RunSSH(ctx, r.Selected, nil, argv...)
}

func (r *Runtime) runShell(ctx context.Context, script string) ([]byte, error) {
	return r.run(ctx, "sh", "-lc", script)
}

func (r *Runtime) ensureRemoteSession() {
	remote := r.Selected.Remote
	if strings.TrimSpace(remote.SSHUser) == "" {
		remote.SSHUser = r.Options.SSHUser
	}
	if strings.TrimSpace(remote.SSHPrivateKeyPath) == "" {
		remote.SSHPrivateKeyPath = r.Options.SSHPrivateKey
	}
	if strings.TrimSpace(remote.RemoteStateDir) == "" {
		remote.RemoteStateDir = "~/.cp9/" + r.Selected.ClusterName
	}
	if strings.TrimSpace(remote.RemoteKubeconfigPath) == "" {
		remote.RemoteKubeconfigPath = "~/.cp9/" + r.Selected.ClusterName + "/kubeconfig.yaml"
	}
}

func (r *Runtime) persist() error {
	if err := host.Save(r.Selected.StateDir, r.Selected); err != nil {
		return fmt.Errorf("persist remote k3d host: %w", err)
	}
	return nil
}

func (r *Runtime) waitForSSH(ctx context.Context) error {
	deadline := time.Now().Add(sshReadyTimeout)
	var lastErr error
	for time.Now().Before(deadline) {
		probeCtx, cancel := context.WithTimeout(ctx, host.SSHTimeout(r.Selected))
		_, err := r.run(probeCtx, "true")
		probeTimedOut := errors.Is(probeCtx.Err(), context.DeadlineExceeded)
		cancel()
		if err == nil {
			return nil
		}
		if ctx.Err() != nil {
			return fmt.Errorf("wait for remote ssh: %w", ctx.Err())
		}
		lastErr = err
		if r.Options.AllowSSMFallback && probeTimedOut &&
			strings.TrimSpace(r.Selected.Remote.SSHTransport) != host.SSHTransportSSM {
			previous := r.Selected.Remote.SSHTransport
			r.Selected.Remote.SSHTransport = host.SSHTransportSSM
			verbose.Logf(
				"remote direct SSH probe timed out; trying SSM transport profile=%s region=%s instance=%s",
				r.Selected.Remote.Profile,
				r.Selected.Remote.Region,
				r.Selected.Remote.InstanceID,
			)
			ssmCtx, ssmCancel := context.WithTimeout(ctx, host.SSHTimeout(r.Selected))
			_, ssmErr := r.run(ssmCtx, "true")
			ssmCancel()
			if ssmErr == nil {
				return r.persist()
			}
			lastErr = fmt.Errorf("direct probe timed out; SSM fallback: %w", ssmErr)
			r.Selected.Remote.SSHTransport = previous
		}
		if err := sleep(ctx, sshRetryDelay); err != nil {
			return fmt.Errorf("wait for remote ssh: %w", err)
		}
	}
	if err := ctx.Err(); err != nil {
		return fmt.Errorf("wait for remote ssh: %w", err)
	}
	return fmt.Errorf("remote ssh did not become reachable before timeout: %w", lastErr)
}

func sleep(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func (r *Runtime) resolvePaths(ctx context.Context) error {
	remote := r.Selected.Remote
	if strings.TrimSpace(remote.SSHUser) == "" {
		output, err := r.run(ctx, "id", "-un")
		if err != nil {
			return fmt.Errorf("resolve remote user: %w", err)
		}
		remote.SSHUser = strings.TrimSpace(string(output))
		if remote.SSHUser == "" {
			return fmt.Errorf("resolve remote user: remote identity is empty")
		}
	}
	needsStateDir := strings.HasPrefix(remote.RemoteStateDir, "~/")
	needsKubeconfig := strings.HasPrefix(remote.RemoteKubeconfigPath, "~/")
	if !needsStateDir && !needsKubeconfig {
		return nil
	}
	output, err := r.runShell(ctx, `printf '%s' "$HOME"`)
	if err != nil {
		return fmt.Errorf("resolve remote home: %w", err)
	}
	homeDir := strings.TrimSpace(string(output))
	if homeDir == "" {
		return fmt.Errorf("resolve remote home: remote $HOME is empty")
	}
	if needsStateDir {
		remote.RemoteStateDir = path.Join(homeDir, strings.TrimPrefix(remote.RemoteStateDir, "~/"))
	}
	if needsKubeconfig {
		remote.RemoteKubeconfigPath = path.Join(homeDir, strings.TrimPrefix(remote.RemoteKubeconfigPath, "~/"))
	}
	return nil
}

func (r *Runtime) resolveCPUCount(ctx context.Context) error {
	if r.Selected.Remote.CPUCount > 0 {
		return nil
	}
	output, err := r.runShell(ctx, `set -eu
if command -v getconf >/dev/null 2>&1; then
  getconf _NPROCESSORS_ONLN
else
  nproc
fi
`)
	if err != nil {
		return fmt.Errorf("resolve remote CPU count: %w", err)
	}
	value := strings.TrimSpace(string(output))
	count, err := strconv.Atoi(value)
	if err != nil || count < 1 {
		return fmt.Errorf("resolve remote CPU count: invalid value %q", value)
	}
	r.Selected.Remote.CPUCount = count
	return nil
}

func (r *Runtime) resolveArchitecture(ctx context.Context) error {
	arch := strings.TrimSpace(r.Selected.Remote.Architecture)
	if arch == "" {
		output, err := r.runShell(ctx, `uname -m`)
		if err != nil {
			return fmt.Errorf("resolve remote architecture: %w", err)
		}
		arch = strings.TrimSpace(string(output))
	}
	switch arch {
	case "aarch64", "arm64":
		arch = "arm64"
	case "x86_64", "amd64":
		arch = "amd64"
	default:
		return fmt.Errorf("resolve remote architecture: unsupported uname -m %q", arch)
	}
	r.Selected.Remote.Architecture = arch
	if image := r.Selected.Remote.Image; image != nil {
		pinned := strings.TrimSpace(image.Architecture)
		if pinned == "x86_64" {
			pinned = "amd64"
		}
		if pinned != arch {
			return fmt.Errorf("pinned remote image architecture %q does not match runtime architecture %q", image.Architecture, arch)
		}
	}
	return nil
}

func (r *Runtime) bootstrapTools(ctx context.Context) error {
	arch := r.Selected.Remote.Architecture
	awsCLIArch := "aarch64"
	tiltArch := arch
	if arch == "amd64" {
		awsCLIArch = "x86_64"
		tiltArch = "x86_64"
	}
	script := fmt.Sprintf(`set -eu
if ! command -v curl >/dev/null 2>&1; then
  if command -v apt-get >/dev/null 2>&1; then sudo apt-get update && sudo apt-get install -y curl ca-certificates;
  elif command -v dnf >/dev/null 2>&1; then sudo dnf install -y curl ca-certificates;
  else sudo yum install -y curl ca-certificates; fi
fi
if ! command -v openssl >/dev/null 2>&1; then
  if command -v apt-get >/dev/null 2>&1; then sudo apt-get update && sudo apt-get install -y openssl;
  elif command -v dnf >/dev/null 2>&1; then sudo dnf install -y openssl;
  else sudo yum install -y openssl; fi
fi
if ! command -v git >/dev/null 2>&1; then
  if command -v apt-get >/dev/null 2>&1; then sudo apt-get update && sudo apt-get install -y git;
  elif command -v dnf >/dev/null 2>&1; then sudo dnf install -y git;
  else sudo yum install -y git; fi
fi
if ! command -v python3 >/dev/null 2>&1 || ! command -v unzip >/dev/null 2>&1; then
  if command -v apt-get >/dev/null 2>&1; then sudo apt-get update && sudo apt-get install -y python3 unzip;
  elif command -v dnf >/dev/null 2>&1; then sudo dnf install -y python3 unzip;
  else sudo yum install -y python3 unzip; fi
fi
if ! command -v rsync >/dev/null 2>&1; then
  if command -v apt-get >/dev/null 2>&1; then sudo apt-get update && sudo apt-get install -y rsync;
  elif command -v dnf >/dev/null 2>&1; then sudo dnf install -y rsync;
  else sudo yum install -y rsync; fi
fi
if ! command -v aws >/dev/null 2>&1 || ! aws --version >/dev/null 2>&1; then
  rm -rf /tmp/cp9-aws-cli /tmp/cp9-awscliv2.zip
  sudo rm -rf /usr/local/aws-cli
  sudo rm -f /usr/local/bin/aws /usr/local/bin/aws_completer
  curl --retry 5 --retry-delay 2 --retry-all-errors -fsSL -o /tmp/cp9-awscliv2.zip https://awscli.amazonaws.com/awscli-exe-linux-%s.zip
  mkdir -p /tmp/cp9-aws-cli
  unzip -q /tmp/cp9-awscliv2.zip -d /tmp/cp9-aws-cli
  sudo sh /tmp/cp9-aws-cli/aws/install --bin-dir /usr/local/bin --install-dir /usr/local/aws-cli
  rm -rf /tmp/cp9-aws-cli /tmp/cp9-awscliv2.zip
fi
if ! command -v docker >/dev/null 2>&1; then
  if command -v apt-get >/dev/null 2>&1; then sudo apt-get update && sudo apt-get install -y docker.io;
  elif command -v dnf >/dev/null 2>&1; then sudo dnf install -y docker;
  else sudo yum install -y docker; fi
  sudo systemctl enable --now docker
  sudo usermod -aG docker "$USER" || true
fi
if ! command -v k3d >/dev/null 2>&1; then
  curl --retry 5 --retry-delay 2 --retry-all-errors -fsSL https://raw.githubusercontent.com/k3d-io/k3d/main/install.sh | bash
fi
if ! command -v go >/dev/null 2>&1 || ! go version >/dev/null 2>&1; then
  curl --retry 5 --retry-delay 2 --retry-all-errors -fsSL -o /tmp/cp9-go.tgz https://go.dev/dl/go%s.linux-%s.tar.gz
  sudo rm -rf /usr/local/go
  sudo tar -C /usr/local -xzf /tmp/cp9-go.tgz
  sudo ln -sf /usr/local/go/bin/go /usr/local/bin/go
  sudo ln -sf /usr/local/go/bin/gofmt /usr/local/bin/gofmt
fi
if ! command -v kubectl >/dev/null 2>&1 || ! kubectl version --client=true >/dev/null 2>&1; then
  curl --retry 5 --retry-delay 2 --retry-all-errors -fsSL -o kubectl https://dl.k8s.io/release/$(curl --retry 5 --retry-delay 2 --retry-all-errors -fsSL https://dl.k8s.io/release/stable.txt)/bin/linux/%s/kubectl
  chmod +x kubectl
  sudo mv kubectl /usr/local/bin/kubectl
fi
if ! command -v tilt >/dev/null 2>&1 || ! tilt version 2>/dev/null | grep -Fq "v%s"; then
  curl --retry 5 --retry-delay 2 --retry-all-errors -fsSL -o /tmp/cp9-tilt.tgz https://github.com/tilt-dev/tilt/releases/download/v%s/tilt.%s.linux.%s.tar.gz
  tar -C /tmp -xzf /tmp/cp9-tilt.tgz tilt
  sudo install -m 0755 /tmp/tilt /usr/local/bin/tilt
fi
if ! command -v helm >/dev/null 2>&1 || [ "$(helm version --short 2>/dev/null | sed 's/^v//; s/+.*//')" != "%s" ]; then
  curl --retry 5 --retry-delay 2 --retry-all-errors -fsSL -o /tmp/cp9-helm.tgz https://get.helm.sh/helm-v%s-linux-%s.tar.gz
  tar -C /tmp -xzf /tmp/cp9-helm.tgz linux-%s/helm
  sudo install -m 0755 /tmp/linux-%s/helm /usr/local/bin/helm
fi
printf 'MaxSessions 64\n' | sudo tee /etc/ssh/sshd_config.d/99-cp9-max-sessions.conf >/dev/null
sudo sshd -t
sudo systemctl reload sshd 2>/dev/null || sudo systemctl reload ssh
sudo sysctl -w fs.inotify.max_user_instances=8192 fs.inotify.max_user_watches=1048576 >/dev/null
printf '%%s\n' 'fs.inotify.max_user_instances=8192' 'fs.inotify.max_user_watches=1048576' | sudo tee /etc/sysctl.d/99-cp9-k3d-inotify.conf >/dev/null
mkdir -p %s
`, awsCLIArch, remoteGoVersion, arch, arch,
		remoteTiltVersion, remoteTiltVersion, remoteTiltVersion, tiltArch,
		remoteHelmVersion, remoteHelmVersion, arch, arch, arch,
		host.ShellEscape(r.Selected.Remote.RemoteStateDir))
	if _, err := r.runShell(ctx, script); err != nil {
		return fmt.Errorf("bootstrap remote tools: %w", err)
	}
	if r.Options.TTL <= 0 {
		return nil
	}
	unit := remoteTTLPrefix + sanitizeUnitName(r.Selected.ClusterName)
	shutdown := fmt.Sprintf(
		"sudo systemctl stop %[1]s.timer %[1]s.service >/dev/null 2>&1 || true\n"+
			"sudo systemctl reset-failed %[1]s.timer %[1]s.service >/dev/null 2>&1 || true\n"+
			"sudo systemd-run --unit %[1]s --on-active=%[2]d /sbin/shutdown -h now",
		host.ShellEscape(unit),
		int(r.Options.TTL.Seconds()),
	)
	if _, err := r.runShell(ctx, shutdown); err != nil {
		return fmt.Errorf("schedule remote ttl shutdown: %w", err)
	}
	return nil
}

func (r *Runtime) preflight(ctx context.Context) error {
	const script = `set -eu
df -BG /
df -BG / | awk 'NR==2 {gsub(/G/, "", $4); if (($4 + 0) < 16) { printf "remote root filesystem has %s GiB free; need at least 16 GiB\n", $4 > "/dev/stderr"; exit 1 } }'
docker info >/dev/null
go version >/dev/null
aws --version >/dev/null
docker pull public.ecr.aws/docker/library/alpine:3.19 >/dev/null
docker image inspect public.ecr.aws/docker/library/alpine:3.19 >/dev/null
preflight_dir="$(mktemp -d)"
trap 'rm -rf "$preflight_dir"' EXIT
printf '%s\n' 'FROM public.ecr.aws/docker/library/alpine:3.19' 'RUN true' > "$preflight_dir/Dockerfile"
docker build -t cp9-remote-k3d-preflight:alpine-3.19 "$preflight_dir" >/dev/null
docker image rm cp9-remote-k3d-preflight:alpine-3.19 >/dev/null 2>&1 || true
`
	if _, err := r.runShell(ctx, script); err != nil {
		return fmt.Errorf("preflight remote docker/disk: %w", err)
	}
	return nil
}

func (r *Runtime) ensureCluster(ctx context.Context) error {
	remote := r.Selected.Remote
	volume := remote.RemoteStateDir + ":" + stub.RecordingNodeStateDir + "@all"
	args := []string{
		"k3d", "cluster", "create", r.Selected.ClusterName,
		"--k3s-arg", "--disable=traefik@server:0",
		"--runtime-ulimit", "nofile=1048576:1048576",
		"--wait",
		"--volume", volume,
	}
	if !r.Options.PublicHTTP && !r.Options.PublicHTTPS && r.Options.O11YHostPort == 0 && r.Selected.DedicatedAPI == nil {
		args = append(args, "--no-lb")
	}
	if r.Options.PublicHTTP {
		args = append(args, "--port", "80:80@server:0")
	}
	if r.Options.PublicHTTPS {
		args = append(args, "--port", "443:443@server:0")
	}
	if r.Options.O11YHostPort > 0 {
		args = append(args, "--port", fmt.Sprintf("%d:%d@server:0", r.Options.O11YHostPort, r.Options.O11YHostPort))
	}
	if api := r.Selected.DedicatedAPI; api != nil {
		if api.HostPort < 1 || api.HostPort > 65535 {
			return fmt.Errorf("prepare remote Dedicated API port mapping: invalid host port %d", api.HostPort)
		}
		if api.NodePort < 30000 || api.NodePort > 32767 {
			return fmt.Errorf("prepare remote Dedicated API port mapping: invalid NodePort %d", api.NodePort)
		}
		args = append(args, "--port", fmt.Sprintf("127.0.0.1:%d:%d@server:0", api.HostPort, api.NodePort))
	}
	if r.Options.Agents > 0 {
		args = append(args, "--agents", strconv.Itoa(r.Options.Agents))
	}
	if r.Options.Memory != "" {
		args = append(args, "--servers-memory", r.Options.Memory)
	}
	if r.Options.K3sImage != "" {
		args = append(args, "--image", r.Options.K3sImage)
	}
	check := fmt.Sprintf(`set -eu
if k3d cluster list --no-headers | awk '{print $1}' | grep -qx %s; then
  k3d kubeconfig get %s > %s
  exit 0
fi
%s
k3d kubeconfig get %s > %s
`, host.ShellEscape(r.Selected.ClusterName), host.ShellEscape(r.Selected.ClusterName), host.ShellEscape(remote.RemoteKubeconfigPath),
		host.ShellEscapeArgs(args...), host.ShellEscape(r.Selected.ClusterName), host.ShellEscape(remote.RemoteKubeconfigPath))
	if _, err := r.runShell(ctx, check); err != nil {
		return fmt.Errorf("ensure remote k3d cluster: %w", err)
	}
	return nil
}

func (r *Runtime) exportKubeconfig(ctx context.Context) error {
	output, err := r.run(ctx, "cat", r.Selected.Remote.RemoteKubeconfigPath)
	if err != nil {
		return fmt.Errorf("read remote kubeconfig: %w", err)
	}
	if err := os.MkdirAll(filepath.Dir(r.Selected.KubeconfigPath), 0o755); err != nil {
		return fmt.Errorf("create kubeconfig dir: %w", err)
	}
	if err := os.WriteFile(r.Selected.KubeconfigPath, output, 0o600); err != nil {
		return fmt.Errorf("write local kubeconfig mirror: %w", err)
	}
	return nil
}

func sanitizeUnitName(name string) string {
	name = strings.ToLower(name)
	return strings.NewReplacer("/", "-", ".", "-", "_", "-").Replace(name)
}