//go:build e2e

package harness

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
)

// SSHTarget is the minimal addressing tuple used by LsblkRootGiB. Unlike
// harness.SSH / Cluster.Node, instances here are addressed by (host, port, key)
// discovered at runtime, which doesn't fit the Cluster model.
type SSHTarget struct {
	User    string
	Host    string
	Port    int
	KeyPath string
}

// InstancePublicSSHHost returns (host, 22) for an SSH connection to inst's
// public IP. No qemu-hostfwd fallback: hostfwd bypasses the OVN datapath (no
// SG ACL, no IGW, no SNAT/DNAT), masking the networking regressions these
// tests exist to catch. No public IP is a hard failure.
func InstancePublicSSHHost(t *testing.T, inst *ec2.Instance) (string, int) {
	t.Helper()
	if inst == nil {
		t.Fatalf("InstancePublicSSHHost: nil instance")
	}
	pub := aws.StringValue(inst.PublicIpAddress)
	if pub == "" || pub == "None" {
		t.Fatalf("InstancePublicSSHHost: instance %s has no public IP; "+
			"qemu-hostfwd fallback is disabled (it bypasses the OVN datapath)",
			aws.StringValue(inst.InstanceId))
	}
	return pub, 22
}

// InstancePrivateIP returns inst's primary private IPv4, or "" when unset or
// unresolvable. Used to populate VPCDiagnosticsOpts.LogicalIP so the datapath
// captures grep both halves of the NAT translation. Non-fatal — it feeds a
// best-effort diagnostic, so a describe miss logs and returns "".
func InstancePrivateIP(t *testing.T, c *AWSClient, instanceID string) string {
	t.Helper()
	out, err := c.EC2.DescribeInstances(&ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(instanceID)},
	})
	if err != nil {
		t.Logf("InstancePrivateIP: describe %s: %v", instanceID, err)
		return ""
	}
	for _, r := range out.Reservations {
		for _, in := range r.Instances {
			return aws.StringValue(in.PrivateIpAddress)
		}
	}
	return ""
}

// LsblkRootGiB SSHes into the VM and returns the root disk size in GiB
// (bytes / 1<<30, rounded down) by running findmnt + lsblk in the guest.
func LsblkRootGiB(t *testing.T, tgt SSHTarget) int {
	t.Helper()
	cmd := `SRC=$(findmnt -n -o SOURCE /); PKN=$(lsblk -n -o PKNAME "$SRC" 2>/dev/null | head -1); DEV=${PKN:-$(basename "$SRC")}; lsblk -b -d -n -o SIZE "/dev/$DEV"`
	raw := strings.TrimSpace(RunSSH(t, tgt, cmd))
	bytesN, err := strconv.ParseInt(raw, 10, 64)
	if err != nil {
		t.Fatalf("ssh lsblk %s@%s:%d: parse %q: %v", tgt.User, tgt.Host, tgt.Port, raw, err)
	}
	return int(bytesN / (1 << 30))
}

// GetSerialConsoleAccessEnabled returns the current account attribute value.
// Wraps ec2.GetSerialConsoleAccessStatus so callers don't have to deal with
// the pointer-bool field.
func GetSerialConsoleAccessEnabled(t *testing.T, c *AWSClient) bool {
	t.Helper()
	out, err := c.EC2.GetSerialConsoleAccessStatus(&ec2.GetSerialConsoleAccessStatusInput{})
	if err != nil {
		t.Fatalf("GetSerialConsoleAccessStatus: %v", err)
	}
	return aws.BoolValue(out.SerialConsoleAccessEnabled)
}

// EnableSerialConsoleAccess flips the account attribute on and returns the
// new value (should be true).
func EnableSerialConsoleAccess(t *testing.T, c *AWSClient) bool {
	t.Helper()
	out, err := c.EC2.EnableSerialConsoleAccess(&ec2.EnableSerialConsoleAccessInput{})
	if err != nil {
		t.Fatalf("EnableSerialConsoleAccess: %v", err)
	}
	return aws.BoolValue(out.SerialConsoleAccessEnabled)
}

// DisableSerialConsoleAccess flips the account attribute off and returns the
// new value (should be false).
func DisableSerialConsoleAccess(t *testing.T, c *AWSClient) bool {
	t.Helper()
	out, err := c.EC2.DisableSerialConsoleAccess(&ec2.DisableSerialConsoleAccessInput{})
	if err != nil {
		t.Fatalf("DisableSerialConsoleAccess: %v", err)
	}
	return aws.BoolValue(out.SerialConsoleAccessEnabled)
}

