//go:build e2e

package multinode

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

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/mulgadc/spinifex/tests/e2e/harness"
	"github.com/stretchr/testify/require"
)

// Spread placement + NAT GW runs as one top-level Test with 6 sub-tests sharing
// one VPC + bastion + private trio + NAT GW graph. Shared setup avoids rebuilding
// the full graph per sub-test. Outer Test is sequential (owns CIDR 10.100.0.0/16 + EIP pool).

// runSpread drives the spread-placement + NAT GW suite.
// Wrapper TestMultinodeSpread in tests_test.go controls execution order.
func runSpread(t *testing.T, fix *Fixture) {
	harness.Phase(t, "Multinode — Spread Placement + NAT Gateway")
	c := fix.AWS

	v := spreadSetupVPC(t, fix)

	t.Run("VPCSetup", func(t *testing.T) {
		// DescribeRouteTables is the explicit oracle for this sub-test; setup already fataled on create errors.
		out, err := c.EC2.DescribeRouteTables(&ec2.DescribeRouteTablesInput{
			RouteTableIds: []*string{aws.String(v.MainRTBID)},
		})
		require.NoError(t, err, "describe main RT")
		require.NotEmpty(t, out.RouteTables, "main RT not found")
		var foundIGW bool
		for _, r := range out.RouteTables[0].Routes {
			if aws.StringValue(r.DestinationCidrBlock) == "0.0.0.0/0" &&
				aws.StringValue(r.GatewayId) == v.IGWID {
				foundIGW = true
				break
			}
		}
		require.Truef(t, foundIGW, "main RT %s missing 0.0.0.0/0 -> %s route", v.MainRTBID, v.IGWID)
	})

	b := spreadSetupBastion(t, fix, v)

	t.Run("BastionSSH", func(t *testing.T) {
		// JUnit marker: SSH + scp completed in spreadSetupBastion.
		t.Logf("bastion %s SSH+scp ok at %s", b.ID, b.PublicIP)
	})

	p := spreadSetupPrivateTrio(t, fix, b)

	t.Run("SpreadPlacement", func(t *testing.T) {
		unique := uniqueCount(p.HostingNodes)
		harness.Detail(t, "unique_hosting_nodes", unique, "want", 3)
		// The NAT-egress scenario below only proves what it claims to (three
		// distinct chassis routing through one NAT GW) if the trio actually
		// landed on three distinct nodes. Two colocated instances is a spread
		// placement group failure, not a "best-effort" outcome to log past.
		require.Equalf(t, 3, unique,
			"spread placement failed: %d unique hosting nodes (want 3), hosts=%v", unique, p.HostingNodes)
	})

	t.Run("PreNATIsolation", func(t *testing.T) {
		for i, privIP := range p.PrivateIPs {
			if _, err := pingViaBastion(b.KeyPath, b.PublicIP, privIP); err == nil {
				t.Fatalf("baseline FAIL: %s (%s) can reach internet WITHOUT NAT GW", p.InstanceIDs[i], privIP)
			}
			harness.Detail(t, "pre_nat", p.InstanceIDs[i], "reachable", false)
		}
	})

	n := spreadSetupNATGW(t, fix, p)

	t.Run("NATGatewayInternet", func(t *testing.T) {
		harness.Step(t, "verify internet via NAT GW (x%d)", len(p.PrivateIPs))
		for i, privIP := range p.PrivateIPs {
			host := p.HostingNodes[i]
			harness.EventuallyErr(t, func() error {
				if _, err := pingViaBastion(b.KeyPath, b.PublicIP, privIP); err != nil {
					return fmt.Errorf("ping via NAT GW (%s on %s): %w", p.InstanceIDs[i], host, err)
				}
				return nil
			}, 90*time.Second, 5*time.Second)
			harness.Detail(t, "post_nat", p.InstanceIDs[i], "node", host, "reachable", true)
		}
	})

	t.Run("NATCleanupOrdering", func(t *testing.T) {
		harness.Step(t, "delete-route 0.0.0.0/0")
		_, err := c.EC2.DeleteRoute(&ec2.DeleteRouteInput{
			RouteTableId:         aws.String(n.PrivRTBID),
			DestinationCidrBlock: aws.String("0.0.0.0/0"),
		})
		require.NoError(t, err, "delete-route 0.0.0.0/0")
		*n.RouteDeleted = true

		harness.Step(t, "disassociate-route-table %s", n.AssocID)
		_, err = c.EC2.DisassociateRouteTable(&ec2.DisassociateRouteTableInput{
			AssociationId: aws.String(n.AssocID),
		})
		require.NoError(t, err, "disassociate private RT")
		*n.AssocReleased = true

		harness.Step(t, "delete-route-table %s", n.PrivRTBID)
		_, err = c.EC2.DeleteRouteTable(&ec2.DeleteRouteTableInput{
			RouteTableId: aws.String(n.PrivRTBID),
		})
		require.NoError(t, err, "delete private RT")
		*n.RTBDeleted = true

		harness.Step(t, "delete-nat-gateway %s", n.NatGWID)
		_, err = c.EC2.DeleteNatGateway(&ec2.DeleteNatGatewayInput{
			NatGatewayId: aws.String(n.NatGWID),
		})
		require.NoError(t, err, "delete-nat-gateway")
		waitForNATGatewayStateFatal(t, c, n.NatGWID, "deleted")
		*n.NatDeleted = true

		harness.Step(t, "release-address %s", n.EIPAllocID)
		_, err = c.EC2.ReleaseAddress(&ec2.ReleaseAddressInput{
			AllocationId: aws.String(n.EIPAllocID),
		})
		require.NoError(t, err, "release-address")
		*n.EIPReleased = true
	})
}

