package mailruntime

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/extension/providers/mail/mailstore"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/mail/protocol"
	"github.com/tidbcloud/AgentTeam/internal/extension/providers/mail/token"
	"github.com/tidbcloud/AgentTeam/internal/store/extension"
	"github.com/tidbcloud/AgentTeam/internal/store/storeerr"
)

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

	fixture := newRuntimeFixture(t)
	for index := 0; index < 50; index++ {
		fixture.imap.searchPage.Messages = append(
			fixture.imap.searchPage.Messages,
			protocol.MessageSummary{
				Ref: protocol.MessageRef{
					Mailbox:     "INBOX",
					UIDValidity: 41,
					UID:         uint32(1000 - index),
				},
				From:      "sender@example.com",
				To:        []string{"recipient@example.com"},
				Subject:   strings.Repeat("s", 998),
				SentAt:    fixture.now.Add(-time.Duration(index) * time.Minute),
				Truncated: false,
			},
		)
	}
	result, err := fixture.service.Invoke(
		context.Background(),
		fixture.invocation("search-messages", map[string]any{
			"filters": map[string]any{"folder": "inbox"},
			"limit":   50,
		}),
	)
	require.NoError(t, err)
	require.LessOrEqual(t, len(result), maxMailResultBytes)
	var output searchOutput
	require.NoError(t, json.Unmarshal(result, &output))
	require.NotEmpty(t, output.Messages)
	require.Less(t, len(output.Messages), 50)
	require.True(t, output.Truncated)
	require.True(t, output.HasMore)
	require.NotEmpty(t, output.NextCursor)
	require.NotEmpty(t, output.Messages[0].AttachmentMessage)
	require.NotEmpty(t, output.Messages[0].ReplyMessage)
	require.NotEmpty(t, output.Messages[0].ForwardMessage)
	require.Equal(t, 1, fixture.imap.searchCalls)

	claims, err := fixture.tokens.Open(
		output.Messages[0].Message,
		token.Binding{
			Kind:           token.KindMessage,
			WorkspaceID:    "wsp_1",
			InstallationID: "appinst_1",
			ConnectionID:   "appconn_1",
			Generation:     fixture.connection.Generation,
			AgentID:        "agt_1",
			Action:         "get-message",
		},
	)
	require.NoError(t, err)
	require.Equal(t, uint32(41), claims.UIDValidity)
	require.NotContains(t, output.Messages[0].Message, "INBOX")
}

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

	fixture := newRuntimeFixture(t)
	ref := protocol.MessageRef{
		Mailbox: "INBOX", UIDValidity: 42, UID: 7,
	}
	fixture.imap.message = protocol.Message{
		Ref:     ref,
		From:    "sender@example.com",
		To:      []string{"recipient@example.com"},
		CC:      []string{"copy@example.com"},
		Subject: "Untrusted",
		SentAt:  fixture.now,
		Body: protocol.Body{
			Format:  "sanitized_html",
			Content: "<pre>safe &lt;script&gt;alert(1)&lt;/script&gt;</pre>",
		},
		HasAttachments: true,
	}
	messageToken := fixture.sealReference(
		t,
		token.KindMessage,
		"get-message",
		ref,
		0,
		"",
	)
	result, err := fixture.service.Invoke(
		context.Background(),
		fixture.invocation("get-message", map[string]any{
			"message": messageToken,
		}),
	)
	require.NoError(t, err)
	var message messageOutput
	require.NoError(t, json.Unmarshal(result, &message))
	require.True(t, message.UntrustedExternalContent)
	require.True(t, message.Body.UntrustedExternalContent)
	require.Equal(t, "sanitized_html", message.Body.Format)
	require.NotContains(t, message.Body.Content, "<script>")
	require.NotEmpty(t, message.AttachmentMessage)
	require.NotEmpty(t, message.ReplyMessage)
	require.NotEmpty(t, message.ForwardMessage)

	fixture.imap.attachments = []protocol.Attachment{{
		Part:      []int{2, 1},
		Filename:  "invoice.pdf",
		MediaType: "application/pdf",
		SizeBytes: 1234,
	}}
	result, err = fixture.service.Invoke(
		context.Background(),
		fixture.invocation("list-attachments", map[string]any{
			"message": message.AttachmentMessage,
		}),
	)
	require.NoError(t, err)
	var attachments attachmentListOutput
	require.NoError(t, json.Unmarshal(result, &attachments))
	require.Len(t, attachments.Attachments, 1)
	require.Equal(t, "invoice.pdf", attachments.Attachments[0].Filename)
	require.True(t, attachments.Attachments[0].UntrustedExternalContent)
	require.NotEmpty(t, attachments.Attachments[0].Attachment)
	require.NotContains(t, string(result), "content")
	require.Equal(t, 1, fixture.imap.attachmentCalls)
}

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

	fixture := newRuntimeFixture(t)
	oldRef := protocol.DraftRef{
		Mailbox: "Drafts", UIDValidity: 52, UID: 9, ModSeq: 77,
	}
	fixture.imap.replaceRef = protocol.DraftRef{
		Mailbox: "Drafts", UIDValidity: 52, UID: 10, ModSeq: 88,
	}
	draftToken := fixture.sealReference(
		t,
		token.KindDraft,
		"update-draft",
		protocol.MessageRef{
			Mailbox: oldRef.Mailbox, UIDValidity: oldRef.UIDValidity,
			UID: oldRef.UID,
		},
		oldRef.ModSeq,
		"",
	)
	result, err := fixture.service.Invoke(
		context.Background(),
		fixture.invocation("update-draft", map[string]any{
			"clientRequestId":  "req_update_1",
			"draft":            draftToken,
			"expectedRevision": 77,
			"to":               []string{"recipient@example.com"},
			"subject":          "Updated",
			"body": map[string]any{
				"format": "text", "content": "body",
			},
		}),
	)
	require.NoError(t, err)
	var output draftMutationOutput
	require.NoError(t, json.Unmarshal(result, &output))
	require.Equal(t, "accepted", output.Status)
	require.Equal(t, uint64(88), output.Revision)
	require.NotEmpty(t, output.Draft)
	require.NotEmpty(t, output.SendDraft)
	require.Equal(t, 1, fixture.imap.replaceCalls)
	require.Equal(t, oldRef, fixture.imap.replacedOldRef)
	require.Equal(
		t,
		"CaseSensitive@example.com",
		fixture.imap.replacedContent.From,
	)
}

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

	fixture := newRuntimeFixture(t)
	sourceRef := protocol.MessageRef{
		Mailbox: "INBOX", UIDValidity: 42, UID: 8,
	}
	fixture.imap.message = protocol.Message{
		Ref:     sourceRef,
		From:    "author@example.com",
		To:      []string{"CaseSensitive@example.com", "team@example.com"},
		CC:      []string{"copy@example.com"},
		Subject: "Status",
		SentAt:  fixture.now,
		Body: protocol.Body{
			Format: "text", Content: "source line",
		},
	}
	fixture.imap.appendRefs = []protocol.DraftRef{
		{Mailbox: "Drafts", UIDValidity: 52, UID: 10, ModSeq: 80},
		{Mailbox: "Drafts", UIDValidity: 52, UID: 11, ModSeq: 81},
	}

	replyToken := fixture.sealReference(
		t,
		token.KindMessage,
		"create-reply-draft",
		sourceRef,
		0,
		"",
	)
	_, err := fixture.service.Invoke(
		context.Background(),
		fixture.invocation("create-reply-draft", map[string]any{
			"clientRequestId": "req_reply_1",
			"message":         replyToken,
			"replyAll":        true,
			"body": map[string]any{
				"format": "text", "content": "my reply",
			},
		}),
	)
	require.NoError(t, err)

	forwardToken := fixture.sealReference(
		t,
		token.KindMessage,
		"create-forward-draft",
		sourceRef,
		0,
		"",
	)
	_, err = fixture.service.Invoke(
		context.Background(),
		fixture.invocation("create-forward-draft", map[string]any{
			"clientRequestId": "req_forward_1",
			"message":         forwardToken,
			"to":              []string{"next@example.com"},
			"body": map[string]any{
				"format": "html", "content": "<p>FYI</p>",
			},
		}),
	)
	require.NoError(t, err)
	require.Equal(t, 2, fixture.imap.messageCalls)
	require.Equal(t, 2, fixture.imap.appendCalls)
	require.Len(t, fixture.imap.appended, 2)
	require.Equal(t, "Re: Status", fixture.imap.appended[0].Subject)
	require.Contains(t, fixture.imap.appended[0].Body, "> source line")
	require.Contains(t, fixture.imap.appended[0].To, "author@example.com")
	require.Contains(t, fixture.imap.appended[0].CC, "team@example.com")
	require.Equal(t, "Fwd: Status", fixture.imap.appended[1].Subject)
	require.Contains(t, fixture.imap.appended[1].Body, "Forwarded message")
	require.NotContains(t, fixture.imap.appended[1].Body, "<script>")
	for _, content := range fixture.imap.appended {
		require.Equal(t, "CaseSensitive@example.com", content.From)
	}
}

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

	for _, action := range []string{
		actionCreateReplyDraft,
		actionCreateForwardDraft,
	} {
		t.Run(action, func(t *testing.T) {
			fixture := newRuntimeFixture(t)
			sourceRef := protocol.MessageRef{
				Mailbox: "INBOX", UIDValidity: 42, UID: 8,
			}
			fixture.imap.message = protocol.Message{
				Ref:       sourceRef,
				From:      "author@example.com",
				To:        []string{"CaseSensitive@example.com"},
				Subject:   "Status",
				SentAt:    fixture.now,
				Truncated: true,
				Body: protocol.Body{
					Format: "text", Content: "partial source",
					Truncated: true,
				},
			}
			token := fixture.sealReference(
				t,
				token.KindMessage,
				action,
				sourceRef,
				0,
				"",
			)
			input := map[string]any{
				"clientRequestId": "req_quote_1",
				"message":         token,
				"body": map[string]any{
					"format": "text", "content": "intro",
				},
			}
			if action == actionCreateForwardDraft {
				input["to"] = []string{"next@example.com"}
			}

			_, err := fixture.service.Invoke(
				context.Background(),
				fixture.invocation(action, input),
			)
			require.Error(t, err)
			require.Equal(t, ErrorConflict, mustErrorCode(t, err))
			require.Equal(t, 1, fixture.imap.messageCalls)
			require.Equal(t, 0, fixture.imap.appendCalls)
		})
	}
}

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

	fixture := newRuntimeFixture(t)
	fixture.imap.appendErr = &protocol.ProtocolError{
		Class:     protocol.ErrorUnknown,
		Operation: "append draft",
	}
	result, err := fixture.service.Invoke(
		context.Background(),
		fixture.invocation("create-draft", map[string]any{
			"clientRequestId": "req_create_1",
			"to":              []string{"recipient@example.com"},
			"subject":         "Subject",
			"body": map[string]any{
				"format": "text", "content": "body",
			},
		}),
	)
	require.NoError(t, err)
	var output draftMutationOutput
	require.NoError(t, json.Unmarshal(result, &output))
	require.Equal(t, "unknown", output.Status)
	require.Equal(t, 1, fixture.imap.appendCalls)
	require.Equal(t, 0, fixture.imap.getDraftCalls)
	require.Empty(t, output.Draft)
}

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

	fixture := newRuntimeFixture(t)
	fixture.imap.appendRefs = []protocol.DraftRef{{
		Mailbox: "Drafts", UIDValidity: 52, UID: 21,
	}}
	result, err := fixture.service.Invoke(
		context.Background(),
		fixture.invocation("create-draft", map[string]any{
			"clientRequestId": "req_create_1",
			"to":              []string{"recipient@example.com"},
			"subject":         "Subject",
			"body": map[string]any{
				"format": "text", "content": "body",
			},
		}),
	)
	require.NoError(t, err)
	var output draftMutationOutput
	require.NoError(t, json.Unmarshal(result, &output))
	require.Equal(t, "accepted", output.Status)
	require.Zero(t, output.Revision)
	require.NotEmpty(t, output.Draft)
	require.Empty(t, output.SendDraft)
	require.Equal(t, 1, fixture.imap.appendCalls)
}