// DiscoverDefaultVPC returns (vpcID, defaultSGID, defaultSubnetID). Fatals if
// the default VPC is absent — its absence indicates a bootstrap failure.
func DiscoverDefaultVPC(t *testing.T, c *AWSClient) (vpcID, sgID, subnetID string) {
	t.Helper()
	vpcs, err := c.EC2.DescribeVpcs(&ec2.DescribeVpcsInput{
		Filters: []*ec2.Filter{{Name: aws.String("is-default"), Values: []*string{aws.String("true")}}},
	})
	if err != nil {
		t.Fatalf("describe-vpcs: %v", err)
	}
	if len(vpcs.Vpcs) == 0 {
		t.Fatalf("no default VPC found")
	}
	vpcID = aws.StringValue(vpcs.Vpcs[0].VpcId)

	sgs, err := c.EC2.DescribeSecurityGroups(&ec2.DescribeSecurityGroupsInput{
		Filters: []*ec2.Filter{
			{Name: aws.String("vpc-id"), Values: []*string{aws.String(vpcID)}},
			{Name: aws.String("group-name"), Values: []*string{aws.String("default")}},
		},
	})
	if err != nil {
		t.Fatalf("describe-security-groups: %v", err)
	}
	if len(sgs.SecurityGroups) == 0 {
		t.Fatalf("default SG missing for VPC %s", vpcID)
	}
	sgID = aws.StringValue(sgs.SecurityGroups[0].GroupId)

	subnets, err := c.EC2.DescribeSubnets(&ec2.DescribeSubnetsInput{
		Filters: []*ec2.Filter{{Name: aws.String("vpc-id"), Values: []*string{aws.String(vpcID)}}},
	})
	if err != nil {
		t.Fatalf("describe-subnets: %v", err)
	}
	for _, s := range subnets.Subnets {
		if aws.BoolValue(s.DefaultForAz) {
			subnetID = aws.StringValue(s.SubnetId)
			break
		}
	}
	// Fall back to the first subnet if none is marked DefaultForAz; callers
	// just need a usable subnet ID.
	if subnetID == "" && len(subnets.Subnets) > 0 {
		subnetID = aws.StringValue(subnets.Subnets[0].SubnetId)
	}
	return vpcID, sgID, subnetID
}

// AuthorizeSSHIngress idempotently authorizes tcp/22 ingress from 0.0.0.0/0
// on sgID.
func AuthorizeSSHIngress(t *testing.T, c *AWSClient, sgID string) {
	t.Helper()
	AuthorizeTCPIngress(t, c, sgID, 22)
}

// AuthorizeTCPIngress idempotently authorizes a single TCP port from 0.0.0.0/0
// on sgID. InvalidPermission.Duplicate is treated as success — the rule was
// already in place on a re-run.
func AuthorizeTCPIngress(t *testing.T, c *AWSClient, sgID string, port int64) {
	t.Helper()
	_, err := c.EC2.AuthorizeSecurityGroupIngress(&ec2.AuthorizeSecurityGroupIngressInput{
		GroupId: aws.String(sgID),
		IpPermissions: []*ec2.IpPermission{{
			IpProtocol: aws.String("tcp"),
			FromPort:   aws.Int64(port),
			ToPort:     aws.Int64(port),
			IpRanges:   []*ec2.IpRange{{CidrIp: aws.String("0.0.0.0/0")}},
		}},
	})
	if err == nil {
		return
	}
	var aerr awserr.Error
	if asErr(err, &aerr) && aerr.Code() == "InvalidPermission.Duplicate" {
		return
	}
	t.Fatalf("authorize-security-group-ingress tcp/%d on %s: %v", port, sgID, err)
}

// AuthorizeICMPIngress idempotently authorizes ICMP (all types) ingress from
// 0.0.0.0/0 on sgID. Same Duplicate-tolerant contract as AuthorizeSSHIngress.
func AuthorizeICMPIngress(t *testing.T, c *AWSClient, sgID string) {
	t.Helper()
	_, err := c.EC2.AuthorizeSecurityGroupIngress(&ec2.AuthorizeSecurityGroupIngressInput{
		GroupId: aws.String(sgID),
		IpPermissions: []*ec2.IpPermission{{
			IpProtocol: aws.String("icmp"),
			FromPort:   aws.Int64(-1),
			ToPort:     aws.Int64(-1),
			IpRanges:   []*ec2.IpRange{{CidrIp: aws.String("0.0.0.0/0")}},
		}},
	})
	if err == nil {
		return
	}
	var aerr awserr.Error
	if asErr(err, &aerr) && aerr.Code() == "InvalidPermission.Duplicate" {
		return
	}
	t.Fatalf("authorize-security-group-ingress icmp on %s: %v", sgID, err)
}