// --- Layered setup helpers ------------------------------------------------

// spreadVPC carries the IDs produced by spreadSetupVPC. Public + private subnet,
// IGW + main RT (with 0.0.0.0/0 IGW route), and bastion+private SGs.
type spreadVPC struct {
	VPCID        string
	PubSubnetID  string
	PrivSubnetID string
	IGWID        string
	BastionSGID  string
	PrivSGID     string
	MainRTBID    string
}

// spreadSetupVPC creates the VPC graph (VPC + subnets + IGW + main RT route + SGs)
// and registers LIFO cleanup against t.
func spreadSetupVPC(t *testing.T, fix *Fixture) spreadVPC {
	t.Helper()
	c := fix.AWS

	harness.Step(t, "create VPC 10.100.0.0/16 + subnets + IGW + SGs")
	vpcOut, err := c.EC2.CreateVpc(&ec2.CreateVpcInput{
		CidrBlock: aws.String("10.100.0.0/16"),
	})
	require.NoError(t, err, "create-vpc 10.100.0.0/16")
	vpcID := aws.StringValue(vpcOut.Vpc.VpcId)
	require.NotEmpty(t, vpcID, "VpcId empty")
	t.Cleanup(func() {
		_, _ = c.EC2.DeleteVpc(&ec2.DeleteVpcInput{VpcId: aws.String(vpcID)})
	})

	pubSubOut, err := c.EC2.CreateSubnet(&ec2.CreateSubnetInput{
		VpcId:     aws.String(vpcID),
		CidrBlock: aws.String("10.100.1.0/24"),
	})
	require.NoError(t, err, "create public subnet")
	pubSubnetID := aws.StringValue(pubSubOut.Subnet.SubnetId)
	t.Cleanup(func() {
		_, _ = c.EC2.DeleteSubnet(&ec2.DeleteSubnetInput{SubnetId: aws.String(pubSubnetID)})
	})

	privSubOut, err := c.EC2.CreateSubnet(&ec2.CreateSubnetInput{
		VpcId:     aws.String(vpcID),
		CidrBlock: aws.String("10.100.2.0/24"),
	})
	require.NoError(t, err, "create private subnet")
	privSubnetID := aws.StringValue(privSubOut.Subnet.SubnetId)
	t.Cleanup(func() {
		_, _ = c.EC2.DeleteSubnet(&ec2.DeleteSubnetInput{SubnetId: aws.String(privSubnetID)})
	})

	_, err = c.EC2.ModifySubnetAttribute(&ec2.ModifySubnetAttributeInput{
		SubnetId:            aws.String(pubSubnetID),
		MapPublicIpOnLaunch: &ec2.AttributeBooleanValue{Value: aws.Bool(true)},
	})
	require.NoError(t, err, "modify-subnet-attribute MapPublicIpOnLaunch")

	igwOut, err := c.EC2.CreateInternetGateway(&ec2.CreateInternetGatewayInput{})
	require.NoError(t, err, "create-internet-gateway")
	igwID := aws.StringValue(igwOut.InternetGateway.InternetGatewayId)
	t.Cleanup(func() {
		_, _ = c.EC2.DetachInternetGateway(&ec2.DetachInternetGatewayInput{
			InternetGatewayId: aws.String(igwID),
			VpcId:             aws.String(vpcID),
		})
		_, _ = c.EC2.DeleteInternetGateway(&ec2.DeleteInternetGatewayInput{
			InternetGatewayId: aws.String(igwID),
		})
	})
	_, err = c.EC2.AttachInternetGateway(&ec2.AttachInternetGatewayInput{
		InternetGatewayId: aws.String(igwID),
		VpcId:             aws.String(vpcID),
	})
	require.NoError(t, err, "attach-internet-gateway")

	mainRT, err := c.EC2.DescribeRouteTables(&ec2.DescribeRouteTablesInput{
		Filters: []*ec2.Filter{
			{Name: aws.String("vpc-id"), Values: []*string{aws.String(vpcID)}},
		},
	})
	require.NoError(t, err, "describe-route-tables (main)")
	require.NotEmpty(t, mainRT.RouteTables, "no main route table for fresh VPC")
	mainRTBID := aws.StringValue(mainRT.RouteTables[0].RouteTableId)
	_, err = c.EC2.CreateRoute(&ec2.CreateRouteInput{
		RouteTableId:         aws.String(mainRTBID),
		DestinationCidrBlock: aws.String("0.0.0.0/0"),
		GatewayId:            aws.String(igwID),
	})
	require.NoError(t, err, "create-route main RT -> IGW")

	bastionSGOut, err := c.EC2.CreateSecurityGroup(&ec2.CreateSecurityGroupInput{
		VpcId:       aws.String(vpcID),
		GroupName:   aws.String("natgw-bastion"),
		Description: aws.String("Spread bastion (SSH ingress from anywhere)"),
	})
	require.NoError(t, err, "create bastion SG")
	bastionSGID := aws.StringValue(bastionSGOut.GroupId)
	t.Cleanup(func() {
		_, _ = c.EC2.DeleteSecurityGroup(&ec2.DeleteSecurityGroupInput{
			GroupId: aws.String(bastionSGID),
		})
	})
	_, err = c.EC2.AuthorizeSecurityGroupIngress(&ec2.AuthorizeSecurityGroupIngressInput{
		GroupId: aws.String(bastionSGID),
		IpPermissions: []*ec2.IpPermission{
			{
				IpProtocol: aws.String("tcp"),
				FromPort:   aws.Int64(22),
				ToPort:     aws.Int64(22),
				IpRanges:   []*ec2.IpRange{{CidrIp: aws.String("0.0.0.0/0")}},
			},
		},
	})
	require.NoError(t, err, "authorize bastion SG tcp/22")

	privSGOut, err := c.EC2.CreateSecurityGroup(&ec2.CreateSecurityGroupInput{
		VpcId:       aws.String(vpcID),
		GroupName:   aws.String("natgw-private"),
		Description: aws.String("Spread private (SSH from bastion-sg, ICMP from VPC)"),
	})
	require.NoError(t, err, "create private SG")
	privSGID := aws.StringValue(privSGOut.GroupId)
	t.Cleanup(func() {
		_, _ = c.EC2.DeleteSecurityGroup(&ec2.DeleteSecurityGroupInput{
			GroupId: aws.String(privSGID),
		})
	})
	_, err = c.EC2.AuthorizeSecurityGroupIngress(&ec2.AuthorizeSecurityGroupIngressInput{
		GroupId: aws.String(privSGID),
		IpPermissions: []*ec2.IpPermission{
			{
				IpProtocol: aws.String("tcp"),
				FromPort:   aws.Int64(22),
				ToPort:     aws.Int64(22),
				UserIdGroupPairs: []*ec2.UserIdGroupPair{
					{GroupId: aws.String(bastionSGID)},
				},
			},
		},
	})
	require.NoError(t, err, "authorize private SG tcp/22 from bastion-sg")
	_, err = c.EC2.AuthorizeSecurityGroupIngress(&ec2.AuthorizeSecurityGroupIngressInput{
		GroupId: aws.String(privSGID),
		IpPermissions: []*ec2.IpPermission{
			{
				IpProtocol: aws.String("icmp"),
				FromPort:   aws.Int64(-1),
				ToPort:     aws.Int64(-1),
				IpRanges:   []*ec2.IpRange{{CidrIp: aws.String("10.100.0.0/16")}},
			},
		},
	})
	require.NoError(t, err, "authorize private SG icmp from VPC")

	harness.Detail(t, "vpc", vpcID, "pub_subnet", pubSubnetID, "priv_subnet", privSubnetID,
		"igw", igwID, "bastion_sg", bastionSGID, "priv_sg", privSGID)

	return spreadVPC{
		VPCID:        vpcID,
		PubSubnetID:  pubSubnetID,
		PrivSubnetID: privSubnetID,
		IGWID:        igwID,
		BastionSGID:  bastionSGID,
		PrivSGID:     privSGID,
		MainRTBID:    mainRTBID,
	}
}

