package protocol

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"testing"

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

type smtpScriptResult struct {
	mu                sync.Mutex
	commands          []string
	message           string
	authBeforeTLS     bool
	appendSentCommand bool
	dataResponse      string
}

func (result *smtpScriptResult) add(command string) {
	result.mu.Lock()
	result.commands = append(result.commands, command)
	result.mu.Unlock()
}

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

	connector, server, certificate := newScriptedConnector(t)
	result := &smtpScriptResult{}
	serverDone := make(chan error, 1)
	go func() {
		serverDone <- serveSMTP(
			server,
			certificate,
			true,
			false,
			result,
		)
	}()

	client := NewSMTPClient(ClientOptions{Connector: connector})
	capabilities, err := client.ProbeSender(
		contextWithTestTimeout(t),
		Endpoint{
			Host:     "mail.example.com",
			Port:     587,
			TLSMode:  "starttls",
			Username: "user",
			Password: "secret",
		},
		"sender@example.com",
	)
	require.NoError(t, err)
	require.NoError(t, <-serverDone)
	require.True(t, capabilities.Has("8BITMIME"))
	require.True(t, capabilities.Has("SMTPUTF8"))
	require.False(t, result.authBeforeTLS)
	requireCommandOrder(t, result.commands,
		"STARTTLS",
		"AUTH ",
		"MAIL FROM:<sender@example.com>",
		"RSET",
	)
}

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

	connector, server, certificate := newScriptedConnector(t)
	result := &smtpScriptResult{}
	serverDone := make(chan error, 1)
	go func() {
		serverDone <- serveSMTP(
			server,
			certificate,
			false,
			false,
			result,
		)
	}()
	client := NewSMTPClient(ClientOptions{Connector: connector})
	send, err := client.Send(
		contextWithTestTimeout(t),
		Endpoint{
			Host:     "mail.example.com",
			Port:     465,
			TLSMode:  "implicit",
			Username: "user",
			Password: "secret",
		},
		FrozenMessage{
			SenderAddress: "sender@example.com",
			To:            []string{"to@example.com"},
			BCC:           []string{"hidden@example.com"},
			Subject:       "Test",
			BodyFormat:    "text",
			Body:          "Body",
		},
	)
	require.NoError(t, err)
	require.Equal(t, ResultAccepted, send.Status)
	require.NoError(t, <-serverDone)
	require.Contains(t, result.message, "Subject: Test")
	require.NotContains(t, result.message, "Bcc:")
	requireCommandOrder(t, result.commands,
		"MAIL FROM:<sender@example.com>",
		"RCPT TO:<to@example.com>",
		"RCPT TO:<hidden@example.com>",
		"DATA",
	)
	require.False(t, result.appendSentCommand)
}

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

	connector, server, certificate := newScriptedConnector(t)
	result := &smtpScriptResult{}
	serverDone := make(chan error, 1)
	go func() {
		serverDone <- serveSMTP(
			server,
			certificate,
			false,
			true,
			result,
		)
	}()
	client := NewSMTPClient(ClientOptions{Connector: connector})
	send, err := client.Send(
		contextWithTestTimeout(t),
		Endpoint{
			Host:     "mail.example.com",
			Port:     465,
			TLSMode:  "implicit",
			Username: "user",
			Password: "secret",
		},
		FrozenMessage{
			SenderAddress: "sender@example.com",
			To:            []string{"to@example.com"},
			Subject:       "Test",
			BodyFormat:    "text",
			Body:          "Body",
		},
	)
	require.Error(t, err)
	require.Equal(t, ResultUnknown, send.Status)
	require.NoError(t, <-serverDone)
}

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

	connector := &staticConnector{}
	client := NewSMTPClient(ClientOptions{Connector: connector})
	send, err := client.Send(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		FrozenMessage{
			SenderAddress: "sender@example.com",
			To:            []string{"to@example.com"},
			Subject:       "Test",
			BodyFormat:    "text",
			Body:          strings.Repeat("x", maxDraftBodyBytes+1),
		},
	)
	require.Error(t, err)
	require.Equal(t, ResultFailed, send.Status)
	require.Zero(t, connector.calls)
}

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

	connector, server, certificate := newScriptedConnector(t)
	result := &smtpScriptResult{dataResponse: "550 message rejected"}
	serverDone := make(chan error, 1)
	go func() {
		serverDone <- serveSMTP(
			server,
			certificate,
			false,
			false,
			result,
		)
	}()
	client := NewSMTPClient(ClientOptions{Connector: connector})
	send, err := client.Send(
		contextWithTestTimeout(t),
		Endpoint{
			Host:     "mail.example.com",
			Port:     465,
			TLSMode:  "implicit",
			Username: "user",
			Password: "secret",
		},
		FrozenMessage{
			SenderAddress: "sender@example.com",
			To:            []string{"to@example.com"},
			Subject:       "Test",
			BodyFormat:    "text",
			Body:          "Body",
		},
	)
	require.Error(t, err)
	require.Equal(t, ResultFailed, send.Status)
	require.Equal(t, ErrorFailed, ErrorClass(err))
	require.NoError(t, <-serverDone)
}

