package protocol

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

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

type imapScriptResult struct {
	mu             sync.Mutex
	commands       []string
	loginBeforeTLS bool
	mailboxCount   int
}

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

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

	connector, server, certificate := newScriptedConnector(t)
	result := &imapScriptResult{}
	serverDone := make(chan error, 1)
	go func() {
		serverDone <- serveIMAP(
			server,
			certificate,
			true,
			result,
		)
	}()
	client := NewIMAPClient(ClientOptions{Connector: connector})
	capabilities, err := client.Probe(
		contextWithTestTimeout(t),
		Endpoint{
			Host:     "mail.example.com",
			Port:     143,
			TLSMode:  "starttls",
			Username: "user",
			Password: "secret",
		},
	)
	require.NoError(t, err)
	require.NoError(t, <-serverDone)
	require.Equal(t, "Drafts", capabilities.DraftMailbox)
	require.True(t, capabilities.Has("UIDPLUS"))
	require.True(t, capabilities.Has("CONDSTORE"))
	require.True(t, capabilities.SafeDraftMutation())
	require.False(t, result.loginBeforeTLS)
	requireIMAPVerbOrder(
		t,
		result.commands,
		"STARTTLS",
		"CAPABILITY",
		"LOGIN",
		"CAPABILITY",
		"LIST",
	)
}

func requireIMAPVerbOrder(
	t *testing.T,
	commands []string,
	verbs ...string,
) {
	t.Helper()
	index := 0
	for _, command := range commands {
		fields := strings.Fields(command)
		if len(fields) >= 2 &&
			index < len(verbs) &&
			strings.EqualFold(fields[1], verbs[index]) {
			index++
		}
	}
	require.Equal(t, len(verbs), index, commands)
}

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

	connector, server, certificate := newScriptedConnector(t)
	result := &imapScriptResult{}
	serverDone := make(chan error, 1)
	go func() {
		serverDone <- serveIMAPWithCapabilities(
			server,
			certificate,
			false,
			"IMAP4rev1 UIDPLUS SPECIAL-USE",
			result,
		)
	}()
	client := NewIMAPClient(ClientOptions{Connector: connector})
	capabilities, err := client.Probe(
		contextWithTestTimeout(t),
		Endpoint{
			Host:     "mail.example.com",
			Port:     993,
			TLSMode:  "implicit",
			Username: "user",
			Password: "secret",
		},
	)
	require.NoError(t, err)
	require.NoError(t, <-serverDone)
	require.Equal(t, "Drafts", capabilities.DraftMailbox)
	require.False(t, capabilities.SafeDraftMutation())
	requireIMAPVerbOrder(
		t,
		result.commands,
		"CAPABILITY",
		"LOGIN",
		"CAPABILITY",
		"LIST",
	)
}

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

	connector, server, certificate := newScriptedConnector(t)
	result := &imapScriptResult{mailboxCount: 300}
	serverDone := make(chan error, 1)
	go func() {
		serverDone <- serveIMAPWithCapabilities(
			server,
			certificate,
			false,
			"IMAP4rev1 UIDPLUS SPECIAL-USE",
			result,
		)
	}()
	client := NewIMAPClient(ClientOptions{Connector: connector})
	_, err := client.Probe(
		contextWithTestTimeout(t),
		Endpoint{
			Host:     "mail.example.com",
			Port:     993,
			TLSMode:  "implicit",
			Username: "user",
			Password: "secret",
		},
	)
	require.Equal(t, ErrorLimit, ErrorClass(err))
	require.NoError(t, <-serverDone)
}

func serveIMAP(
	raw net.Conn,
	certificate tls.Certificate,
	startTLS bool,
	result *imapScriptResult,
) error {
	return serveIMAPWithCapabilities(
		raw,
		certificate,
		startTLS,
		"IMAP4rev1 UIDPLUS CONDSTORE SPECIAL-USE",
		result,
	)
}

func serveIMAPWithCapabilities(
	raw net.Conn,
	certificate tls.Certificate,
	startTLS bool,
	capabilities string,
	result *imapScriptResult,
) 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("* OK IMAP ready\r\n"); err != nil {
		return err
	}
	if err := writer.Flush(); err != nil {
		return err
	}
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}
		command := strings.TrimSpace(line)
		result.add(command)
		fields := strings.Fields(command)
		if len(fields) < 2 {
			return fmt.Errorf("invalid IMAP command %q", command)
		}
		tag := fields[0]
		verb := strings.ToUpper(fields[1])
		switch verb {
		case "CAPABILITY":
			exposed := capabilities
			if startTLS && !secured {
				exposed += " STARTTLS"
			}
			_, err = fmt.Fprintf(
				writer,
				"* CAPABILITY %s\r\n%s OK CAPABILITY\r\n",
				exposed,
				tag,
			)
		case "STARTTLS":
			if secured {
				return fmt.Errorf("duplicate IMAP STARTTLS")
			}
			_, err = fmt.Fprintf(writer, "%s OK begin TLS\r\n", tag)
			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 "LOGIN":
			if !secured {
				result.loginBeforeTLS = true
				return fmt.Errorf("LOGIN before TLS")
			}
			_, err = fmt.Fprintf(writer, "%s OK LOGIN\r\n", tag)
		case "LIST":
			count := result.mailboxCount
			if count == 0 {
				count = 1
			}
			for index := 0; index < count && err == nil; index++ {
				attribute := "\\HasNoChildren"
				name := fmt.Sprintf("Mailbox-%03d", index)
				if index == 0 {
					attribute += " \\Drafts"
					name = "Drafts"
				}
				_, err = fmt.Fprintf(
					writer,
					"* LIST (%s) \"/\" %q\r\n",
					attribute,
					name,
				)
			}
			if err == nil {
				_, err = fmt.Fprintf(
					writer,
					"%s OK LIST\r\n",
					tag,
				)
			}
		case "LOGOUT":
			_, err = fmt.Fprintf(
				writer,
				"* BYE logout\r\n%s OK LOGOUT\r\n",
				tag,
			)
			if err == nil {
				err = writer.Flush()
			}
			return err
		default:
			return fmt.Errorf("unexpected IMAP command %q", command)
		}
		if err != nil {
			return err
		}
		if err := writer.Flush(); err != nil {
			return err
		}
	}
}