// spreadBastion carries the bastion VM IDs plus the VPC it lives in.
type spreadBastion struct {
	VPC      spreadVPC
	ID       string
	PublicIP string
	KeyName  string
	KeyPath  string
}

// spreadSetupBastion launches the bastion in the public subnet, waits for SSH, and copies
// the test PEM to /tmp/key.pem so the bastion can re-key into private VMs.
func spreadSetupBastion(t *testing.T, fix *Fixture, v spreadVPC) spreadBastion {
	t.Helper()
	c := fix.AWS

	amiID := needAMI(t, fix, needArch(t, fix))
	instType, _ := needInstanceTypeArch(t, fix)
	keyName, keyPath := needKeyPair(t, fix)

	harness.Step(t, "launch bastion in public subnet")
	bastionOut, err := c.EC2.RunInstances(&ec2.RunInstancesInput{
		ImageId:          aws.String(amiID),
		InstanceType:     aws.String(instType),
		KeyName:          aws.String(keyName),
		SubnetId:         aws.String(v.PubSubnetID),
		SecurityGroupIds: []*string{aws.String(v.BastionSGID)},
		MinCount:         aws.Int64(1),
		MaxCount:         aws.Int64(1),
	})
	require.NoError(t, err, "run-instances bastion")
	require.NotEmpty(t, bastionOut.Instances, "run-instances bastion returned no Instances")
	bastionID := aws.StringValue(bastionOut.Instances[0].InstanceId)
	require.NotEmpty(t, bastionID, "bastion InstanceId empty")
	t.Cleanup(func() {
		_, _ = c.EC2.TerminateInstances(&ec2.TerminateInstancesInput{
			InstanceIds: []*string{aws.String(bastionID)},
		})
		_ = waitForInstanceStateBest(c, bastionID, "terminated", 5*time.Minute)
	})

	bastion := harness.WaitForInstanceState(t, c, bastionID, "running",
		harness.WithTimeout(120*time.Second), harness.WithPoll(2*time.Second))
	bastionPubIP := aws.StringValue(bastion.PublicIpAddress)
	require.NotEmptyf(t, bastionPubIP,
		"bastion %s has no PublicIpAddress (pool mode required for phase 11)", bastionID)
	harness.Detail(t, "bastion", bastionID, "bastion_pub_ip", bastionPubIP)

	// Dump datapath diagnostics on failure to help localise ARP/geneve SYN drops.
	t.Cleanup(func() {
		if t.Failed() {
			dumpPlacementNATDiag(t, fix, bastionPubIP, bastionID)
		}
	})

	waitForBastionSSH(t, bastionPubIP, keyPath, 5*time.Minute)
	scpKeyToBastion(t, keyPath, bastionPubIP)

	return spreadBastion{
		VPC:      v,
		ID:       bastionID,
		PublicIP: bastionPubIP,
		KeyName:  keyName,
		KeyPath:  keyPath,
	}
}