// EnsureDefaultSGOpen authorizes tcp/22 + ICMP on the default SG. Default SGs
// only admit same-SG members; without this, runner-originated SSH/ping is dropped
// by the OVN port-group ACL.
func EnsureDefaultSGOpen(t *testing.T, c *AWSClient) {
	t.Helper()
	_, sgID, _ := DiscoverDefaultVPC(t, c)
	AuthorizeSSHIngress(t, c, sgID)
	AuthorizeICMPIngress(t, c, sgID)
}

// AttachVolumeWait attaches volID to instanceID at device and blocks until the
// volume reaches "in-use". Used by the data-durability tests where the attach
// must complete before the guest can format the disk.
func AttachVolumeWait(t *testing.T, c *AWSClient, volID, instanceID, device string) {
	t.Helper()
	_, err := c.EC2.AttachVolume(&ec2.AttachVolumeInput{
		VolumeId:   aws.String(volID),
		InstanceId: aws.String(instanceID),
		Device:     aws.String(device),
	})
	if err != nil {
		t.Fatalf("attach-volume %s -> %s as %s: %v", volID, instanceID, device, err)
	}
	WaitForVolumeState(t, c, volID, ec2.VolumeStateInUse, WithPoll(500*time.Millisecond))
}

// DetachVolumeWait detaches volID and blocks until it reaches "available".
func DetachVolumeWait(t *testing.T, c *AWSClient, volID string) {
	t.Helper()
	// A volume that is already on its way to available is the outcome asked
	// for, not a failure. Teardown runs this after a terminate that releases
	// the volume itself, so refusing an already-released volume fails a test
	// whose assertions all passed.
	if _, err := c.EC2.DetachVolume(&ec2.DetachVolumeInput{VolumeId: aws.String(volID)}); err != nil {
		if !isIncorrectVolumeState(err) {
			t.Fatalf("detach-volume %s: %v", volID, err)
		}
		t.Logf("detach-volume %s: already released, waiting for available", volID)
	}
	WaitForVolumeState(t, c, volID, "available", WithPoll(500*time.Millisecond))
}

// isIncorrectVolumeState reports whether err is the refusal a detach gets for a
// volume that is not attached. It is the answer for one already released, and
// for one still mid-detach, both of which are what a teardown wanted.
func isIncorrectVolumeState(err error) bool {
	var aerr awserr.Error
	return asErr(err, &aerr) && aerr.Code() == "IncorrectState"
}

// RegisterVolumeTeardown force-detaches then deletes volID at test end. It does
// not call t.Fatal, so a cleanup hiccup can't mask the test's real result, but a
// teardown that actually leaves the volume behind fails the test — a silent leak
// is worse than a noisy one.
func RegisterVolumeTeardown(t *testing.T, c *AWSClient, volID string) {
	t.Helper()
	t.Cleanup(func() {
		if err := teardownVolume(c.EC2, volID); err != nil {
			t.Errorf("LEAKED volume %s: %v", volID, err)
		}
	})
}

// Teardown windows. Detach and terminate are both asynchronous, and a volume is
// only released once the guest holding it is down, so these have to outlast a
// slow QEMU shutdown rather than a single API round trip.
const (
	volumeTeardownTimeout   = 90 * time.Second
	instanceTeardownTimeout = 3 * time.Minute
	teardownPollInterval    = 2 * time.Second
)