type runtimeFixture struct {
	now          time.Time
	key          []byte
	connection   mailstore.Connection
	connections  *sequenceConnectionLoader
	imap         *recordingIMAP
	smtp         *recordingSMTP
	attachments  *recordingAttachmentLoader
	materializer *recordingMaterializer
	authorizer   *recordingAuthorizer
	tokens       *token.Codec
	service      *Service
}

type recordingAuthorizer struct {
	mu      sync.Mutex
	calls   int
	err     error
	last    InvocationAuthorization
	afterOK bool
}

func (authorizer *recordingAuthorizer) ReauthorizeInvocation(
	_ context.Context,
	authorization InvocationAuthorization,
) error {
	authorizer.mu.Lock()
	defer authorizer.mu.Unlock()
	authorizer.calls++
	authorizer.last = authorization
	if authorizer.afterOK && authorizer.calls >= 2 {
		return storeerr.ErrForbidden
	}
	return authorizer.err
}

func newRuntimeFixture(t *testing.T) *runtimeFixture {
	t.Helper()
	now := time.Date(2026, 8, 15, 12, 0, 0, 0, time.UTC)
	key := []byte(strings.Repeat("k", 32))
	connection := statelessTestConnection(t, key)
	connections := &sequenceConnectionLoader{
		sequence: []mailstore.Connection{connection},
	}
	tokens, err := token.New(
		token.Keyring{Current: token.Key{
			ID: "runtime", Material: key,
		}},
		"test",
		func() time.Time { return now },
	)
	require.NoError(t, err)
	imap := &recordingIMAP{}
	smtp := &recordingSMTP{}
	attachments := &recordingAttachmentLoader{}
	materializer := &recordingMaterializer{}
	authorizer := &recordingAuthorizer{}
	service := New(nil, key, Options{
		Connections:  connections,
		IMAP:         imap,
		SMTP:         smtp,
		Attachments:  attachments,
		Materializer: materializer,
		Authorizer:   authorizer,
		Tokens:       tokens,
		Now:          func() time.Time { return now },
		TokenTTL:     10 * time.Minute,
	})
	return &runtimeFixture{
		now: now, key: key, connection: connection,
		connections: connections, imap: imap, smtp: smtp,
		attachments: attachments, materializer: materializer,
		authorizer: authorizer,
		tokens:     tokens, service: service,
	}
}