// spreadPrivate carries the spread-placement trio: PG name, 3 private instance
// IDs, their private IPs, and the cluster node each one landed on.
type spreadPrivate struct {
	Bastion      spreadBastion
	InstanceIDs  []string
	PrivateIPs   []string
	HostingNodes []string
	PGName       string
}

// spreadSetupPrivateTrio creates the spread placement group + 3 private VMs,
// waits for SSH-via-bastion, and resolves each instance's hosting node and private IP.
func spreadSetupPrivateTrio(t *testing.T, fix *Fixture, b spreadBastion) spreadPrivate {
	t.Helper()
	c := fix.AWS

	amiID := needAMI(t, fix, needArch(t, fix))
	instType, _ := needInstanceTypeArch(t, fix)

	harness.Step(t, "create spread placement group + 3 private VMs")
	pgName := "nat-spread"
	_, err := c.EC2.CreatePlacementGroup(&ec2.CreatePlacementGroupInput{
		GroupName: aws.String(pgName),
		Strategy:  aws.String("spread"),
	})
	require.NoError(t, err, "create-placement-group nat-spread")
	t.Cleanup(func() {
		_, _ = c.EC2.DeletePlacementGroup(&ec2.DeletePlacementGroupInput{
			GroupName: aws.String(pgName),
		})
	})

	privOut, err := c.EC2.RunInstances(&ec2.RunInstancesInput{
		ImageId:          aws.String(amiID),
		InstanceType:     aws.String(instType),
		KeyName:          aws.String(b.KeyName),
		SubnetId:         aws.String(b.VPC.PrivSubnetID),
		SecurityGroupIds: []*string{aws.String(b.VPC.PrivSGID)},
		MinCount:         aws.Int64(3),
		MaxCount:         aws.Int64(3),
		Placement: &ec2.Placement{
			GroupName: aws.String(pgName),
		},
	})
	require.NoError(t, err, "run-instances private x3")
	require.Lenf(t, privOut.Instances, 3, "run-instances private returned %d instances", len(privOut.Instances))

	privIDs := make([]string, 0, 3)
	for _, inst := range privOut.Instances {
		id := aws.StringValue(inst.InstanceId)
		require.NotEmpty(t, id, "private InstanceId empty")
		privIDs = append(privIDs, id)
		idCopy := id
		t.Cleanup(func() {
			_, _ = c.EC2.TerminateInstances(&ec2.TerminateInstancesInput{
				InstanceIds: []*string{aws.String(idCopy)},
			})
			_ = waitForInstanceStateBest(c, idCopy, "terminated", 5*time.Minute)
		})
	}
	harness.Detail(t, "private_ids", privIDs)

	for _, id := range privIDs {
		harness.WaitForInstanceState(t, c, id, "running",
			harness.WithTimeout(120*time.Second), harness.WithPoll(2*time.Second))
	}

	hostingNodes := make([]string, 0, len(privIDs))
	for _, id := range privIDs {
		n := harness.InstanceHostingNode(t, fix.Cluster, id)
		if n == nil {
			hostingNodes = append(hostingNodes, "unknown")
			harness.Detail(t, "instance", id, "host", "unknown")
			continue
		}
		hostingNodes = append(hostingNodes, n.Name)
		harness.Detail(t, "instance", id, "host", n.Name)
	}

	privIPs := make([]string, 0, len(privIDs))
	eniIDs := make([]string, 0, len(privIDs))
	for _, id := range privIDs {
		out, err := c.EC2.DescribeInstances(&ec2.DescribeInstancesInput{
			InstanceIds: []*string{aws.String(id)},
		})
		require.NoError(t, err, "describe-instances private")
		require.NotEmpty(t, out.Reservations, "no Reservations for %s", id)
		require.NotEmpty(t, out.Reservations[0].Instances, "no Instances for %s", id)
		inst := out.Reservations[0].Instances[0]
		ip := aws.StringValue(inst.PrivateIpAddress)
		require.NotEmptyf(t, ip, "%s has no PrivateIpAddress", id)
		privIPs = append(privIPs, ip)
		eni := ""
		if len(inst.NetworkInterfaces) > 0 {
			eni = aws.StringValue(inst.NetworkInterfaces[0].NetworkInterfaceId)
		}
		eniIDs = append(eniIDs, eni)
	}
	harness.Detail(t, "private_ips", privIPs)

	// The failure this phase is most likely to hit is a guest that never answers
	// SSH, and the NAT/EIP dump says nothing about one. Registered before the
	// wait so it fires on the guest that failed, and covers all three so a
	// working guest is available as the control.
	t.Cleanup(func() {
		if t.Failed() {
			dumpSpreadGuestDiag(t, fix, privIDs, privIPs, eniIDs, hostingNodes)
		}
	})

	harness.Step(t, "wait for private SSH via bastion (x%d)", len(privIPs))
	for i, privIP := range privIPs {
		waitForSSHViaBastion(t, b.KeyPath, b.PublicIP, privIP, fmt.Sprintf("priv #%d (%s)", i, privIDs[i]))
	}

	return spreadPrivate{
		Bastion:      b,
		InstanceIDs:  privIDs,
		PrivateIPs:   privIPs,
		HostingNodes: hostingNodes,
		PGName:       pgName,
	}
}