// teardownVolume force-detaches volID if it is still attached, waits for it to
// reach "available", then deletes it. It returns an error only when the volume
// is still present afterwards — that is, only when it has genuinely leaked.
//
// Deletion has to be self-sufficient rather than rely on teardown ordering.
// Fixture cleanups drain LIFO, so a suite that creates its instance before its
// volume tears the volume down first, while it is still attached: DeleteVolume
// answers VolumeInUse and the volume is stranded. Registration order cannot
// express that dependency, so the teardown resolves it itself and works for any
// creation order.
func teardownVolume(ec2c ec2iface.EC2API, volID string) error {
	forced := false
	deadline := time.Now().Add(volumeTeardownTimeout)
	for {
		out, err := ec2c.DescribeVolumes(&ec2.DescribeVolumesInput{
			VolumeIds: []*string{aws.String(volID)},
		})
		if err != nil {
			if volumeGone(err) {
				return nil
			}
			return fmt.Errorf("describe volume %s: %w", volID, err)
		}
		if len(out.Volumes) == 0 {
			return nil
		}

		switch state := aws.StringValue(out.Volumes[0].State); state {
		case ec2.VolumeStateDeleting, ec2.VolumeStateDeleted:
			return nil
		case ec2.VolumeStateAvailable:
			if _, err := ec2c.DeleteVolume(&ec2.DeleteVolumeInput{
				VolumeId: aws.String(volID),
			}); err != nil && !volumeGone(err) {
				return fmt.Errorf("delete volume %s: %w", volID, err)
			}
			return nil
		default:
			// Force, because by teardown the attaching guest is usually already
			// terminating and there is nothing left to acknowledge a graceful
			// detach. Issued once, then we wait for the state to settle.
			if !forced {
				// Bound the call itself: the SDK's built-in retry (4 attempts
				// against the gateway's 30s NATS budget) can burn ~119s on a
				// slow detach, more than the whole deadline, leaving no time
				// for the state poll below and misreporting a slow detach as
				// an unexplained leak. Deriving the context from the same
				// deadline this loop already honours means a bounded-out
				// detach still leaves the poll loop below a chance to observe
				// the volume settle.
				detachCtx, cancel := context.WithDeadline(context.Background(), deadline)
				_, _ = ec2c.DetachVolumeWithContext(detachCtx, &ec2.DetachVolumeInput{
					VolumeId: aws.String(volID),
					Force:    aws.Bool(true),
				})
				cancel()
				forced = true
				if time.Now().After(deadline) {
					return fmt.Errorf("volume %s still %s after %s, not deleted", volID, state, volumeTeardownTimeout)
				}
				continue
			}
			if time.Now().After(deadline) {
				return fmt.Errorf("volume %s still %s after %s, not deleted", volID, state, volumeTeardownTimeout)
			}
		}
		time.Sleep(teardownPollInterval)
	}
}

// teardownInstance terminates instID and waits until it reports "terminated".
// The wait is the point: TerminateInstances returns as soon as the request is
// accepted, so a volume teardown running straight afterwards would still find
// its volume attached to a guest that has not gone away yet.
func teardownInstance(ec2c ec2iface.EC2API, instID string) error {
	if _, err := ec2c.TerminateInstances(&ec2.TerminateInstancesInput{
		InstanceIds: []*string{aws.String(instID)},
	}); err != nil && !instanceGone(err) {
		return fmt.Errorf("terminate instance %s: %w", instID, err)
	}

	deadline := time.Now().Add(instanceTeardownTimeout)
	for {
		out, err := ec2c.DescribeInstances(&ec2.DescribeInstancesInput{
			InstanceIds: []*string{aws.String(instID)},
		})
		if err != nil {
			if instanceGone(err) {
				return nil
			}
			return fmt.Errorf("describe instance %s: %w", instID, err)
		}
		if len(out.Reservations) == 0 || len(out.Reservations[0].Instances) == 0 {
			return nil
		}
		state := aws.StringValue(out.Reservations[0].Instances[0].State.Name)
		if state == ec2.InstanceStateNameTerminated {
			return nil
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("instance %s still %s after %s", instID, state, instanceTeardownTimeout)
		}
		time.Sleep(teardownPollInterval)
	}
}

// volumeGone and instanceGone report whether err is the API's way of saying the
// resource no longer exists, which is success for a teardown.
func volumeGone(err error) bool { return errCodeIs(err, "InvalidVolume.NotFound") }

func instanceGone(err error) bool { return errCodeIs(err, "InvalidInstanceID.NotFound") }

func errCodeIs(err error, code string) bool {
	var ae awserr.Error
	return asErr(err, &ae) && ae.Code() == code
}

// asErr is a thin errors.As that lets the helper stay testify-free.
func asErr(err error, target *awserr.Error) bool {
	if a, ok := err.(awserr.Error); ok {
		*target = a
		return true
	}
	return false
}