//go:build e2e

package harness

import (
	"context"
	"fmt"
	"maps"
	"os/exec"
	"regexp"
	"slices"
	"strconv"
	"strings"
	"testing"
	"time"
)

// Guest-side data-durability helpers shared by the volume-durability,
// snapshot-restore, and CreateImage e2e tests. They drive a real guest over SSH
// to format, write, sync, and checksum bytes so the assembled
// QEMU↔viperblock↔predastore I/O path is exercised end-to-end — something the
// per-layer unit tests cannot cover.

// guestSentinelFile is the file written into a freshly formatted data volume.
const guestSentinelFile = "e2e-sentinel.bin"

// guestExecTimeout bounds a single guest command by default (mkfs/dd are
// sub-second on the small payloads most callers use; the ceiling only guards
// against a hung SSH). Callers driving a larger payload — enough data that
// sustaining it within this ceiling would require an unrealistic transfer
// rate — must use GuestExecTimeout with a budget sized to their own payload
// instead of assuming this default is high enough.
const guestExecTimeout = 2 * time.Minute

// sha256RE matches a bare sha256 digest so the checksum can be lifted out of
// command output that may also carry the file path (`<sha>␠␠<path>`).
var sha256RE = regexp.MustCompile(`\b[0-9a-f]{64}\b`)

// GuestExec runs cmd over SSH against tgt with the default guestExecTimeout
// budget and returns combined stdout+stderr plus the run error. It never
// calls t.Fatal, so callers can branch on an expected non-zero exit.
func GuestExec(tgt SSHTarget, cmd string) (string, error) {
	return GuestExecTimeout(tgt, cmd, guestExecTimeout)
}

// GuestExecTimeout is GuestExec with a caller-supplied timeout, for commands
// whose payload size makes the default guestExecTimeout unrealistic (e.g. a
// multi-gigabyte write against a network-backed volume). If the command is
// still running when timeout elapses, the returned error is wrapped with
// context.DeadlineExceeded so callers can distinguish "the command ran and
// failed" from "the command never got the chance to finish" via errors.Is —
// a distinction that matters wherever the two would otherwise look like the
// same red test for different reasons.
func GuestExecTimeout(tgt SSHTarget, cmd string, timeout time.Duration) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	out, err := exec.CommandContext(ctx, "ssh", guestSSHArgs(tgt, 5, cmd)...).CombinedOutput()
	if err != nil && ctx.Err() == context.DeadlineExceeded {
		return string(out), fmt.Errorf("guest command exceeded its %s timeout (still running when the deadline hit): %w: %w", timeout, context.DeadlineExceeded, err)
	}
	return string(out), err
}

// guestDiskSet returns the set of whole-disk device names (lsblk TYPE=disk)
// currently visible in the guest, keyed by bare name (e.g. "vda", "vdc").
func guestDiskSet(tgt SSHTarget) (map[string]struct{}, error) {
	out, err := GuestExec(tgt, "lsblk -dn -o NAME,TYPE")
	if err != nil {
		return nil, fmt.Errorf("lsblk: %w\n%s", err, out)
	}
	set := map[string]struct{}{}
	for _, line := range strings.Split(strings.TrimSpace(out), "\n") {
		f := strings.Fields(line)
		if len(f) >= 2 && f[1] == "disk" {
			set[f[0]] = struct{}{}
		}
	}
	return set, nil
}

// GuestDiskSet is the t.Fatal-on-error wrapper around guestDiskSet, used to
// snapshot the guest's disk inventory before an AttachVolume.
func GuestDiskSet(t *testing.T, tgt SSHTarget) map[string]struct{} {
	t.Helper()
	set, err := guestDiskSet(tgt)
	if err != nil {
		t.Fatalf("GuestDiskSet: %v", err)
	}
	return set
}