// spreadNAT carries the NAT GW + private RT IDs. The *bool fields are already-torn-down
// latches so deferred cleanup skips explicit teardown steps that already ran.
type spreadNAT struct {
	Private       spreadPrivate
	NatGWID       string
	EIPAllocID    string
	EIP           string
	PrivRTBID     string
	AssocID       string
	NatDeleted    *bool
	EIPReleased   *bool
	RTBDeleted    *bool
	AssocReleased *bool
	RouteDeleted  *bool
}

// spreadSetupNATGW allocates an EIP, creates the NAT GW in the public subnet, builds the
// private RT + route, and waits for the NAT GW to reach `available`.
func spreadSetupNATGW(t *testing.T, fix *Fixture, p spreadPrivate) spreadNAT {
	t.Helper()
	c := fix.AWS

	harness.Step(t, "allocate EIP + create NAT GW + private RT")
	eipOut, err := c.EC2.AllocateAddress(&ec2.AllocateAddressInput{
		Domain: aws.String("vpc"),
	})
	require.NoError(t, err, "allocate-address vpc")
	natAllocID := aws.StringValue(eipOut.AllocationId)
	natPubIP := aws.StringValue(eipOut.PublicIp)
	require.NotEmpty(t, natAllocID, "EIP AllocationId empty")
	eipReleased := new(bool)
	t.Cleanup(func() {
		if !*eipReleased {
			_, _ = c.EC2.ReleaseAddress(&ec2.ReleaseAddressInput{
				AllocationId: aws.String(natAllocID),
			})
		}
	})
	harness.Detail(t, "eip", natPubIP, "alloc", natAllocID)

	natOut, err := c.EC2.CreateNatGateway(&ec2.CreateNatGatewayInput{
		SubnetId:     aws.String(p.Bastion.VPC.PubSubnetID),
		AllocationId: aws.String(natAllocID),
	})
	require.NoError(t, err, "create-nat-gateway")
	natGWID := aws.StringValue(natOut.NatGateway.NatGatewayId)
	require.NotEmpty(t, natGWID, "NatGatewayId empty")
	natDeleted := new(bool)
	t.Cleanup(func() {
		if !*natDeleted {
			_, _ = c.EC2.DeleteNatGateway(&ec2.DeleteNatGatewayInput{
				NatGatewayId: aws.String(natGWID),
			})
			_ = waitForNATGatewayStateBest(c, natGWID, "deleted", 5*time.Minute)
		}
	})
	harness.Detail(t, "nat_gw", natGWID)
	waitForNATGatewayStateFatal(t, c, natGWID, "available")

	privRTOut, err := c.EC2.CreateRouteTable(&ec2.CreateRouteTableInput{
		VpcId: aws.String(p.Bastion.VPC.VPCID),
	})
	require.NoError(t, err, "create-route-table (private)")
	privRTBID := aws.StringValue(privRTOut.RouteTable.RouteTableId)
	require.NotEmpty(t, privRTBID, "private RouteTableId empty")
	rtbDeleted := new(bool)
	t.Cleanup(func() {
		if !*rtbDeleted {
			_, _ = c.EC2.DeleteRouteTable(&ec2.DeleteRouteTableInput{
				RouteTableId: aws.String(privRTBID),
			})
		}
	})

	rtbAssocOut, err := c.EC2.AssociateRouteTable(&ec2.AssociateRouteTableInput{
		RouteTableId: aws.String(privRTBID),
		SubnetId:     aws.String(p.Bastion.VPC.PrivSubnetID),
	})
	require.NoError(t, err, "associate private RT")
	rtbAssocID := aws.StringValue(rtbAssocOut.AssociationId)
	assocReleased := new(bool)
	t.Cleanup(func() {
		if !*assocReleased {
			_, _ = c.EC2.DisassociateRouteTable(&ec2.DisassociateRouteTableInput{
				AssociationId: aws.String(rtbAssocID),
			})
		}
	})

	_, err = c.EC2.CreateRoute(&ec2.CreateRouteInput{
		RouteTableId:         aws.String(privRTBID),
		DestinationCidrBlock: aws.String("0.0.0.0/0"),
		NatGatewayId:         aws.String(natGWID),
	})
	require.NoError(t, err, "create-route 0.0.0.0/0 -> NAT GW")
	routeDeleted := new(bool)
	t.Cleanup(func() {
		if !*routeDeleted {
			_, _ = c.EC2.DeleteRoute(&ec2.DeleteRouteInput{
				RouteTableId:         aws.String(privRTBID),
				DestinationCidrBlock: aws.String("0.0.0.0/0"),
			})
		}
	})

	return spreadNAT{
		Private:       p,
		NatGWID:       natGWID,
		EIPAllocID:    natAllocID,
		EIP:           natPubIP,
		PrivRTBID:     privRTBID,
		AssocID:       rtbAssocID,
		NatDeleted:    natDeleted,
		EIPReleased:   eipReleased,
		RTBDeleted:    rtbDeleted,
		AssocReleased: assocReleased,
		RouteDeleted:  routeDeleted,
	}
}