func (fixture *runtimeFixture) invocation(
	action string,
	input any,
) Invocation {
	encoded, err := json.Marshal(input)
	if err != nil {
		panic(err)
	}
	return Invocation{
		WorkspaceID:           "wsp_1",
		AgentID:               "agt_1",
		AgentMemberID:         "mem_agent_1",
		ServiceInstallationID: "appinst_1",
		ServiceConnectionID:   "appconn_1",
		ServiceConnectionGeneration: extension.RuntimeConnectionGeneration(
			fixture.connection.ID,
			fixture.connection.ExternalSubjectRef,
			mailstore.CredentialEnvelopeDigest(
				fixture.connection.EncryptedCredentials,
			),
			fixture.connection.Config,
		),
		ServiceInstallationExtensionID: "instext_1",
		AuthorizedByMemberID:           "mem_human_1",
		ExtensionName:                  "work-mail",
		Action:                         action,
		DefinitionVersion:              1,
		DefinitionDigest:               "sha256:" + strings.Repeat("d", 64),
		CatalogVersion:                 3,
		RequestHash:                    "sha256:" + strings.Repeat("e", 64),
		Input:                          encoded,
	}
}

func (fixture *runtimeFixture) sealReference(
	t *testing.T,
	kind string,
	action string,
	ref protocol.MessageRef,
	revision uint64,
	part string,
) string {
	t.Helper()
	token, err := fixture.tokens.Seal(
		token.Claims{
			Kind:           kind,
			WorkspaceID:    "wsp_1",
			InstallationID: "appinst_1",
			ConnectionID:   "appconn_1",
			Generation:     fixture.connection.Generation,
			AgentID:        "agt_1",
			Action:         action,
			Mailbox:        ref.Mailbox,
			UIDValidity:    ref.UIDValidity,
			UID:            ref.UID,
			Revision:       revision,
			Part:           part,
		},
		10*time.Minute,
	)
	require.NoError(t, err)
	return token
}