func serveSMTP(
	raw net.Conn,
	certificate tls.Certificate,
	startTLS bool,
	dropAfterData bool,
	result *smtpScriptResult,
) error {
	defer raw.Close()
	connection := raw
	secured := !startTLS
	if secured {
		tlsConnection := tls.Server(raw, &tls.Config{
			Certificates: []tls.Certificate{certificate},
			MinVersion:   tls.VersionTLS12,
		})
		if err := tlsConnection.Handshake(); err != nil {
			return err
		}
		connection = tlsConnection
	}
	reader := bufio.NewReader(connection)
	writer := bufio.NewWriter(connection)
	if _, err := writer.WriteString("220 mail.example.com ready\r\n"); err != nil {
		return err
	}
	if err := writer.Flush(); err != nil {
		return err
	}
	authenticated := false
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
		command := strings.TrimSpace(line)
		result.add(command)
		upper := strings.ToUpper(command)
		switch {
		case strings.HasPrefix(upper, "EHLO "):
			if startTLS && !secured {
				_, err = writer.WriteString(
					"250-mail.example.com\r\n250 STARTTLS\r\n",
				)
			} else {
				_, err = writer.WriteString(
					"250-mail.example.com\r\n" +
						"250-AUTH PLAIN\r\n" +
						"250-8BITMIME\r\n" +
						"250 SMTPUTF8\r\n",
				)
			}
		case upper == "STARTTLS":
			if secured {
				return fmt.Errorf("duplicate STARTTLS")
			}
			_, err = writer.WriteString("220 begin TLS\r\n")
			if err == nil {
				err = writer.Flush()
			}
			if err != nil {
				return err
			}
			tlsConnection := tls.Server(raw, &tls.Config{
				Certificates: []tls.Certificate{certificate},
				MinVersion:   tls.VersionTLS12,
			})
			if err := tlsConnection.Handshake(); err != nil {
				return err
			}
			connection = tlsConnection
			reader = bufio.NewReader(connection)
			writer = bufio.NewWriter(connection)
			secured = true
			continue
		case strings.HasPrefix(upper, "AUTH "):
			if !secured {
				result.authBeforeTLS = true
				return fmt.Errorf("AUTH before TLS")
			}
			authenticated = true
			_, err = writer.WriteString("235 authenticated\r\n")
		case strings.HasPrefix(upper, "MAIL FROM:"):
			if !authenticated {
				return fmt.Errorf("MAIL before AUTH")
			}
			_, err = writer.WriteString("250 sender accepted\r\n")
		case strings.HasPrefix(upper, "RCPT TO:"):
			_, err = writer.WriteString("250 recipient accepted\r\n")
		case upper == "RSET":
			_, err = writer.WriteString("250 reset\r\n")
		case upper == "DATA":
			_, err = writer.WriteString("354 end with dot\r\n")
			if err == nil {
				err = writer.Flush()
			}
			if err != nil {
				return err
			}
			var message strings.Builder
			for {
				dataLine, readErr := reader.ReadString('\n')
				if readErr != nil {
					return readErr
				}
				if dataLine == ".\r\n" {
					break
				}
				message.WriteString(dataLine)
			}
			result.message = message.String()
			if dropAfterData {
				return nil
			}
			response := result.dataResponse
			if response == "" {
				response = "250 queued"
			}
			_, err = writer.WriteString(response + "\r\n")
		case upper == "QUIT":
			_, err = writer.WriteString("221 bye\r\n")
			if err == nil {
				err = writer.Flush()
			}
			return err
		default:
			return fmt.Errorf("unexpected SMTP command %q", command)
		}
		if err != nil {
			return err
		}
		if err := writer.Flush(); err != nil {
			return err
		}
	}
}