// --- helpers --------------------------------------------------------------

// uniqueCount returns the number of distinct non-empty, non-"unknown" strings in xs.
func uniqueCount(xs []string) int {
	seen := map[string]bool{}
	for _, x := range xs {
		if x == "" || x == "unknown" {
			continue
		}
		seen[x] = true
	}
	return len(seen)
}

// waitForBastionSSH polls SSH until the handshake completes. Uses a 5min budget to cover
// both cloud-init and OVN flow install on the external switch lagging instance "running".
// Logs an ICMP probe each iteration to distinguish "datapath down" from "sshd not up yet".
func waitForBastionSSH(t *testing.T, host, keyPath string, timeout time.Duration) {
	t.Helper()
	harness.Step(t, "wait bastion SSH %s", host)
	harness.EventuallyErr(t, func() error {
		pingCtx, pingCancel := context.WithTimeout(context.Background(), 3*time.Second)
		pingOut, pingErr := exec.CommandContext(pingCtx, "ping", "-c", "1", "-W", "2", host).CombinedOutput()
		pingCancel()
		pingState := "ok"
		if pingErr != nil {
			pingState = fmt.Sprintf("fail (%v)", pingErr)
		}

		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		args := []string{
			"-o", "StrictHostKeyChecking=no",
			"-o", "UserKnownHostsFile=/dev/null",
			"-o", "LogLevel=ERROR",
			"-o", "ConnectTimeout=5",
			"-o", "BatchMode=yes",
			"-i", keyPath,
			"ubuntu@" + host,
			"true",
		}
		if out, err := exec.CommandContext(ctx, "ssh", args...).CombinedOutput(); err != nil {
			return fmt.Errorf("ssh %s: %w\n  ping=%s\n  ping_output=%s\n  ssh_output=%s",
				host, err, pingState, strings.TrimSpace(string(pingOut)), string(out))
		}
		return nil
	}, timeout, 5*time.Second)
}

// scpKeyToBastion copies the test PEM to /tmp/key.pem on the bastion and
// chmods it so the bastion can re-key into private VMs.
func scpKeyToBastion(t *testing.T, keyPath, host string) {
	t.Helper()
	harness.Step(t, "scp key -> bastion:/tmp/key.pem")
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	scpArgs := []string{
		"-o", "StrictHostKeyChecking=no",
		"-o", "UserKnownHostsFile=/dev/null",
		"-o", "LogLevel=ERROR",
		"-o", "ConnectTimeout=10",
		"-o", "BatchMode=yes",
		"-i", keyPath,
		keyPath,
		"ubuntu@" + host + ":/tmp/key.pem",
	}
	if out, err := exec.CommandContext(ctx, "scp", scpArgs...).CombinedOutput(); err != nil {
		t.Fatalf("scp %s -> %s:/tmp/key.pem: %v\n%s", keyPath, host, err, string(out))
	}
	// chmod 600: sshd refuses keys with loose permissions.
	ctx2, cancel2 := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel2()
	sshArgs := []string{
		"-o", "StrictHostKeyChecking=no",
		"-o", "UserKnownHostsFile=/dev/null",
		"-o", "LogLevel=ERROR",
		"-o", "ConnectTimeout=10",
		"-o", "BatchMode=yes",
		"-i", keyPath,
		"ubuntu@" + host,
		"chmod 600 /tmp/key.pem",
	}
	if out, err := exec.CommandContext(ctx2, "ssh", sshArgs...).CombinedOutput(); err != nil {
		t.Fatalf("chmod /tmp/key.pem: %v\n%s", err, string(out))
	}
}

// sshViaBastion runs cmd on privIP through bastionHost using BatchMode +
// ConnectTimeout=10 nested-ssh — mirrors the bash priv_hop_cmd helper.
// Returns combined stdout+stderr and the exit error.
func sshViaBastion(keyPath, bastionHost, privIP, cmd string) (string, error) {
	hop := fmt.Sprintf(
		"ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null "+
			"-o LogLevel=ERROR -o ConnectTimeout=10 -o BatchMode=yes "+
			"-i /tmp/key.pem ubuntu@%s '%s'",
		privIP, cmd,
	)
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	args := []string{
		"-o", "StrictHostKeyChecking=no",
		"-o", "UserKnownHostsFile=/dev/null",
		"-o", "LogLevel=ERROR",
		"-o", "ConnectTimeout=10",
		"-o", "BatchMode=yes",
		"-i", keyPath,
		"ubuntu@" + bastionHost,
		hop,
	}
	out, err := exec.CommandContext(ctx, "ssh", args...).CombinedOutput()
	return string(out), err
}