type sequenceConnectionLoader struct {
	mu       sync.Mutex
	sequence []mailstore.Connection
	calls    int
}

func (loader *sequenceConnectionLoader) LoadActiveConnection(
	_ context.Context,
	workspaceID string,
	installationID string,
	connectionID string,
) (mailstore.Connection, error) {
	loader.mu.Lock()
	defer loader.mu.Unlock()
	loader.calls++
	if len(loader.sequence) == 0 {
		return mailstore.Connection{}, errors.New("no connection")
	}
	index := loader.calls - 1
	if index >= len(loader.sequence) {
		index = len(loader.sequence) - 1
	}
	connection := loader.sequence[index]
	if connection.WorkspaceID != workspaceID ||
		connection.ServiceInstallationID != installationID ||
		connection.ID != connectionID {
		return mailstore.Connection{}, fmt.Errorf("wrong connection identity")
	}
	return connection, nil
}

type recordingIMAP struct {
	mu sync.Mutex

	searchPage      protocol.SearchPage
	searchErr       error
	searchCalls     int
	message         protocol.Message
	messageErr      error
	messageCalls    int
	attachments     []protocol.Attachment
	attachmentErr   error
	attachmentCalls int
	attachmentBlobs map[string][]byte
	appendRefs      []protocol.DraftRef
	appendErr       error
	appendCalls     int
	appended        []protocol.DraftContent
	draft           protocol.Draft
	getDraftErr     error
	getDraftCalls   int
	replaceRef      protocol.DraftRef
	replaceErr      error
	replaceCalls    int
	replacedOldRef  protocol.DraftRef
	replacedContent protocol.DraftContent
}

