package protocol

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

	"github.com/emersion/go-imap"
	"github.com/stretchr/testify/require"
)

const scriptedDraftMailbox = "Lokalisert/Utkast"

type highLevelIMAPScript struct {
	mu                 *sync.Mutex
	commands           []string
	appendedMessage    []byte
	dropAfterAppend    bool
	dropAfterStore     bool
	modifiedOnStore    bool
	omitDraftUse       bool
	oldMissing         bool
	oldModSeq          uint64
	omitAppendedModSeq bool
	searchCount        int
	attachmentBody     []byte
	attachmentRawSize  uint32
	attachmentFilename string
	bodyContent        []byte
	bodyEncoding       string
	bodyMediaType      string
	bodyRawSize        uint32
	bodySizeSet        bool
	appendChunkSize    int
	appendChunkDelay   time.Duration
}

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

func (script *highLevelIMAPScript) commandContaining(value string) string {
	script.mu.Lock()
	defer script.mu.Unlock()
	for _, command := range script.commands {
		if strings.Contains(strings.ToUpper(command), strings.ToUpper(value)) {
			return command
		}
	}
	return ""
}

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

	client, script, done := newHighLevelIMAPClient(t, highLevelIMAPScript{
		searchCount: 60,
	})
	page, err := client.Search(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		SearchRequest{Folder: FolderInbox, Limit: 50},
	)
	require.NoError(t, err)
	require.NoError(t, <-done)
	require.Len(t, page.Messages, 50)
	require.Equal(t, uint32(60), page.Messages[0].Ref.UID)
	require.Equal(t, uint32(11), page.Messages[49].Ref.UID)
	require.Equal(t, uint32(11), page.BeforeUID)
	require.Contains(t, script.commandContaining("UID FETCH"), "10:60")
	require.NotContains(t, script.commandContaining("UID FETCH"), "1:60")
}

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

	client, _, done := newHighLevelIMAPClient(t, highLevelIMAPScript{})
	capabilities, err := client.Probe(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
	)
	require.NoError(t, err)
	require.NoError(t, <-done)
	require.Equal(t, imap.InboxName, capabilities.Mailbox(FolderInbox))
	require.Equal(
		t,
		scriptedDraftMailbox,
		capabilities.Mailbox(FolderDrafts),
	)
	require.Equal(t, "Outbox/Sent", capabilities.Mailbox(FolderSent))
	require.Equal(t, "Cabinet", capabilities.Mailbox(FolderArchive))

	client, script, done := newHighLevelIMAPClient(t, highLevelIMAPScript{
		omitDraftUse: true,
	})
	_, err = client.AppendDraft(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		testDraftContent(),
	)
	require.Equal(t, ErrorUnsupported, ErrorClass(err))
	require.NoError(t, <-done)
	require.Empty(t, script.commandContaining("APPEND"))
}

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

	connector, server, certificate := newScriptedConnector(t)
	done := make(chan error, 1)
	go func() {
		defer server.Close()
		connection := tls.Server(server, &tls.Config{
			Certificates: []tls.Certificate{certificate},
			MinVersion:   tls.VersionTLS12,
		})
		if err := connection.Handshake(); err != nil {
			done <- err
			return
		}
		// A syntactically valid but oversized greeting: the reader only
		// stops on budget exhaustion, not on a parse error, so the
		// observable exhaustion seam classifies the failure as a limit.
		payload := "* OK " + strings.Repeat("x", maxProtocolBytes+1) + "\r\n"
		_, _ = io.WriteString(connection, payload)
		done <- nil
	}()
	client := NewIMAPClient(ClientOptions{Connector: connector})
	_, err := client.Probe(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
	)
	require.Error(t, err)
	require.Equal(t, ErrorLimit, ErrorClass(err))
	require.NoError(t, <-done)
}

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

	client, script, done := newHighLevelIMAPClient(t, highLevelIMAPScript{})
	message, err := client.GetMessage(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		MessageRef{
			Mailbox:     imap.InboxName,
			UIDValidity: 777,
			UID:         7,
		},
	)
	require.NoError(t, err)
	require.NoError(t, <-done)
	require.Equal(t, "Hello", message.Body.Content)
	require.True(t, message.HasAttachments)
	require.Contains(t, script.commandContaining("BODY.PEEK"), "BODY.PEEK[1]")
	require.NotContains(t, script.commandContaining("BODY.PEEK"), "[2]")

	client, script, done = newHighLevelIMAPClient(t, highLevelIMAPScript{})
	attachments, err := client.ListAttachments(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		MessageRef{
			Mailbox:     imap.InboxName,
			UIDValidity: 777,
			UID:         7,
		},
	)
	require.NoError(t, err)
	require.NoError(t, <-done)
	require.Equal(t, []Attachment{{
		Part:      []int{2},
		Filename:  "report.pdf",
		MediaType: "application/pdf",
		SizeBytes: 123,
	}}, attachments)
	require.Empty(t, script.commandContaining("BODY.PEEK"))
}

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

	client, script, done := newHighLevelIMAPClient(t, highLevelIMAPScript{})
	ref, err := client.AppendDraft(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		testDraftContent(),
	)
	require.NoError(t, err)
	require.NoError(t, <-done)
	require.Equal(t, DraftRef{
		Mailbox:     scriptedDraftMailbox,
		UIDValidity: 777,
		UID:         42,
		ModSeq:      11,
	}, ref)
	require.Contains(t, script.commandContaining("APPEND"), scriptedDraftMailbox)
	require.Contains(t, string(script.appendedMessage), "Subject: Project update")
	require.Empty(t, script.commandContaining("APPEND \"SENT"))
}

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

	client, _, done := newHighLevelIMAPClient(t, highLevelIMAPScript{})
	draft, err := client.GetDraft(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		DraftRef{
			Mailbox:     scriptedDraftMailbox,
			UIDValidity: 777,
			UID:         7,
			ModSeq:      9,
		},
	)
	require.NoError(t, err)
	require.NoError(t, <-done)
	require.Equal(t, uint64(9), draft.Ref.ModSeq)
	require.Equal(t, uint32(7), draft.Ref.UID)
	require.Equal(t, uint32(777), draft.Ref.UIDValidity)
	require.Equal(t, "sender@example.com", draft.Content.From)
	require.Equal(t, "Hello", draft.Content.Body)
}

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

	addresses := make([]*imap.Address, 0, maxAddresses+1)
	for index := 0; index <= maxAddresses; index++ {
		addresses = append(addresses, &imap.Address{
			MailboxName: fmt.Sprintf("hidden-%d", index),
			HostName:    "example.com",
		})
	}
	content, truncated, err := draftContent(
		&imap.Envelope{
			From: []*imap.Address{{
				MailboxName: "sender",
				HostName:    "example.com",
			}},
			To:      []*imap.Address{addresses[0]},
			Bcc:     addresses,
			Subject: strings.Repeat("s", maxSubjectBytes+1),
		},
		Body{Format: "text", Content: "body"},
	)
	require.NoError(t, err)
	require.True(t, truncated)
	require.Len(t, content.BCC, maxAddresses)
	require.Len(t, content.Subject, maxSubjectBytes)
}

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

	client, script, done := newHighLevelIMAPClient(t, highLevelIMAPScript{})
	ref, err := client.ReplaceDraft(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		DraftRef{
			Mailbox:     scriptedDraftMailbox,
			UIDValidity: 777,
			UID:         7,
			ModSeq:      9,
		},
		testDraftContent(),
	)
	require.NoError(t, err)
	require.NoError(t, <-done)
	require.Equal(t, uint32(42), ref.UID)
	require.Contains(
		t,
		script.commandContaining("UNCHANGEDSINCE"),
		"(UNCHANGEDSINCE 9)",
	)
	require.Contains(t, script.commandContaining("UID EXPUNGE"), "7")
}

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

	for _, test := range []struct {
		name   string
		script highLevelIMAPScript
	}{
		{name: "changed", script: highLevelIMAPScript{oldModSeq: 10}},
		{name: "missing", script: highLevelIMAPScript{oldMissing: true}},
	} {
		test := test
		t.Run(test.name, func(t *testing.T) {
			client, script, done := newHighLevelIMAPClient(t, test.script)
			_, err := client.ReplaceDraft(
				contextWithTestTimeout(t),
				testIMAPEndpoint(),
				DraftRef{
					Mailbox:     scriptedDraftMailbox,
					UIDValidity: 777,
					UID:         7,
					ModSeq:      9,
				},
				testDraftContent(),
			)
			require.Equal(t, ErrorConflict, ErrorClass(err))
			require.NoError(t, <-done)
			require.Empty(t, script.commandContaining("APPEND"))
		})
	}
}

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

	client, _, done := newHighLevelIMAPClient(t, highLevelIMAPScript{
		dropAfterAppend: true,
	})
	_, err := client.AppendDraft(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		testDraftContent(),
	)
	require.Equal(t, ErrorUnknown, ErrorClass(err))
	require.NoError(t, <-done)
}

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

	for _, test := range []struct {
		name      string
		script    highLevelIMAPScript
		wantClass string
	}{
		{
			name:      "changed_after_append",
			script:    highLevelIMAPScript{modifiedOnStore: true},
			wantClass: ErrorUnknown,
		},
		{
			name:      "disconnect_after_store",
			script:    highLevelIMAPScript{dropAfterStore: true},
			wantClass: ErrorUnknown,
		},
	} {
		test := test
		t.Run(test.name, func(t *testing.T) {
			client, script, done := newHighLevelIMAPClient(t, test.script)
			_, err := client.ReplaceDraft(
				contextWithTestTimeout(t),
				testIMAPEndpoint(),
				DraftRef{
					Mailbox:     scriptedDraftMailbox,
					UIDValidity: 777,
					UID:         7,
					ModSeq:      9,
				},
				testDraftContent(),
			)
			require.Equal(t, test.wantClass, ErrorClass(err))
			require.NoError(t, <-done)
			require.NotEmpty(t, script.commandContaining("APPEND"))
			require.NotEmpty(t, script.commandContaining("UNCHANGEDSINCE"))
		})
	}
}

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

	client, _, done := newHighLevelIMAPClient(t, highLevelIMAPScript{
		omitAppendedModSeq: true,
	})
	ref, err := client.AppendDraft(
		contextWithTestTimeout(t),
		testIMAPEndpoint(),
		testDraftContent(),
	)
	require.NoError(t, err)
	require.NoError(t, <-done)
	require.Equal(t, uint32(777), ref.UIDValidity)
	require.Equal(t, uint32(42), ref.UID)
	require.Zero(t, ref.ModSeq)
}