// waitForSSHViaBastion polls hostname-through-bastion until sshd is up on the private VM.
func waitForSSHViaBastion(t *testing.T, keyPath, bastionHost, privIP, label string) {
	t.Helper()
	harness.EventuallyErr(t, func() error {
		out, err := sshViaBastion(keyPath, bastionHost, privIP, "hostname")
		if err != nil {
			return fmt.Errorf("%s: bastion->%s hostname: %w\n%s", label, privIP, err, out)
		}
		if strings.TrimSpace(out) == "" {
			return fmt.Errorf("%s: empty hostname response", label)
		}
		return nil
	}, 3*time.Minute, 5*time.Second)
}

// pingViaBastion sends a single 8.8.8.8 ping through the bastion hop. Used
// twice: pre-NAT (expect err) and post-NAT (expect nil after a poll).
func pingViaBastion(keyPath, bastionHost, privIP string) (string, error) {
	return sshViaBastion(keyPath, bastionHost, privIP, "ping -c 1 -W 3 8.8.8.8")
}

// --- NAT GW state pollers (promote to harness/poll.go when a third consumer appears)

func waitForNATGatewayStateFatal(t *testing.T, c *harness.AWSClient, id, target string) {
	t.Helper()
	harness.EventuallyErr(t, func() error {
		out, err := c.EC2.DescribeNatGateways(&ec2.DescribeNatGatewaysInput{
			NatGatewayIds: []*string{aws.String(id)},
		})
		if err != nil {
			return fmt.Errorf("describe-nat-gateways %s: %w", id, err)
		}
		if len(out.NatGateways) == 0 {
			if target == "deleted" {
				return nil
			}
			return fmt.Errorf("%s not found", id)
		}
		state := aws.StringValue(out.NatGateways[0].State)
		if state == target {
			return nil
		}
		if state == "failed" {
			msg := aws.StringValue(out.NatGateways[0].FailureMessage)
			return fmt.Errorf("%s entered failed state: %s", id, msg)
		}
		return fmt.Errorf("%s state=%s want=%s", id, state, target)
	}, 5*time.Minute, 2*time.Second)
}

func waitForNATGatewayStateBest(c *harness.AWSClient, id, target string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	for {
		out, err := c.EC2.DescribeNatGateways(&ec2.DescribeNatGatewaysInput{
			NatGatewayIds: []*string{aws.String(id)},
		})
		if err == nil {
			if len(out.NatGateways) == 0 && target == "deleted" {
				return nil
			}
			if len(out.NatGateways) > 0 && aws.StringValue(out.NatGateways[0].State) == target {
				return nil
			}
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("nat gw %s did not reach %s within %s", id, target, timeout)
		}
		time.Sleep(2 * time.Second)
	}
}

// dumpPlacementNATDiag fans out datapath probes (ARP, xfrm, OVS/OVN state) to every
// cluster node on failure and writes results under t's artifact dir. Best-effort:
// SSH errors are logged but never re-fail the test.
// dumpSpreadGuestDiag captures, for every spread guest, the evidence needed to
// tell "the guest never booted" from "the guest booted and its port never
// bound": its console, its logical switch port in NB and its port binding in SB,
// and the OVS interface backing its tap on the node it landed on.
func dumpSpreadGuestDiag(t *testing.T, fix *Fixture, ids, ips, enis, nodes []string) {
	t.Helper()
	if fix.Env == nil || fix.Cluster == nil || len(fix.Cluster.Nodes) == 0 {
		t.Logf("dumpSpreadGuestDiag: no env or cluster nodes; skipping")
		return
	}
	artifactDir := fix.ArtifactDir(t)
	ssh := harness.NewPeerSSH()

	for i, id := range ids {
		ip, eni, node := at(ips, i), at(enis, i), at(nodes, i)
		t.Logf("dumpSpreadGuestDiag: guest %s ip=%s eni=%s node=%s", id, ip, eni, node)

		console, err := harness.InstanceConsole(fix.AWS, id)
		if err != nil {
			console = fmt.Sprintf("(get-console-output failed: %v)", err)
		} else if console == "" {
			console = "(console empty — the guest may never have started)"
		}
		harness.DumpFile(t, artifactDir, fmt.Sprintf("guest-%s-console.log", id), []byte(console))

		port := "port-" + eni
		probes := []struct{ name, cmd string }{
			{"nb-lsp", fmt.Sprintf("sudo ovn-nbctl --no-leader-only list logical_switch_port %s 2>&1 || true", port)},
			{"sb-port-binding", fmt.Sprintf("sudo ovn-sbctl --no-leader-only find port_binding logical_port=%s 2>&1 || true", port)},
			{"ovs-interface", fmt.Sprintf("sudo ovs-vsctl --columns=name,ofport,external_ids,link_state find interface external_ids:iface-id=%s 2>&1 || true", port)},
			{"ovs-ports", "sudo ovs-vsctl list-ports br-int 2>&1 || true"},
			{"tap-link", fmt.Sprintf("ip -d link show | grep -A2 %s 2>&1 || true", eni)},
		}

		// Every node, not just the hosting one: a port bound on the wrong chassis
		// is exactly the failure a hosting-node-only dump cannot show.
		for _, n := range fix.Cluster.Nodes {
			for _, p := range probes {
				ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
				out, rerr := ssh.Run(ctx, n.Addr, p.cmd)
				cancel()
				header := fmt.Sprintf("# guest=%s eni=%s expected_node=%s host=%s cmd=%s\n",
					id, eni, node, n.Addr, p.cmd)
				if rerr != nil {
					header += fmt.Sprintf("# (ssh error: %v)\n", rerr)
				}
				name := fmt.Sprintf("guest-%s-%s-%s.log", id, n.Name, p.name)
				harness.DumpFile(t, artifactDir, name, append([]byte(header), out...))
			}
		}
	}
}