func (*recordingIMAP) Probe(
	context.Context,
	protocol.Endpoint,
) (protocol.Capabilities, error) {
	return protocol.Capabilities{}, nil
}

func (imap *recordingIMAP) Search(
	_ context.Context,
	_ protocol.Endpoint,
	_ protocol.SearchRequest,
) (protocol.SearchPage, error) {
	imap.mu.Lock()
	defer imap.mu.Unlock()
	imap.searchCalls++
	return imap.searchPage, imap.searchErr
}

func (imap *recordingIMAP) GetMessage(
	_ context.Context,
	_ protocol.Endpoint,
	_ protocol.MessageRef,
) (protocol.Message, error) {
	imap.mu.Lock()
	defer imap.mu.Unlock()
	imap.messageCalls++
	return imap.message, imap.messageErr
}

func (imap *recordingIMAP) ListAttachments(
	_ context.Context,
	_ protocol.Endpoint,
	_ protocol.MessageRef,
) ([]protocol.Attachment, error) {
	imap.mu.Lock()
	defer imap.mu.Unlock()
	imap.attachmentCalls++
	return append([]protocol.Attachment(nil), imap.attachments...),
		imap.attachmentErr
}

func (imap *recordingIMAP) AppendDraft(
	_ context.Context,
	_ protocol.Endpoint,
	content protocol.DraftContent,
) (protocol.DraftRef, error) {
	imap.mu.Lock()
	defer imap.mu.Unlock()
	imap.appendCalls++
	imap.appended = append(imap.appended, content)
	if imap.appendErr != nil {
		return protocol.DraftRef{}, imap.appendErr
	}
	index := imap.appendCalls - 1
	if index < len(imap.appendRefs) {
		return imap.appendRefs[index], nil
	}
	return protocol.DraftRef{
		Mailbox: "Drafts", UIDValidity: 52,
		UID: uint32(20 + imap.appendCalls), ModSeq: uint64(90 + imap.appendCalls),
	}, nil
}