func requireCommandOrder(
	t *testing.T,
	commands []string,
	prefixes ...string,
) {
	t.Helper()
	index := 0
	for _, command := range commands {
		if index < len(prefixes) &&
			strings.HasPrefix(
				strings.ToUpper(command),
				strings.ToUpper(prefixes[index]),
			) {
			index++
		}
	}
	require.Equal(t, len(prefixes), index, commands)
}

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

	connector, server, certificate := newScriptedConnector(t)
	result := &smtpScriptResult{}
	serverDone := make(chan error, 1)
	go func() {
		serverDone <- serveSMTP(
			server,
			certificate,
			false,
			false,
			result,
		)
	}()
	client := NewSMTPClient(ClientOptions{Connector: connector})
	blob := []byte("%PDF-1.4 smtp-attachment")
	send, err := client.Send(
		contextWithTestTimeout(t),
		Endpoint{
			Host:     "mail.example.com",
			Port:     465,
			TLSMode:  "implicit",
			Username: "user",
			Password: "secret",
		},
		FrozenMessage{
			SenderAddress: "sender@example.com",
			To:            []string{"to@example.com"},
			BCC:           []string{"hidden@example.com"},
			Subject:       "With attachment",
			BodyFormat:    "text",
			Body:          "Body",
			Attachments: []FrozenAttachment{{
				Filename:  "report.pdf",
				MediaType: "application/pdf",
				Blob:      blob,
			}},
		},
	)
	require.NoError(t, err)
	require.Equal(t, ResultAccepted, send.Status)
	require.NoError(t, <-serverDone)
	require.Contains(t, result.message, "Subject: With attachment")
	require.Contains(t, result.message, "multipart/mixed")
	require.Contains(t, result.message, "Content-Disposition: attachment; filename=report.pdf")
	require.NotContains(t, result.message, "Bcc:")
	requireCommandOrder(t, result.commands,
		"MAIL FROM:<sender@example.com>",
		"RCPT TO:<to@example.com>",
		"RCPT TO:<hidden@example.com>",
		"DATA",
	)
	require.False(t, result.appendSentCommand)
}

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

	connector, server, certificate := newScriptedConnector(t)
	result := &smtpScriptResult{}
	serverDone := make(chan error, 1)
	go func() {
		serverDone <- serveSMTP(
			server,
			certificate,
			false,
			false,
			result,
		)
	}()
	client := NewSMTPClient(ClientOptions{Connector: connector})
	blob := make([]byte, 3<<20)
	for index := range blob {
		blob[index] = 'S'
	}
	send, err := client.Send(
		contextWithTestTimeout(t),
		Endpoint{
			Host:     "mail.example.com",
			Port:     465,
			TLSMode:  "implicit",
			Username: "user",
			Password: "secret",
		},
		FrozenMessage{
			SenderAddress: "sender@example.com",
			To:            []string{"to@example.com"},
			Subject:       "Large attachment",
			BodyFormat:    "text",
			Body:          "Body",
			Attachments: []FrozenAttachment{{
				Filename:  "large.bin",
				MediaType: "application/octet-stream",
				Blob:      blob,
			}},
		},
	)
	require.NoError(t, err)
	require.Equal(t, ResultAccepted, send.Status)
	require.NoError(t, <-serverDone)
	// The encoded base64 body is 4 MiB, far above the retired 2 MiB global
	// connection budget; the whole message must reach the DATA writer.
	require.Contains(t, result.message, "Subject: Large attachment")
	require.Greater(t, len(result.message), 4*1024*1024)
}

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

	connector, _, _ := newScriptedConnector(t)
	client := NewSMTPClient(ClientOptions{Connector: connector})
	blob := make([]byte, maxSMTPMessageBytes)
	send, err := client.Send(
		contextWithTestTimeout(t),
		Endpoint{
			Host:     "mail.example.com",
			Port:     465,
			TLSMode:  "implicit",
			Username: "user",
			Password: "secret",
		},
		FrozenMessage{
			SenderAddress: "sender@example.com",
			To:            []string{"to@example.com"},
			Subject:       "Over budget",
			BodyFormat:    "text",
			Body:          "Body",
			Attachments: []FrozenAttachment{{
				Filename:  "huge.bin",
				MediaType: "application/octet-stream",
				Blob:      blob,
			}},
		},
	)
	require.Error(t, err)
	require.Equal(t, ResultFailed, send.Status)
	require.Equal(t, ErrorInvalid, ErrorClass(err))
}