// at reads s[i] tolerantly: a diagnostic must not panic on a short slice.
func at(s []string, i int) string {
	if i < len(s) {
		return s[i]
	}
	return ""
}

func dumpPlacementNATDiag(t *testing.T, fix *Fixture, bastionPubIP, bastionID string) {
	t.Helper()
	if fix.Env == nil {
		t.Logf("dumpPlacementNATDiag: no env; skipping")
		return
	}
	if fix.Cluster == nil || len(fix.Cluster.Nodes) == 0 {
		t.Logf("dumpPlacementNATDiag: no cluster nodes; skipping")
		return
	}
	artifactDir := fix.ArtifactDir(t)

	t.Logf("dumpPlacementNATDiag: bastion=%s eip=%s nodes=%d -> %s",
		bastionID, bastionPubIP, len(fix.Cluster.Nodes), artifactDir)

	probes := []struct {
		name string
		cmd  string
	}{
		{"arp", "arp -an"},
		{"ip-route-eip", fmt.Sprintf("ip route get %s 2>&1 || true", bastionPubIP)},
		{"ip-addr", "ip -br addr"},
		{"ovs-show", "sudo ovs-vsctl show"},
		{"ovs-bridge-mappings", "sudo ovs-vsctl get Open_vSwitch . external_ids:ovn-bridge-mappings 2>&1 || true"},
		{"ovn-sb-chassis", "sudo ovn-sbctl --no-leader-only list chassis 2>&1 || true"},
		{"ovn-nb-nat", "sudo ovn-nbctl --no-leader-only list nat 2>&1 || true"},
		{"ovn-nb-gw-chassis", "sudo ovn-nbctl --no-leader-only list gateway_chassis 2>&1 || true"},
		{"ovn-nb-lrp", "sudo ovn-nbctl --no-leader-only list logical_router_port 2>&1 || true"},
		{"ip-link-geneve", "ip -d link show type geneve 2>&1 || true"},
	}

	ssh := harness.NewPeerSSH()
	for _, n := range fix.Cluster.Nodes {
		for _, p := range probes {
			ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
			out, err := ssh.Run(ctx, n.Addr, p.cmd)
			cancel()
			header := fmt.Sprintf("# host=%s cmd=%s\n", n.Addr, p.cmd)
			if err != nil {
				header += fmt.Sprintf("# (ssh error: %v)\n", err)
			}
			name := fmt.Sprintf("diag-%s-%s.log", n.Name, p.name)
			harness.DumpFile(t, artifactDir, name, append([]byte(header), out...))
		}
	}

	localProbes := []struct {
		name string
		cmd  string
		args []string
	}{
		{"local-ping-eip", "ping", []string{"-c", "3", "-W", "2", bastionPubIP}},
		{"local-traceroute-eip", "traceroute", []string{"-n", "-w", "2", "-m", "12", bastionPubIP}},
		{"local-ip-route-eip", "ip", []string{"route", "get", bastionPubIP}},
		{"local-arp", "arp", []string{"-an"}},
	}
	for _, p := range localProbes {
		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
		out, err := exec.CommandContext(ctx, p.cmd, p.args...).CombinedOutput()
		cancel()
		header := fmt.Sprintf("$ %s %s\n", p.cmd, strings.Join(p.args, " "))
		if err != nil {
			header += fmt.Sprintf("(exit error: %v)\n", err)
		}
		harness.DumpFile(t, artifactDir, p.name+".log", append([]byte(header), out...))
	}
}

// waitForInstanceStateBest is the cleanup-safe analogue of harness.WaitForInstanceState: polls without t.Fatal.
func waitForInstanceStateBest(c *harness.AWSClient, id, target string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	for {
		out, err := c.EC2.DescribeInstances(&ec2.DescribeInstancesInput{
			InstanceIds: []*string{aws.String(id)},
		})
		if err == nil && len(out.Reservations) > 0 && len(out.Reservations[0].Instances) > 0 {
			if aws.StringValue(out.Reservations[0].Instances[0].State.Name) == target {
				return nil
			}
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("instance %s did not reach %s within %s", id, target, timeout)
		}
		time.Sleep(2 * time.Second)
	}
}