// WaitForNewGuestDisk polls the guest until a whole-disk device absent from
// `before` appears, returning its bare name (e.g. "vdc"). Proves a hotplugged
// volume actually reached the guest kernel, not just the EC2 control plane.
func WaitForNewGuestDisk(t *testing.T, tgt SSHTarget, before map[string]struct{}, timeout time.Duration) string {
	t.Helper()
	var found string
	EventuallyErr(t, func() error {
		now, err := guestDiskSet(tgt)
		if err != nil {
			return err
		}
		for name := range now {
			if _, ok := before[name]; !ok {
				found = name
				return nil
			}
		}
		names := slices.Collect(maps.Keys(now))
		return fmt.Errorf("no new disk yet (visible: %s)", strings.Join(names, ","))
	}, timeout, 2*time.Second)
	return found
}

// LsblkDeviceBytes returns the raw byte size of /dev/<dev> as the guest kernel
// sees it. Byte-exact rather than GiB-rounded (unlike LsblkRootGiB) so a resize
// that lands short by less than 1 GiB still fails a caller's comparison.
func LsblkDeviceBytes(t *testing.T, tgt SSHTarget, dev string) int64 {
	t.Helper()
	out, err := GuestExec(tgt, fmt.Sprintf("lsblk -b -d -n -o SIZE /dev/%s", dev))
	if err != nil {
		t.Fatalf("LsblkDeviceBytes(%s): %v\n%s", dev, err, out)
	}
	// Read the last non-empty line rather than the whole output: an SSH banner
	// or sudo notice ahead of the size would otherwise break the parse.
	lines := strings.Split(strings.TrimSpace(out), "\n")
	raw := strings.TrimSpace(lines[len(lines)-1])
	size, err := strconv.ParseInt(raw, 10, 64)
	if err != nil {
		t.Fatalf("LsblkDeviceBytes(%s): parse %q: %v\nfull output:\n%s", dev, raw, err, out)
	}
	return size
}

// GuestFormatWriteSentinel formats /dev/<dev> as ext4 with the given label,
// mounts it, writes sizeMiB of random data to the sentinel file, fsyncs, and
// returns the file's sha256. Unmounts before returning so the volume can be
// detached cleanly.
func GuestFormatWriteSentinel(t *testing.T, tgt SSHTarget, dev, label string, sizeMiB int) string {
	t.Helper()
	mnt := "/mnt/" + label
	script := strings.Join([]string{
		fmt.Sprintf("sudo mkfs.ext4 -F -L %s /dev/%s >/dev/null 2>&1", label, dev),
		fmt.Sprintf("sudo mkdir -p %s", mnt),
		fmt.Sprintf("sudo mount /dev/%s %s", dev, mnt),
		fmt.Sprintf("sudo dd if=/dev/urandom of=%s/%s bs=1M count=%d conv=fsync status=none", mnt, guestSentinelFile, sizeMiB),
		"sync",
		fmt.Sprintf("sudo sha256sum %s/%s", mnt, guestSentinelFile),
		fmt.Sprintf("sudo umount %s", mnt),
	}, " && ")
	out, err := GuestExec(tgt, script)
	if err != nil {
		t.Fatalf("GuestFormatWriteSentinel(%s,%s): %v\n%s", dev, label, err, out)
	}
	return mustSha(t, out)
}

// GuestReadSentinelSha mounts source (a device path such as "/dev/vdc" or a
// "/dev/disk/by-label/<label>" path) at a temp mountpoint, returns the
// sentinel file's sha256, and unmounts.
func GuestReadSentinelSha(t *testing.T, tgt SSHTarget, source, label string) string {
	t.Helper()
	mnt := "/mnt/" + label
	script := strings.Join([]string{
		fmt.Sprintf("sudo mkdir -p %s", mnt),
		fmt.Sprintf("sudo mount %s %s", source, mnt),
		fmt.Sprintf("sudo sha256sum %s/%s", mnt, guestSentinelFile),
		fmt.Sprintf("sudo umount %s", mnt),
	}, " && ")
	out, err := GuestExec(tgt, script)
	if err != nil {
		t.Fatalf("GuestReadSentinelSha(%s): %v\n%s", source, err, out)
	}
	return mustSha(t, out)
}

// mustSha extracts the single sha256 digest from command output, failing the
// test if none is present.
func mustSha(t *testing.T, out string) string {
	t.Helper()
	sha := sha256RE.FindString(out)
	if sha == "" {
		t.Fatalf("no sha256 digest in guest output:\n%s", out)
	}
	return sha
}