func newHighLevelIMAPClient(
	t *testing.T,
	initial highLevelIMAPScript,
) (*IMAPClient, *highLevelIMAPScript, <-chan error) {
	t.Helper()
	connector, server, certificate := newScriptedConnector(t)
	script := &initial
	script.mu = new(sync.Mutex)
	done := make(chan error, 1)
	go func() {
		done <- serveHighLevelIMAP(server, certificate, script)
	}()
	return NewIMAPClient(ClientOptions{Connector: connector}), script, done
}

func testIMAPEndpoint() Endpoint {
	return Endpoint{
		Host:     "mail.example.com",
		Port:     993,
		TLSMode:  "implicit",
		Username: "user",
		Password: "secret",
	}
}

func testDraftContent() DraftContent {
	return DraftContent{
		From:       "sender@example.com",
		To:         []string{"to@example.com"},
		Subject:    "Project update",
		BodyFormat: "text",
		Body:       "Hello",
	}
}

func serveHighLevelIMAP(
	raw net.Conn,
	certificate tls.Certificate,
	script *highLevelIMAPScript,
) error {
	defer raw.Close()
	connection := tls.Server(raw, &tls.Config{
		Certificates: []tls.Certificate{certificate},
		MinVersion:   tls.VersionTLS12,
	})
	if err := connection.Handshake(); err != nil {
		return err
	}
	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)
		script.add(command)
		fields := strings.Fields(command)
		if len(fields) < 2 {
			return fmt.Errorf("invalid IMAP command")
		}
		tag := fields[0]
		verb := strings.ToUpper(fields[1])
		switch verb {
		case "CAPABILITY":
			err = writeTagged(
				writer,
				tag,
				"* CAPABILITY IMAP4rev1 UIDPLUS CONDSTORE SPECIAL-USE\r\n",
				"OK CAPABILITY",
			)
		case "LOGIN":
			err = writeTagged(writer, tag, "", "OK LOGIN")
		case "LIST":
			draftAttributes := "\\HasNoChildren \\Drafts"
			draftMailbox := scriptedDraftMailbox
			if script.omitDraftUse {
				draftAttributes = "\\HasNoChildren"
				draftMailbox = "Drafts"
			}
			err = writeTagged(
				writer,
				tag,
				"* LIST ("+draftAttributes+") \"/\" \""+
					draftMailbox+"\"\r\n"+
					"* LIST (\\HasNoChildren \\Sent) \"/\" \"Outbox/Sent\"\r\n"+
					"* LIST (\\HasNoChildren \\Archive) \"/\" \"Cabinet\"\r\n",
				"OK LIST",
			)
		case "SELECT", "EXAMINE":
			err = writeTagged(
				writer,
				tag,
				"* FLAGS (\\Answered \\Flagged \\Deleted \\Seen \\Draft)\r\n"+
					"* 2 EXISTS\r\n"+
					"* OK [UIDVALIDITY 777] valid\r\n",
				"OK [READ-WRITE] SELECT",
			)
		case "APPEND":
			err = readScriptedAppend(reader, writer, tag, command, script)
			if script.dropAfterAppend {
				return err
			}
		case "UID":
			err = serveUIDCommand(writer, tag, command, script)
			if script.dropAfterStore &&
				strings.Contains(strings.ToUpper(command), " STORE ") {
				return err
			}
		case "LOGOUT":
			if _, err = writer.WriteString(
				"* BYE logout\r\n" + tag + " OK LOGOUT\r\n",
			); err == nil {
				err = writer.Flush()
			}
			return err
		default:
			return fmt.Errorf("unexpected IMAP command")
		}
		if err != nil {
			return err
		}
	}
}