func (imap *recordingIMAP) GetAttachment(
	_ context.Context,
	_ protocol.Endpoint,
	_ protocol.MessageRef,
	part []int,
) (protocol.AttachmentContent, error) {
	imap.mu.Lock()
	defer imap.mu.Unlock()
	key := formatPartTest(part)
	for _, attachment := range imap.attachments {
		if formatPartTest(attachment.Part) == key {
			return protocol.AttachmentContent{
				Part:      append([]int(nil), part...),
				Filename:  attachment.Filename,
				MediaType: attachment.MediaType,
				Blob:      imap.attachmentBlobs[key],
			}, nil
		}
	}
	for _, attachment := range imap.draft.Attachments {
		if formatPartTest(attachment.Part) == key {
			return protocol.AttachmentContent{
				Part:      append([]int(nil), part...),
				Filename:  attachment.Filename,
				MediaType: attachment.MediaType,
				Blob:      imap.attachmentBlobs[key],
			}, nil
		}
	}
	return protocol.AttachmentContent{}, imap.attachmentErr
}

func formatPartTest(part []int) string {
	values := make([]string, 0, len(part))
	for _, value := range part {
		values = append(values, fmt.Sprint(value))
	}
	return strings.Join(values, ".")
}

func (imap *recordingIMAP) GetDraft(
	_ context.Context,
	_ protocol.Endpoint,
	_ protocol.DraftRef,
) (protocol.Draft, error) {
	imap.mu.Lock()
	defer imap.mu.Unlock()
	imap.getDraftCalls++
	return imap.draft, imap.getDraftErr
}

func (imap *recordingIMAP) ReplaceDraft(
	_ context.Context,
	_ protocol.Endpoint,
	ref protocol.DraftRef,
	content protocol.DraftContent,
) (protocol.DraftRef, error) {
	imap.mu.Lock()
	defer imap.mu.Unlock()
	imap.replaceCalls++
	imap.replacedOldRef = ref
	imap.replacedContent = content
	return imap.replaceRef, imap.replaceErr
}

type recordingSMTP struct {
	mu     sync.Mutex
	calls  int
	frozen protocol.FrozenMessage
	status string
	err    error
}

func (smtp *recordingSMTP) Send(
	_ context.Context,
	_ protocol.Endpoint,
	frozen protocol.FrozenMessage,
) (protocol.SendResult, error) {
	smtp.mu.Lock()
	defer smtp.mu.Unlock()
	smtp.calls++
	smtp.frozen = frozen
	if smtp.status == "" {
		smtp.status = protocol.ResultAccepted
	}
	return protocol.SendResult{Status: smtp.status}, smtp.err
}

func mustErrorCode(t *testing.T, err error) ErrorCode {
	t.Helper()
	code, ok := ErrorCodeOf(err)
	require.True(t, ok)
	return code
}

func connectionWithoutCapability(
	t *testing.T,
	connection mailstore.Connection,
	removed string,
) mailstore.Connection {
	t.Helper()
	var config storedConnectionConfig
	require.NoError(t, json.Unmarshal(connection.Config, &config))
	filtered := config.IMAP.Capabilities[:0]
	for _, capability := range config.IMAP.Capabilities {
		if capability != removed {
			filtered = append(filtered, capability)
		}
	}
	config.IMAP.Capabilities = filtered
	digest, err := config.ComputeDigest()
	require.NoError(t, err)
	config.CapabilityDigest = digest
	connection.Config, err = json.Marshal(config)
	require.NoError(t, err)
	connection.Generation = "sha256:" + strings.Repeat("f", 64)
	return connection
}