package ec2k3d

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
)

type o11yPrivateDNSWaitEC2 struct {
	ec2API

	calls    int
	fail     bool
	verifyAt int
}

func (f *o11yPrivateDNSWaitEC2) DescribeVpcEndpointServiceConfigurations(
	context.Context,
	*ec2.DescribeVpcEndpointServiceConfigurationsInput,
	...func(*ec2.Options),
) (*ec2.DescribeVpcEndpointServiceConfigurationsOutput, error) {
	f.calls++
	state := ec2types.DnsNameStatePendingVerification
	if f.fail {
		state = ec2types.DnsNameStateFailed
	} else if f.calls >= f.verifyAt {
		state = ec2types.DnsNameStateVerified
	}
	return &ec2.DescribeVpcEndpointServiceConfigurationsOutput{
		ServiceConfigurations: []ec2types.ServiceConfiguration{{
			ServiceId: awssdk.String("vpce-svc-proof"),
			PrivateDnsNameConfiguration: &ec2types.PrivateDnsNameConfiguration{
				State: state,
			},
		}},
	}, nil
}

func TestWaitForO11YEndpointServicePrivateDNSUsesCallerDeadline(t *testing.T) {
	client := &o11yPrivateDNSWaitEC2{verifyAt: 4}
	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()

	started := time.Now()
	err := waitForO11YEndpointServicePrivateDNS(
		ctx,
		client,
		"vpce-svc-proof",
		10*time.Millisecond,
	)
	if err != nil {
		t.Fatalf("wait for private DNS: %v", err)
	}
	if elapsed := time.Since(started); elapsed <= 5*time.Millisecond {
		t.Fatalf("wait returned after %s, want verification after the old inner deadline", elapsed)
	}
	if client.calls != 4 {
		t.Fatalf("describe calls = %d, want 4", client.calls)
	}
}

func TestWaitForO11YEndpointServicePrivateDNSFailsClosed(t *testing.T) {
	client := &o11yPrivateDNSWaitEC2{fail: true}
	err := waitForO11YEndpointServicePrivateDNS(
		context.Background(), client, "vpce-svc-proof", time.Millisecond,
	)
	if err == nil || !strings.Contains(err.Error(), "private DNS verification failed") {
		t.Fatalf("wait for failed private DNS = %v, want verification failure", err)
	}
}

func TestWaitForO11YEndpointServicePrivateDNSStopsWithCallerContext(t *testing.T) {
	client := &o11yPrivateDNSWaitEC2{verifyAt: 100}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()

	err := waitForO11YEndpointServicePrivateDNS(
		ctx, client, "vpce-svc-proof", time.Hour,
	)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("wait for pending private DNS = %v, want caller deadline", err)
	}
}