func readScriptedAppend(
	reader *bufio.Reader,
	writer *bufio.Writer,
	tag string,
	command string,
	script *highLevelIMAPScript,
) error {
	start := strings.LastIndex(command, "{")
	end := strings.LastIndex(command, "}")
	if start < 0 || end <= start {
		return fmt.Errorf("APPEND literal is missing")
	}
	size, err := strconv.Atoi(strings.TrimSuffix(
		command[start+1:end],
		"+",
	))
	if err != nil {
		return err
	}
	if _, err := writer.WriteString("+ continue\r\n"); err != nil {
		return err
	}
	if err := writer.Flush(); err != nil {
		return err
	}
	payload := make([]byte, size+2)
	if script.appendChunkDelay > 0 && script.appendChunkSize > 0 {
		// Simulate a slow link: consume the literal in throttled chunks so
		// the append spans longer than any metadata command timeout.
		buffer := bufio.NewReader(reader)
		for read := 0; read < len(payload); {
			chunk := script.appendChunkSize
			if remaining := len(payload) - read; remaining < chunk {
				chunk = remaining
			}
			if _, err := io.ReadFull(buffer, payload[read:read+chunk]); err != nil {
				return err
			}
			read += chunk
			// A cancellable timer paces the simulated slow link instead of a
			// bare fixed sleep.
			timer := time.NewTimer(script.appendChunkDelay)
			<-timer.C
		}
	} else if _, err := io.ReadFull(reader, payload); err != nil {
		return err
	}
	script.appendedMessage = append([]byte(nil), payload[:size]...)
	if script.dropAfterAppend {
		return nil
	}
	return writeTagged(writer, tag, "", "OK [APPENDUID 777 42] APPEND")
}

