package transport

import (
	"context"
	"crypto/tls"
	"errors"
	"net"
	"net/netip"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestDialerPinsValidatedAddressAndPreservesTLSServerName(
	t *testing.T,
) {
	t.Parallel()

	resolver := &fakeResolver{answers: [][]netip.Addr{
		{netip.MustParseAddr("8.8.8.8")},
		{netip.MustParseAddr("10.0.0.1")},
	}}
	var network string
	var address string
	client, server := net.Pipe()
	defer server.Close()
	dialer := &Dialer{
		Resolver: resolver,
		DialContext: func(
			_ context.Context,
			gotNetwork string,
			gotAddress string,
		) (net.Conn, error) {
			network = gotNetwork
			address = gotAddress
			return client, nil
		},
	}
	result, err := dialer.ResolveAndDial(
		context.Background(),
		Target{
			Protocol: ProtocolIMAP,
			Host:     "IMAP.Example.COM",
			Port:     993,
			TLSMode:  TLSModeImplicit,
		},
	)
	require.NoError(t, err)
	defer result.Conn.Close()
	require.Equal(t, 1, resolver.calls)
	require.Equal(t, "tcp", network)
	require.Equal(t, "8.8.8.8:993", address)
	require.Equal(t, "imap.example.com", result.TLSConfig.ServerName)
	require.Equal(t, uint16(tls.VersionTLS12), result.TLSConfig.MinVersion)
	require.False(t, result.TLSConfig.InsecureSkipVerify)
}

func TestDialerRejectsForbiddenAnswerBeforeDial(t *testing.T) {
	t.Parallel()

	resolver := &fakeResolver{answers: [][]netip.Addr{{
		netip.MustParseAddr("169.254.169.254"),
	}}}
	dialCalls := 0
	dialer := &Dialer{
		Resolver: resolver,
		DialContext: func(
			context.Context,
			string,
			string,
		) (net.Conn, error) {
			dialCalls++
			return nil, errors.New("must not dial")
		},
	}
	_, err := dialer.ResolveAndDial(
		context.Background(),
		Target{
			Protocol: ProtocolSMTP,
			Host:     "smtp.example.com",
			Port:     587,
			TLSMode:  TLSModeStartTLS,
		},
	)
	require.ErrorContains(t, err, "forbidden")
	require.Zero(t, dialCalls)
}

func TestDialerTotalDeadlineCancelsBlockingResolver(t *testing.T) {
	resolver := &blockingResolver{}
	dialCalls := 0
	dialer := &Dialer{
		Resolver:     resolver,
		TotalTimeout: 50 * time.Millisecond,
		DialContext: func(
			context.Context,
			string,
			string,
		) (net.Conn, error) {
			dialCalls++
			return nil, errors.New("must not dial")
		},
	}

	started := time.Now()
	_, err := dialer.ResolveAndDial(
		context.Background(),
		Target{
			Protocol: ProtocolIMAP,
			Host:     "imap.example.com",
			Port:     993,
			TLSMode:  TLSModeImplicit,
		},
	)

	require.ErrorIs(t, err, context.DeadlineExceeded)
	require.Equal(t, 1, resolver.calls)
	require.True(t, resolver.sawDeadline)
	require.Zero(t, dialCalls)
	require.Less(t, time.Since(started), time.Second)
}

func TestDialerTotalDeadlineStopsAfterBlockingPublicAddress(
	t *testing.T,
) {
	resolver := &fakeResolver{answers: [][]netip.Addr{{
		netip.MustParseAddr("8.8.8.8"),
		netip.MustParseAddr("1.1.1.1"),
		netip.MustParseAddr("9.9.9.9"),
	}}}
	var dialed []string
	dialer := &Dialer{
		Resolver:     resolver,
		TotalTimeout: 50 * time.Millisecond,
		DialContext: func(
			ctx context.Context,
			_ string,
			address string,
		) (net.Conn, error) {
			dialed = append(dialed, address)
			if _, ok := ctx.Deadline(); !ok {
				return nil, errors.New("dial context has no deadline")
			}
			<-ctx.Done()
			return nil, ctx.Err()
		},
	}

	_, err := dialer.ResolveAndDial(
		context.Background(),
		Target{
			Protocol: ProtocolSMTP,
			Host:     "smtp.example.com",
			Port:     587,
			TLSMode:  TLSModeStartTLS,
		},
	)

	require.ErrorIs(t, err, context.DeadlineExceeded)
	require.Equal(t, []string{"8.8.8.8:587"}, dialed)
}

func TestDialerCapsPublicAddressAttempts(t *testing.T) {
	resolver := &fakeResolver{answers: [][]netip.Addr{{
		netip.MustParseAddr("8.8.8.8"),
		netip.MustParseAddr("1.1.1.1"),
		netip.MustParseAddr("9.9.9.9"),
		netip.MustParseAddr("208.67.222.222"),
		netip.MustParseAddr("8.8.4.4"),
		netip.MustParseAddr("1.0.0.1"),
	}}}
	var dialed []string
	dialer := &Dialer{
		Resolver:     resolver,
		TotalTimeout: time.Second,
		DialContext: func(
			_ context.Context,
			_ string,
			address string,
		) (net.Conn, error) {
			dialed = append(dialed, address)
			return nil, errors.New("public blackhole")
		},
	}

	_, err := dialer.ResolveAndDial(
		context.Background(),
		Target{
			Protocol: ProtocolIMAP,
			Host:     "imap.example.com",
			Port:     993,
			TLSMode:  TLSModeImplicit,
		},
	)

	require.ErrorContains(t, err, "public blackhole")
	require.Equal(t, []string{
		"8.8.8.8:993",
		"1.1.1.1:993",
		"9.9.9.9:993",
		"208.67.222.222:993",
	}, dialed)
}

func TestValidateTargetAllowsOnlyMailSubmissionPorts(t *testing.T) {
	t.Parallel()

	for _, target := range []Target{
		{
			Protocol: ProtocolIMAP,
			Host:     "imap.example.com",
			Port:     993,
			TLSMode:  TLSModeImplicit,
		},
		{
			Protocol: ProtocolIMAP,
			Host:     "imap.example.com",
			Port:     143,
			TLSMode:  TLSModeStartTLS,
		},
		{
			Protocol: ProtocolSMTP,
			Host:     "smtp.example.com",
			Port:     465,
			TLSMode:  TLSModeImplicit,
		},
		{
			Protocol: ProtocolSMTP,
			Host:     "smtp.example.com",
			Port:     587,
			TLSMode:  TLSModeStartTLS,
		},
	} {
		require.NoError(t, ValidateTarget(target))
	}
	for _, target := range []Target{
		{
			Protocol: ProtocolSMTP,
			Host:     "smtp.example.com",
			Port:     25,
			TLSMode:  TLSModeStartTLS,
		},
		{
			Protocol: ProtocolIMAP,
			Host:     "imap.example.com",
			Port:     143,
			TLSMode:  TLSModeImplicit,
		},
	} {
		require.Error(t, ValidateTarget(target))
	}
}

type blockingResolver struct {
	calls       int
	sawDeadline bool
}

func (resolver *blockingResolver) LookupNetIP(
	ctx context.Context,
	_ string,
	_ string,
) ([]netip.Addr, error) {
	resolver.calls++
	_, resolver.sawDeadline = ctx.Deadline()
	<-ctx.Done()
	return nil, ctx.Err()
}