package mailruntime

import (
	"context"
	"encoding/json"
	"errors"
	stdhtml "html"
	"strconv"
	"strings"

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

type searchInput struct {
	Filters SearchFilters `json:"filters"`
	Limit   int           `json:"limit,omitempty"`
	Cursor  string        `json:"cursor,omitempty"`
}

type resourceInput struct {
	Message string `json:"message"`
}

type messageMetadataOutput struct {
	Message                  string   `json:"message"`
	AttachmentMessage        string   `json:"attachmentMessage"`
	ReplyMessage             string   `json:"replyMessage"`
	ForwardMessage           string   `json:"forwardMessage"`
	From                     string   `json:"from"`
	To                       []string `json:"to"`
	Subject                  string   `json:"subject"`
	SentAt                   string   `json:"sentAt"`
	SizeBytes                uint32   `json:"sizeBytes"`
	Unread                   bool     `json:"unread"`
	HasAttachments           bool     `json:"hasAttachments"`
	Truncated                bool     `json:"truncated"`
	UntrustedExternalContent bool     `json:"untrustedExternalContent"`
}

type searchOutput struct {
	Messages   []messageMetadataOutput `json:"messages"`
	NextCursor string                  `json:"nextCursor,omitempty"`
	HasMore    bool                    `json:"hasMore"`
	Truncated  bool                    `json:"truncated"`
}

type projectedSearchMessage struct {
	output messageMetadataOutput
	uid    uint32
}

type messageOutput struct {
	Message                  string     `json:"message"`
	AttachmentMessage        string     `json:"attachmentMessage,omitempty"`
	ReplyMessage             string     `json:"replyMessage,omitempty"`
	ForwardMessage           string     `json:"forwardMessage,omitempty"`
	From                     string     `json:"from"`
	To                       []string   `json:"to"`
	CC                       []string   `json:"cc,omitempty"`
	Subject                  string     `json:"subject"`
	SentAt                   string     `json:"sentAt"`
	Body                     bodyOutput `json:"body"`
	HasAttachments           bool       `json:"hasAttachments"`
	Truncated                bool       `json:"truncated"`
	UntrustedExternalContent bool       `json:"untrustedExternalContent"`
}

type bodyOutput struct {
	Format                   string `json:"format"`
	Content                  string `json:"content"`
	Truncated                bool   `json:"truncated"`
	UntrustedExternalContent bool   `json:"untrustedExternalContent"`
}

type attachmentOutput struct {
	Attachment               string `json:"attachment"`
	Filename                 string `json:"filename"`
	MediaType                string `json:"mediaType"`
	SizeBytes                uint32 `json:"sizeBytes"`
	UntrustedExternalContent bool   `json:"untrustedExternalContent"`
}

type attachmentListOutput struct {
	Attachments []attachmentOutput `json:"attachments"`
	Truncated   bool               `json:"truncated"`
}

func (service *Service) searchMessages(
	ctx context.Context,
	invocation Invocation,
) (json.RawMessage, error) {
	var input searchInput
	if err := decodeStrictJSON(invocation.Input, &input); err != nil {
		return nil, mailError(
			ErrorValidation,
			"mail search input is invalid",
			err,
		)
	}
	if input.Limit == 0 {
		input.Limit = defaultSearchLimit
	}
	if input.Limit < 1 ||
		input.Limit > maxSearchResults ||
		len(input.Cursor) > 2048 ||
		validateSearchFilters(input.Filters) != nil {
		return nil, mailError(
			ErrorValidation,
			"mail search input is invalid",
			nil,
		)
	}
	runtime, err := service.resolveRuntime(ctx, invocation)
	if err != nil {
		return nil, err
	}
	filterDigest, err := sha256JSON(input.Filters)
	if err != nil {
		return nil, mailError(
			ErrorInternal,
			"mail search filter could not be bound",
			err,
		)
	}
	var beforeUID uint32
	if input.Cursor != "" {
		beforeUID, err = service.openSearchCursor(
			input.Cursor,
			runtime,
			filterDigest,
			input.Limit,
		)
		if err != nil {
			return nil, err
		}
	}
	runtime, err = service.recheckRuntime(ctx, runtime)
	if err != nil {
		return nil, err
	}
	folder := input.Filters.Folder
	if folder == "spam" {
		folder = protocol.FolderJunk
	}
	page, err := service.imap.Search(
		ctx,
		runtime.material.IMAP,
		protocol.SearchRequest{
			Folder:         folder,
			Text:           input.Filters.Text,
			From:           append([]string(nil), input.Filters.From...),
			To:             append([]string(nil), input.Filters.To...),
			ReceivedAfter:  input.Filters.ReceivedAfter,
			ReceivedBefore: input.Filters.ReceivedBefore,
			Unread:         input.Filters.Unread,
			HasAttachments: input.Filters.HasAttachments,
			Limit:          input.Limit,
			BeforeUID:      beforeUID,
		},
	)
	if err != nil {
		return nil, classifyProtocolError(err)
	}
	if len(page.Messages) > input.Limit ||
		len(page.Messages) > maxSearchResults {
		return nil, mailError(
			ErrorUpstream,
			"mailbox returned too many search results",
			nil,
		)
	}
	if _, err := service.recheckRuntime(ctx, runtime); err != nil {
		return nil, err
	}
	projected, outputTruncated, err := service.projectSearchMessages(
		runtime,
		page.Messages,
	)
	if err != nil {
		return nil, err
	}
	return service.marshalSearchOutput(
		runtime,
		page,
		projected,
		outputTruncated,
		filterDigest,
		input.Limit,
	)
}

func (service *Service) projectSearchMessages(
	runtime runtimeContext,
	messages []protocol.MessageSummary,
) ([]projectedSearchMessage, bool, error) {
	projected := make(
		[]projectedSearchMessage,
		0,
		len(messages),
	)
	outputTruncated := false
	for _, message := range messages {
		if err := validateMessageSummary(message); err != nil {
			return nil, false, mailError(
				ErrorUpstream,
				"mailbox search result is invalid",
				err,
			)
		}
		messageToken, err := service.sealReference(
			runtime,
			token.KindMessage,
			actionGetMessage,
			message.Ref,
			0,
			"",
		)
		if err != nil {
			return nil, false, err
		}
		attachmentMessage, err := service.sealReference(
			runtime,
			token.KindMessage,
			actionListAttachments,
			message.Ref,
			0,
			"",
		)
		if err != nil {
			return nil, false, err
		}
		replyMessage, err := service.sealReference(
			runtime,
			token.KindMessage,
			actionCreateReplyDraft,
			message.Ref,
			0,
			"",
		)
		if err != nil {
			return nil, false, err
		}
		forwardMessage, err := service.sealReference(
			runtime,
			token.KindMessage,
			actionCreateForwardDraft,
			message.Ref,
			0,
			"",
		)
		if err != nil {
			return nil, false, err
		}
		outputTruncated = outputTruncated || message.Truncated
		projected = append(projected, projectedSearchMessage{
			uid: message.Ref.UID,
			output: messageMetadataOutput{
				Message:                  messageToken,
				AttachmentMessage:        attachmentMessage,
				ReplyMessage:             replyMessage,
				ForwardMessage:           forwardMessage,
				From:                     message.From,
				To:                       append([]string{}, message.To...),
				Subject:                  message.Subject,
				SentAt:                   message.SentAt.UTC().Format(timeFormat),
				SizeBytes:                message.SizeBytes,
				Unread:                   message.Unread,
				HasAttachments:           message.HasAttachments,
				Truncated:                message.Truncated,
				UntrustedExternalContent: true,
			},
		})
	}
	return projected, outputTruncated, nil
}

func (service *Service) marshalSearchOutput(
	runtime runtimeContext,
	page protocol.SearchPage,
	projected []projectedSearchMessage,
	outputTruncated bool,
	filterDigest string,
	limit int,
) (json.RawMessage, error) {
	for {
		nextBeforeUID := page.BeforeUID
		if len(projected) < len(page.Messages) && len(projected) > 0 {
			nextBeforeUID = projected[len(projected)-1].uid
		}
		nextCursor, err := service.sealSearchCursor(
			runtime,
			nextBeforeUID,
			filterDigest,
			limit,
		)
		if err != nil {
			return nil, err
		}
		messages := make(
			[]messageMetadataOutput,
			0,
			len(projected),
		)
		for _, message := range projected {
			messages = append(messages, message.output)
		}
		output := searchOutput{
			Messages:   messages,
			NextCursor: nextCursor,
			HasMore:    nextCursor != "",
			Truncated: outputTruncated ||
				len(projected) < len(page.Messages),
		}
		encoded, err := marshalBounded(output)
		if !errors.Is(err, errMailResultTooLarge) {
			return encoded, err
		}
		if len(projected) <= 1 {
			return nil, mailError(
				ErrorUpstream,
				"mail search result could not fit the response budget",
				err,
			)
		}
		projected = projected[:len(projected)-1]
	}
}

func (service *Service) getMessage(
	ctx context.Context,
	invocation Invocation,
) (json.RawMessage, error) {
	var input resourceInput
	if err := decodeStrictJSON(invocation.Input, &input); err != nil ||
		input.Message == "" {
		return nil, mailError(
			ErrorValidation,
			"mail message input is invalid",
			err,
		)
	}
	runtime, err := service.resolveRuntime(ctx, invocation)
	if err != nil {
		return nil, err
	}
	ref, _, err := service.openReference(
		input.Message,
		runtime,
		token.KindMessage,
		actionGetMessage,
	)
	if err != nil {
		return nil, err
	}
	runtime, err = service.recheckRuntime(ctx, runtime)
	if err != nil {
		return nil, err
	}
	message, err := service.imap.GetMessage(
		ctx,
		runtime.material.IMAP,
		ref,
	)
	if err != nil {
		return nil, classifyProtocolError(err)
	}
	if message.Ref != ref {
		return nil, mailError(
			ErrorUpstream,
			"mailbox returned the wrong message",
			nil,
		)
	}
	if err := validateMessage(message); err != nil {
		return nil, mailError(
			ErrorUpstream,
			"mailbox message is invalid",
			err,
		)
	}
	if _, err := service.recheckRuntime(ctx, runtime); err != nil {
		return nil, err
	}
	body, err := normalizeReadBody(message.Body)
	if err != nil {
		return nil, mailError(
			ErrorUpstream,
			"mailbox message body is invalid",
			err,
		)
	}
	output, err := service.projectMessageOutput(
		runtime,
		input.Message,
		ref,
		message,
		body,
	)
	if err != nil {
		return nil, err
	}
	return marshalMessageOutput(output, body)
}

func (service *Service) projectMessageOutput(
	runtime runtimeContext,
	messageToken string,
	ref protocol.MessageRef,
	message protocol.Message,
	body protocol.Body,
) (messageOutput, error) {
	attachmentMessage, err := service.sealReference(
		runtime,
		token.KindMessage,
		actionListAttachments,
		ref,
		0,
		"",
	)
	if err != nil {
		return messageOutput{}, err
	}
	replyMessage, err := service.sealReference(
		runtime,
		token.KindMessage,
		actionCreateReplyDraft,
		ref,
		0,
		"",
	)
	if err != nil {
		return messageOutput{}, err
	}
	forwardMessage, err := service.sealReference(
		runtime,
		token.KindMessage,
		actionCreateForwardDraft,
		ref,
		0,
		"",
	)
	if err != nil {
		return messageOutput{}, err
	}
	return messageOutput{
		Message:           messageToken,
		AttachmentMessage: attachmentMessage,
		ReplyMessage:      replyMessage,
		ForwardMessage:    forwardMessage,
		From:              message.From,
		To:                append([]string{}, message.To...),
		CC:                append([]string(nil), message.CC...),
		Subject:           message.Subject,
		SentAt:            message.SentAt.UTC().Format(timeFormat),
		Body: bodyOutput{
			Format:                   body.Format,
			Content:                  body.Content,
			Truncated:                body.Truncated,
			UntrustedExternalContent: true,
		},
		HasAttachments:           message.HasAttachments,
		Truncated:                message.Truncated || body.Truncated,
		UntrustedExternalContent: true,
	}, nil
}

func marshalMessageOutput(
	output messageOutput,
	body protocol.Body,
) (json.RawMessage, error) {
	if encoded, err := marshalBounded(output); err == nil {
		return encoded, nil
	}
	output.Truncated = true
	output.Body.Truncated = true
	plain := body.Content
	if body.Format == "sanitized_html" {
		converted, err := htmlToPlainText(body.Content)
		if err != nil {
			return nil, mailError(
				ErrorUpstream,
				"mailbox message body is invalid",
				err,
			)
		}
		plain = converted
	}
	for limit := len(plain) / 2; limit >= 0; limit /= 2 {
		content := boundedUTF8(plain, limit)
		if body.Format == "sanitized_html" {
			content = "<pre>" + stdhtml.EscapeString(content) + "</pre>"
		}
		output.Body.Content = content
		encoded, fitErr := marshalBounded(output)
		if fitErr == nil {
			return encoded, nil
		}
		if limit == 0 {
			break
		}
	}
	return nil, mailError(
		ErrorUpstream,
		"mail message result could not fit the response budget",
		errMailResultTooLarge,
	)
}

func (service *Service) listAttachments(
	ctx context.Context,
	invocation Invocation,
) (json.RawMessage, error) {
	var input resourceInput
	if err := decodeStrictJSON(invocation.Input, &input); err != nil ||
		input.Message == "" {
		return nil, mailError(
			ErrorValidation,
			"mail attachment input is invalid",
			err,
		)
	}
	runtime, err := service.resolveRuntime(ctx, invocation)
	if err != nil {
		return nil, err
	}
	ref, _, err := service.openReference(
		input.Message,
		runtime,
		token.KindMessage,
		actionListAttachments,
	)
	if err != nil {
		return nil, err
	}
	runtime, err = service.recheckRuntime(ctx, runtime)
	if err != nil {
		return nil, err
	}
	attachments, err := service.imap.ListAttachments(
		ctx,
		runtime.material.IMAP,
		ref,
	)
	if err != nil {
		return nil, classifyProtocolError(err)
	}
	if len(attachments) > maxAttachmentMetadata {
		return nil, mailError(
			ErrorUpstream,
			"mailbox returned too many attachments",
			nil,
		)
	}
	if _, err := service.recheckRuntime(ctx, runtime); err != nil {
		return nil, err
	}
	projected := make([]attachmentOutput, 0, len(attachments))
	metadataTruncated := false
	for _, attachment := range attachments {
		if len(attachment.Part) == 0 ||
			attachment.Filename == "" ||
			attachment.MediaType == "" {
			return nil, mailError(
				ErrorUpstream,
				"mailbox attachment metadata is invalid",
				nil,
			)
		}
		part := formatPart(attachment.Part)
		token, err := service.sealReference(
			runtime,
			token.KindAttachment,
			actionGetAttachment,
			ref,
			0,
			part,
		)
		if err != nil {
			return nil, err
		}
		projected = append(projected, attachmentOutput{
			Attachment:               token,
			Filename:                 attachment.Filename,
			MediaType:                attachment.MediaType,
			SizeBytes:                attachment.SizeBytes,
			UntrustedExternalContent: true,
		})
		metadataTruncated = metadataTruncated || attachment.Truncated
	}
	for {
		output := attachmentListOutput{
			Attachments: projected,
			Truncated: metadataTruncated ||
				len(projected) < len(attachments),
		}
		encoded, err := marshalBounded(output)
		if !errors.Is(err, errMailResultTooLarge) {
			return encoded, err
		}
		if len(projected) <= 1 {
			return nil, mailError(
				ErrorUpstream,
				"mail attachment result could not fit the response budget",
				err,
			)
		}
		projected = projected[:len(projected)-1]
	}
}

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