func serveUIDCommand(
	writer *bufio.Writer,
	tag string,
	command string,
	script *highLevelIMAPScript,
) error {
	upper := strings.ToUpper(command)
	switch {
	case strings.Contains(upper, " SEARCH"):
		var response strings.Builder
		response.WriteString("* SEARCH")
		for uid := 1; uid <= script.searchCount; uid++ {
			fmt.Fprintf(&response, " %d", uid)
		}
		response.WriteString("\r\n")
		return writeTagged(writer, tag, response.String(), "OK SEARCH")
	case strings.Contains(upper, " STORE "):
		if script.dropAfterStore {
			return nil
		}
		if script.modifiedOnStore {
			return writeTagged(
				writer,
				tag,
				"",
				"OK [MODIFIED 7] STORE",
			)
		}
		message := scriptedMessage(1, 7, 10, true, false)
		if err := writeIMAPFetch(writer, message); err != nil {
			return err
		}
		return writeTagged(writer, tag, "", "OK STORE")
	case strings.Contains(upper, " EXPUNGE "):
		return writeTagged(writer, tag, "* 1 EXPUNGE\r\n", "OK UID EXPUNGE")
	case strings.Contains(upper, " FETCH "):
		return serveUIDFetch(writer, tag, command, script)
	default:
		return fmt.Errorf("unexpected UID command")
	}
}

func serveUIDFetch(
	writer *bufio.Writer,
	tag string,
	command string,
	script *highLevelIMAPScript,
) error {
	fields := strings.Fields(command)
	if len(fields) < 4 {
		return fmt.Errorf("UID FETCH sequence is missing")
	}
	sequence, err := imap.ParseSeqSet(fields[3])
	if err != nil {
		return err
	}
	upper := strings.ToUpper(command)
	bodyFetch := strings.Contains(upper, "BODY.PEEK")
	for _, item := range sequence.Set {
		stop := item.Stop
		if stop == 0 {
			stop = item.Start
		}
		for uid := item.Start; uid <= stop; uid++ {
			if uid == 7 && script.oldMissing {
				continue
			}
			modSeq := uint64(5)
			if uid == 7 {
				modSeq = script.oldModSeq
				if modSeq == 0 {
					modSeq = 9
				}
			}
			if uid == 42 {
				modSeq = 11
			}
			message := scriptedMessage(uid, uid, modSeq, false, bodyFetch)
			applyScriptedPartOverrides(script, message)
			if bodyFetch && strings.Contains(upper, "[2]") {
				section, sectionErr := imap.ParseBodySectionName(
					imap.FetchItem("BODY[2]<0>"),
				)
				if sectionErr != nil {
					return sectionErr
				}
				message.Items[section.FetchItem()] = nil
				message.Body[section] = bytes.NewBuffer(script.attachmentBody)
			}
			if uid == 42 && script.omitAppendedModSeq {
				delete(message.Items, imap.FetchItem("MODSEQ"))
			}
			if err := writeIMAPFetch(writer, message); err != nil {
				return err
			}
			if uid == ^uint32(0) {
				break
			}
		}
	}
	return writeTagged(writer, tag, "", "OK FETCH")
}

func scriptedMessage(
	sequence uint32,
	uid uint32,
	modSeq uint64,
	deleted bool,
	withBody bool,
) *imap.Message {
	textPart := &imap.BodyStructure{
		MIMEType:    "text",
		MIMESubType: "plain",
		Params:      map[string]string{"charset": "utf-8"},
		Encoding:    "7BIT",
		Size:        5,
		Lines:       1,
		Extended:    true,
	}
	attachment := &imap.BodyStructure{
		MIMEType:          "application",
		MIMESubType:       "pdf",
		Params:            map[string]string{"name": "report.pdf"},
		Encoding:          "BASE64",
		Size:              123,
		Extended:          true,
		Disposition:       "attachment",
		DispositionParams: map[string]string{"filename": "report.pdf"},
	}
	structure := &imap.BodyStructure{
		MIMEType:    "multipart",
		MIMESubType: "mixed",
		Parts:       []*imap.BodyStructure{textPart, attachment},
		Params:      map[string]string{},
		Extended:    true,
	}
	items := []imap.FetchItem{
		imap.FetchUid,
		imap.FetchEnvelope,
		imap.FetchFlags,
		imap.FetchInternalDate,
		imap.FetchRFC822Size,
		imap.FetchBodyStructure,
		imap.FetchItem("MODSEQ"),
	}
	var section *imap.BodySectionName
	if withBody {
		section, _ = imap.ParseBodySectionName(
			imap.FetchItem("BODY[1]<0>"),
		)
		items = append(items, section.FetchItem())
	}
	message := imap.NewMessage(sequence, items)
	message.Uid = uid
	message.Envelope = &imap.Envelope{
		Date:    time.Date(2026, 8, 15, 10, 0, 0, 0, time.UTC),
		Subject: "Project update",
		From: []*imap.Address{{
			MailboxName: "sender",
			HostName:    "example.com",
		}},
		To: []*imap.Address{{
			MailboxName: "to",
			HostName:    "example.com",
		}},
	}
	message.Flags = []string{imap.DraftFlag}
	if deleted {
		message.Flags = append(message.Flags, imap.DeletedFlag)
	}
	message.InternalDate = message.Envelope.Date
	message.Size = 128
	message.BodyStructure = structure
	message.Items[imap.FetchItem("MODSEQ")] = []interface{}{
		imap.RawString(strconv.FormatUint(modSeq, 10)),
	}
	if withBody {
		message.Body[section] = bytes.NewBufferString("Hello")
	}
	return message
}

func writeIMAPFetch(
	writer *bufio.Writer,
	message *imap.Message,
) error {
	response := imap.NewUntaggedResp([]interface{}{
		message.SeqNum,
		imap.RawString("FETCH"),
		message.Format(),
	})
	if err := response.WriteTo(imap.NewWriter(writer)); err != nil {
		return err
	}
	return writer.Flush()
}

func writeTagged(
	writer *bufio.Writer,
	tag string,
	untagged string,
	status string,
) error {
	if _, err := writer.WriteString(
		untagged + tag + " " + status + "\r\n",
	); err != nil {
		return err
	}
	return writer.Flush()
}