package imap

import (
	"bytes"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"mime/quotedprintable"
	"net/mail"
	"paraclub-ai-mailer/config"
	"paraclub-ai-mailer/internal/logger"
	"strings"
	"time"

	"github.com/emersion/go-imap"
	"github.com/emersion/go-imap/client"
	"github.com/sirupsen/logrus"
)

// Custom IMAP flag to mark emails as AI-processed
const AIProcessedFlag = "$AIProcessed"

type IMAPClient struct {
	client *client.Client
	config config.IMAPConfig
}

// MessageLiteral implements imap.Literal for draft messages
type MessageLiteral struct {
	content []byte
	pos     int
}

func (m *MessageLiteral) Read(p []byte) (n int, err error) {
	if m.pos >= len(m.content) {
		return 0, io.EOF
	}
	n = copy(p, m.content[m.pos:])
	m.pos += n
	if m.pos >= len(m.content) {
		err = io.EOF
	}
	return
}

func (m *MessageLiteral) Len() int {
	return len(m.content)
}

func New(cfg config.IMAPConfig) (*IMAPClient, error) {
	addr := fmt.Sprintf("%s:%d", cfg.Server, cfg.Port)
	var c *client.Client
	var err error

	if cfg.UseTLS {
		c, err = client.DialTLS(addr, nil)
	} else {
		c, err = client.Dial(addr)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to IMAP server: %v", err)
	}

	if err := c.Login(cfg.Username, cfg.Password); err != nil {
		return nil, fmt.Errorf("failed to login: %v", err)
	}

	return &IMAPClient{
		client: c,
		config: cfg,
	}, nil
}

type Email struct {
	ID      string
	Subject string
	From    string
	Date    time.Time
	Body    string
}

func (ic *IMAPClient) ensureConnection() error {
	if ic.client == nil {
		return ic.reconnect()
	}

	// Try to check connection by NOOPing
	if err := ic.client.Noop(); err != nil {
		return ic.reconnect()
	}
	return nil
}

func (ic *IMAPClient) reconnect() error {
	addr := fmt.Sprintf("%s:%d", ic.config.Server, ic.config.Port)
	var c *client.Client
	var err error

	if ic.config.UseTLS {
		c, err = client.DialTLS(addr, nil)
	} else {
		c, err = client.Dial(addr)
	}
	if err != nil {
		return fmt.Errorf("failed to connect to IMAP server: %v", err)
	}

	if err := c.Login(ic.config.Username, ic.config.Password); err != nil {
		c.Close()
		return fmt.Errorf("failed to login: %v", err)
	}

	if ic.client != nil {
		ic.client.Close()
	}
	ic.client = c
	return nil
}

func (ic *IMAPClient) FetchUnprocessedEmails() ([]Email, error) {
	if err := ic.ensureConnection(); err != nil {
		return nil, fmt.Errorf("failed to ensure IMAP connection: %v", err)
	}

	// Make sure we have a valid mailbox name
	if ic.config.MailboxIn == "" {
		return nil, fmt.Errorf("mailbox_in not configured")
	}

	// Select mailbox and get status
	mbox, err := ic.client.Select(ic.config.MailboxIn, false)
	if err != nil {
		return nil, fmt.Errorf("failed to select mailbox %s: %v", ic.config.MailboxIn, err)
	}

	// If mailbox is empty, return early
	if mbox.Messages == 0 {
		return nil, nil
	}

	// Get all messages in the inbox that haven't been seen yet and haven't been AI-processed
	criteria := imap.NewSearchCriteria()
	criteria.WithoutFlags = []string{"\\Seen", AIProcessedFlag}

	uids, err := ic.client.Search(criteria)
	if err != nil {
		return nil, fmt.Errorf("search failed: %v", err)
	}

	if len(uids) == 0 {
		return nil, nil
	}

	seqSet := new(imap.SeqSet)
	seqSet.AddNum(uids...)

	// Fetch both envelope and body
	section := &imap.BodySectionName{Peek: true}
	items := []imap.FetchItem{imap.FetchEnvelope, section.FetchItem()}

	messages := make(chan *imap.Message, 10)
	done := make(chan error, 1)

	go func() {
		done <- ic.client.Fetch(seqSet, items, messages)
	}()

	var emails []Email
	for msg := range messages {
		// Skip if message doesn't have an envelope
		if msg.Envelope == nil {
			continue
		}

		// Skip if no message ID
		if msg.Envelope.MessageId == "" {
			continue
		}

		r := msg.GetBody(section)
		if r == nil {
			continue
		}

		var bodyBuilder strings.Builder
		if _, err := io.Copy(&bodyBuilder, r); err != nil {
			continue
		}

		// Make sure we have a valid from address
		var from string
		if len(msg.Envelope.From) > 0 {
			from = msg.Envelope.From[0].Address()
		}

		emails = append(emails, Email{
			ID:      msg.Envelope.MessageId,
			Subject: msg.Envelope.Subject,
			From:    from,
			Date:    msg.Envelope.Date,
			Body:    bodyBuilder.String(),
		})
	}

	if err := <-done; err != nil {
		return nil, fmt.Errorf("fetch failed: %v", err)
	}

	return emails, nil
}

func (ic *IMAPClient) SaveDraft(email Email, response string) (err error) {
	defer func() {
		if r := recover(); r != nil {
			logger.WithFields(logrus.Fields{
				"panic":        r,
				"emailSubject": email.Subject,
				"draftBox":     ic.config.DraftBox,
			}).Error("Panic occurred in SaveDraft")
			err = fmt.Errorf("panic in SaveDraft: %v", r)
		}
	}()

	if err := ic.ensureConnection(); err != nil {
		return fmt.Errorf("failed to ensure IMAP connection: %v", err)
	}

	// Get proper folder path and ensure it exists
	draftFolder, err := ic.ensureFolder(ic.config.DraftBox)
	if err != nil {
		return fmt.Errorf("failed to ensure draft folder: %s, error: %v", ic.config.DraftBox, err)
	}
	logger.WithField("draftFolder", draftFolder).Debug("Ensured draft folder path")

	logger.WithField("draftFolder", draftFolder).Debug("Attempting to select draft folder")
	_, err = ic.client.Select(draftFolder, false)
	if err != nil {
		return fmt.Errorf("failed to select draft box '%s': %v", draftFolder, err)
	}
	logger.WithField("draftFolder", draftFolder).Debug("Successfully selected draft folder")

	// Format the draft message with HTML response and original email headers + content
	draft := fmt.Sprintf("From: %s\r\n"+
		"To: %s\r\n"+
		"Subject: Re: %s\r\n"+
		"MIME-Version: 1.0\r\n"+
		"Content-Type: text/html; charset=UTF-8\r\n"+
		"\r\n"+
		"%s\r\n"+
		"<br><br>\r\n"+
		"<div style=\"border-top: 1px solid #B5C4DF; padding-top: 10px; margin-top: 10px;\">\r\n"+
		"<div style=\"font-family: Arial, sans-serif; color: #666666; margin-bottom: 10px;\">\r\n"+
		"<b>From:</b> %s<br>\r\n"+
		"<b>Subject:</b> %s<br>\r\n"+
		"<b>Date:</b> %s<br>\r\n"+
		"</div>\r\n"+
		"<div style=\"margin-left: 10px; padding-left: 10px; border-left: 2px solid #B5C4DF;\">\r\n"+
		"%s\r\n"+
		"</div></div>",
		ic.config.Username,
		email.From,
		email.Subject,
		response,
		email.From,
		email.Subject,
		email.Date.Format("Mon, 02 Jan 2006 15:04:05 -0700"),
		ExtractMessageContent(email.Body))

	literal := &MessageLiteral{
		content: []byte(draft),
		pos:     0,
	}

	// Save the draft to the proper folder path
	flags := []string{"\\Draft"}
	logger.WithFields(logrus.Fields{
		"draftFolder": draftFolder,
		"flags":       flags,
	}).Debug("Attempting to append draft")
	if err := ic.client.Append(draftFolder, flags, time.Now(), literal); err != nil {
		return fmt.Errorf("failed to append draft to '%s': %v", draftFolder, err)
	}
	logger.WithField("draftFolder", draftFolder).Debug("Successfully appended draft")

	return nil
}

// ExtractMessageContent attempts to extract just the message content
// by removing email headers and MIME boundaries
func ExtractMessageContent(body string) string {
	logger.WithField("bodyLength", len(body)).Debug("Starting message content extraction")
	logger.WithField("rawInputBody", body).Debug("extractMessageContent: Raw input body")
	msg, err := mail.ReadMessage(strings.NewReader(body))
	if err != nil {
		logger.WithFields(logrus.Fields{
			"error":      err,
			"bodyLength": len(body),
		}).Debug("Failed to parse email message, falling back to simple extraction")
		// When ReadMessage fails, the body is raw, so header stripping is needed.
		return cleanMessageContent(fallbackExtractContent(body), true)
	}

	contentTypeHeader := msg.Header.Get("Content-Type")
	logger.WithField("contentTypeHeader", contentTypeHeader).Debug("Got Content-Type header")
	logger.WithField("parsedContentTypeHeader", contentTypeHeader).Debug("extractMessageContent: Parsed Content-Type header")

	mediaType, params, err := mime.ParseMediaType(contentTypeHeader)
	if err != nil {
		logger.WithFields(logrus.Fields{
			"error":             err,
			"contentTypeHeader": contentTypeHeader,
		}).Debug("Failed to parse Content-Type header, falling back to simple extraction")
		// When ParseMediaType fails, the body is raw, so header stripping is needed.
		return cleanMessageContent(fallbackExtractContent(body), true)
	}

	logger.WithFields(logrus.Fields{
		"mediaType": mediaType,
		"params":    params,
	}).Debug("Parsed message Content-Type")
	logger.WithFields(logrus.Fields{
		"mediaType": mediaType,
		"params":    params,
	}).Debug("extractMessageContent: Parsed mediaType and params")

	var content string
	if strings.HasPrefix(mediaType, "multipart/") {
		// For multipart, the handling of Content-Transfer-Encoding will be done within handleMultipartMessage for each part
		logger.Debug("extractMessageContent: Handling as multipart message")
		content = handleMultipartMessage(msg.Body, params["boundary"])
		logger.WithField("contentFromMultipart", content).Debug("extractMessageContent: Content after handleMultipartMessage")
	} else {
		// For single part, handle Content-Transfer-Encoding here
		var partReader io.Reader = msg.Body
		transferEncoding := strings.ToLower(msg.Header.Get("Content-Transfer-Encoding"))
		if transferEncoding == "quoted-printable" {
			partReader = quotedprintable.NewReader(msg.Body)
		}
		// Add handling for "base64" if needed in the future
		// else if transferEncoding == "base64" {
		// 	partReader = base64.NewDecoder(base64.StdEncoding, msg.Body)
		// }
		logger.Debug("extractMessageContent: Handling as single part message")
		content = handleSinglePartMessage(partReader)
		logger.WithField("contentFromSinglePart", content).Debug("extractMessageContent: Content after handleSinglePartMessage")
	}

	if content == "" {
		logger.Debug("No content extracted, falling back to simple extraction")
		logger.Debug("extractMessageContent: No content from primary extraction, falling back.")
		// When primary extraction yields no content, the body is raw, so header stripping is needed.
		return cleanMessageContent(fallbackExtractContent(body), true)
	}

	// Clean up the content
	logger.WithField("contentBeforeClean", content).Debug("extractMessageContent: Content before cleanMessageContent")
	// When ReadMessage succeeds, 'content' is already the message body (or part body),
	// so no further header stripping should be done.
	content = cleanMessageContent(content, false)
	logger.WithField("contentAfterClean", content).Debug("extractMessageContent: Content after cleanMessageContent")

	logger.WithField("contentLength", len(content)).Debug("Successfully extracted and cleaned message content")
	return content
}

func handleMultipartMessage(reader io.Reader, boundary string) string {
	logger.WithField("boundary", boundary).Debug("handleMultipartMessage: Starting with boundary")
	if boundary == "" {
		logger.Debug("No boundary found in multipart message")
		return ""
	}

	mReader := multipart.NewReader(reader, boundary)
	var textContent string
	partIndex := 0

	for {
		part, err := mReader.NextPart()
		if err == io.EOF {
			break
		}
		if err != nil {
			logger.WithError(err).Debug("Error reading multipart part")
			return ""
		}

		contentType := part.Header.Get("Content-Type")
		contentTransferEncoding := strings.ToLower(part.Header.Get("Content-Transfer-Encoding"))
		logger.WithFields(logrus.Fields{
			"partIndex":            partIndex,
			"partContentType":      contentType,
			"partTransferEncoding": contentTransferEncoding,
			"partHeaders":          part.Header,
		}).Debug("handleMultipartMessage: Processing part")

		if strings.HasPrefix(contentType, "text/plain") {
			var partReader io.Reader = part
			if contentTransferEncoding == "quoted-printable" {
				partReader = quotedprintable.NewReader(part)
			}
			// Add handling for "base64" if needed in the future
			// else if contentTransferEncoding == "base64" {
			// 	partReader = base64.NewDecoder(base64.StdEncoding, part)
			// }

			buf := new(bytes.Buffer)
			if _, err := buf.ReadFrom(partReader); err != nil {
				logger.WithError(err).WithField("partIndex", partIndex).Debug("Failed to read from partReader in multipart")
				continue // Or handle error more robustly
			}
			textContent = buf.String()
			// Assuming we only care about the first text/plain part found
			// If multiple text/plain parts could exist and need concatenation, this logic would need adjustment.
			break
		}
		partIndex++
	}

	logger.WithField("textContentResult", textContent).Debug("handleMultipartMessage: Returning textContent")
	return textContent
}

func handleSinglePartMessage(reader io.Reader) string {
	logger.Debug("handleSinglePartMessage: Starting")
	buf := new(bytes.Buffer)
	if _, err := buf.ReadFrom(reader); err != nil {
		logger.WithError(err).Debug("Failed to read message body")
		return ""
	}
	content := buf.String()
	logger.WithField("readContent", content).Debug("handleSinglePartMessage: Content read from reader")
	return content
}

func cleanMessageContent(content string, performHeaderStripping bool) string {
	logger.WithField("inputContentLength", len(content)).Debug("cleanMessageContent: Starting")
	logger.WithField("performHeaderStripping", performHeaderStripping).Debug("cleanMessageContent: performHeaderStripping flag")

	if performHeaderStripping {
		// Remove any remaining email headers that might be in the body
		lines := strings.Split(content, "\n")
		var cleanLines []string
		headerSection := true

		for _, line := range lines {
			trimmed := strings.TrimSpace(line)

			// Empty line marks the end of headers
			if headerSection && trimmed == "" {
				headerSection = false
				continue
			}

			// Skip header lines
			if headerSection && (strings.Contains(trimmed, ":") || trimmed == "") {
				continue
			}

			// Add non-header lines
			// This condition was originally !headerSection, but if we are past the headers,
			// we should always add the line. If headerSection is still true here, it means
			// it's the first line of content after potential headers were skipped.
			cleanLines = append(cleanLines, line)
		}
		content = strings.Join(cleanLines, "\n")
		logger.WithField("contentAfterHeaderStripLength", len(content)).Debug("cleanMessageContent: Content after header stripping")
	} else {
		logger.Debug("cleanMessageContent: Skipping header stripping")
	}

	// Convert newlines to HTML breaks for display
	// First, normalize all newlines (\r\n or \n) to just \n
	content = strings.ReplaceAll(content, "\r\n", "\n")
	// Then, replace each \n with <br>\n
	content = strings.ReplaceAll(content, "\n", "<br>\n")
	logger.WithField("contentAfterNewlineConversionLength", len(content)).Debug("cleanMessageContent: Content after newline conversion")

	// Remove any remaining email signature markers
	content = strings.Split(content, "\n-- ")[0]
	content = strings.Split(content, "<br>-- ")[0]

	finalContent := strings.TrimSpace(content)
	logger.WithField("finalCleanedContentLength", len(finalContent)).Debug("cleanMessageContent: Returning final cleaned content")
	return finalContent
}

// fallbackExtractContent is the previous implementation used as fallback
func fallbackExtractContent(body string) string {
	logger.WithField("bodyLength", len(body)).Debug("Using fallback content extraction method")
	logger.WithField("rawInputBodyFallbackLength", len(body)).Debug("fallbackExtractContent: Raw input body")
	var content string
	parts := strings.Split(body, "\r\n\r\n")
	if len(parts) > 1 {
		content = strings.Join(parts[1:], "\r\n\r\n")
		logger.WithFields(logrus.Fields{
			"contentLength": len(content),
			"partsCount":    len(parts),
		}).Debug("Successfully extracted content using fallback method (from parts)")
		logger.WithField("extractedContentFallbackLength", len(content)).Debug("fallbackExtractContent: Content from splitting parts")
	} else {
		content = body
		logger.WithFields(logrus.Fields{
			"contentLength": len(content),
			"fullBody":      true,
		}).Debug("Using full body as content in fallback method")
		logger.WithField("finalContentFallbackLength", len(content)).Debug("fallbackExtractContent: Final content from full body (no parts split)")
	}
	// Newline conversion and signature stripping will be handled by cleanMessageContent,
	// which is now called by the callers of fallbackExtractContent.
	return content
}

// ensureFolder makes sure a folder exists and returns its full path using proper delimiters
func (ic *IMAPClient) ensureFolder(folderName string) (path string, err error) {
	defer func() {
		if r := recover(); r != nil {
			logger.WithFields(logrus.Fields{
				"panic":      r,
				"folderName": folderName,
			}).Error("Panic occurred in ensureFolder")
			err = fmt.Errorf("panic in ensureFolder: %v", r)
		}
	}()
	logger.WithField("folderNameInput", folderName).Debug("Ensuring folder exists (ensureFolder start)")

	var folderPath string
	isGmail := strings.ToLower(ic.config.Server) == "imap.gmail.com"

	if isGmail {
		logger.Debug("ensureFolder: Detected Gmail server. Using Gmail-specific logic.")
		// Gmail always uses '/' as a delimiter and folder names are typically as provided.
		delimiter := "/"
		folderPath = folderName // Gmail folder names in config should already be correct.
		logger.WithFields(logrus.Fields{
			"folderName":        folderName,
			"assumedDelimiter":  delimiter,
			"derivedFolderPath": folderPath,
		}).Debug("ensureFolder: Gmail - folder path set")

		// For Gmail, don't try to CREATE system folders like "[Gmail]/..." or "INBOX"
		// For "INBOX/Subfolder", CREATE is fine as it creates a label.
		if strings.HasPrefix(folderPath, "[Gmail]/") || folderPath == "INBOX" {
			logger.WithField("folderPath", folderPath).Debug("ensureFolder: Gmail - Skipping CREATE for system folder.")
			return folderPath, nil
		}
		// For other paths like "INBOX/Done" or "MyCustomLabel", attempt CREATE.
		logger.WithField("folderPath", folderPath).Debug("ensureFolder: Gmail - Attempting to create folder (label).")
		if createErr := ic.client.Create(folderPath); createErr != nil {
			if !strings.Contains(strings.ToLower(createErr.Error()), "already exists") &&
				!strings.Contains(strings.ToLower(createErr.Error()), "mailbox exists") {
				logger.WithFields(logrus.Fields{
					"folder": folderPath,
					"error":  createErr,
				}).Warn("ensureFolder: Gmail - Folder creation failed (and not 'already exists')")
				// Unlike non-Gmail, if CREATE fails for a non-system folder, it's more likely a real issue.
				// However, subsequent SELECT/APPEND will ultimately determine usability.
				// For now, we'll proceed and let later operations fail if it's critical.
			} else {
				logger.WithFields(logrus.Fields{
					"folder": folderPath,
					"error":  createErr,
				}).Debug("ensureFolder: Gmail - Folder creation failed (likely already exists).")
			}
		} else {
			logger.WithField("folderPath", folderPath).Info("ensureFolder: Gmail - Successfully created folder/label or it already existed.")
		}
		return folderPath, nil

	} else {
		logger.Debug("ensureFolder: Non-Gmail server. Using generic logic.")
		// Generic logic for non-Gmail servers (existing logic with timeouts)
		logger.Debug("ensureFolder: Creating channels mailboxes and listDone")
		mailboxes := make(chan *imap.MailboxInfo, 10)
		listDone := make(chan error, 1)
		logger.Debug("ensureFolder: Channels created. Attempting to start List goroutine")

		go func() {
			logger.Debug("ensureFolder: List goroutine started")
			defer func() {
				if r := recover(); r != nil {
					logger.WithFields(logrus.Fields{
						"panic":      r,
						"folderName": folderName,
					}).Error("Panic occurred in ensureFolder's List goroutine")
					listDone <- fmt.Errorf("panic in List goroutine: %v", r)
				}
			}()
			logger.Debug("ensureFolder: List goroutine: Entering select for client.List or timeout")
			select {
			case listDone <- ic.client.List("", "*", mailboxes):
				logger.Debug("ensureFolder: List goroutine: client.List call completed and sent to listDone")
			case <-time.After(30 * time.Second):
				logger.Error("Timeout occurred during client.List operation in goroutine")
				listDone <- fmt.Errorf("client.List timeout after 30 seconds in goroutine")
			}
			logger.Debug("ensureFolder: List goroutine finished")
		}()
		logger.Debug("ensureFolder: List goroutine launched.")

		var delimiter string
		var receivedMailbox bool
		logger.Debug("ensureFolder: Attempting to range over mailboxes channel")
		for m := range mailboxes {
			receivedMailbox = true
			logger.WithField("mailboxName", m.Name).Debug("ensureFolder: Received mailbox from channel")
			delimiter = m.Delimiter
			logger.WithField("delimiter", delimiter).Debug("ensureFolder: Got delimiter from mailbox info")
			break
		}
		if !receivedMailbox {
			logger.Debug("ensureFolder: Did not receive any mailboxes from channel (it might have been empty or closed early).")
		}
		logger.Debug("ensureFolder: Finished ranging over mailboxes. Attempting to read from listDone channel with timeout")

		select {
		case listErr := <-listDone:
			if listErr != nil {
				logger.WithError(listErr).Error("ensureFolder: Error received from listDone channel")
				return "", fmt.Errorf("failed to list mailboxes to determine delimiter: %v", listErr)
			}
			logger.Debug("ensureFolder: Successfully read from listDone channel.")
		case <-time.After(35 * time.Second):
			logger.Error("ensureFolder: Timeout waiting for listDone channel")
			return "", fmt.Errorf("timeout waiting for LIST operation to complete")
		}

		if delimiter == "" {
			logger.Debug("ensureFolder: Delimiter is still empty after processing listDone and mailboxes.")
			delimiter = "/"
			logger.Debug("No delimiter returned by server, using fallback '/'")
		}
		logger.WithField("delimiter", delimiter).Debug("Determined mailbox delimiter")

		if delimiter != "/" {
			folderPath = strings.ReplaceAll(folderName, "/", delimiter)
		} else {
			folderPath = folderName
		}
		logger.WithFields(logrus.Fields{
			"originalFolderName": folderName,
			"serverDelimiter":    delimiter,
			"derivedFolderPath":  folderPath,
		}).Debug("Derived folder path using server delimiter")

		logger.WithField("folderPath", folderPath).Debug("Attempting to create folder if it doesn't exist")
		if createErr := ic.client.Create(folderPath); createErr != nil {
			if !strings.Contains(strings.ToLower(createErr.Error()), "already exists") &&
				!strings.Contains(strings.ToLower(createErr.Error()), "mailbox exists") {
				logger.WithFields(logrus.Fields{
					"folder": folderPath,
					"error":  createErr,
				}).Warn("Folder creation failed (and not 'already exists')")
			} else {
				logger.WithFields(logrus.Fields{
					"folder": folderPath,
					"error":  createErr,
				}).Debug("Folder creation failed (likely already exists)")
			}
		} else {
			logger.WithField("folderPath", folderPath).Info("Successfully created folder or it already existed")
		}
		return folderPath, nil
	}
}

func (ic *IMAPClient) MarkAsProcessed(email Email) (err error) {
	defer func() {
		if r := recover(); r != nil {
			logger.WithFields(logrus.Fields{
				"panic":        r,
				"emailSubject": email.Subject,
				"processedBox": ic.config.ProcessedBox,
			}).Error("Panic occurred in MarkAsProcessed")
			err = fmt.Errorf("panic in MarkAsProcessed: %v", r)
		}
	}()

	if err := ic.ensureConnection(); err != nil {
		return fmt.Errorf("failed to ensure IMAP connection: %v", err)
	}

	if ic.config.ProcessedBox == "" {
		return fmt.Errorf("processed_box not configured")
	}

	// Get proper folder path and ensure it exists
	processedFolder, err := ic.ensureFolder(ic.config.ProcessedBox)
	if err != nil {
		return fmt.Errorf("failed to ensure processed folder: %s, error: %v", ic.config.ProcessedBox, err)
	}
	logger.WithField("processedFolder", processedFolder).Debug("Ensured processed folder path")

	// Select source mailbox
	_, err = ic.client.Select(ic.config.MailboxIn, false)
	if err != nil {
		return fmt.Errorf("failed to select source mailbox: %v", err)
	}

	// Find the email by Message-Id
	criteria := imap.NewSearchCriteria()
	criteria.Header.Set("Message-Id", email.ID)

	uids, err := ic.client.Search(criteria)
	if err != nil {
		return fmt.Errorf("search failed: %v", err)
	}

	if len(uids) == 0 {
		return fmt.Errorf("email not found")
	}

	// Move the message to processed folder
	seqSet := new(imap.SeqSet)
	seqSet.AddNum(uids...)

	// Mark as read before moving
	item := imap.FormatFlagsOp(imap.AddFlags, true)
	flags := []interface{}{imap.SeenFlag}
	if err := ic.client.Store(seqSet, item, flags, nil); err != nil {
		return fmt.Errorf("failed to mark message as read: %v", err)
	}

	// Copy to processed folder using the proper path
	if err := ic.client.Copy(seqSet, processedFolder); err != nil {
		return fmt.Errorf("failed to copy to processed folder: %v", err)
	}

	// Delete from source folder
	item = imap.FormatFlagsOp(imap.AddFlags, true)
	flags = []interface{}{imap.DeletedFlag}
	if err := ic.client.Store(seqSet, item, flags, nil); err != nil {
		return fmt.Errorf("failed to mark message as deleted: %v", err)
	}

	if err := ic.client.Expunge(nil); err != nil {
		return fmt.Errorf("failed to expunge message: %v", err)
	}

	return nil
}

// MarkAsAIProcessed marks an email with a custom flag to indicate AI has processed it
// This prevents reprocessing the same email if subsequent operations (SaveDraft, MarkAsProcessed) fail
func (ic *IMAPClient) MarkAsAIProcessed(email Email) (err error) {
	defer func() {
		if r := recover(); r != nil {
			logger.WithFields(logrus.Fields{
				"panic":        r,
				"emailSubject": email.Subject,
			}).Error("Panic occurred in MarkAsAIProcessed")
			err = fmt.Errorf("panic in MarkAsAIProcessed: %v", r)
		}
	}()

	if err := ic.ensureConnection(); err != nil {
		return fmt.Errorf("failed to ensure IMAP connection: %v", err)
	}

	// Select source mailbox
	_, err = ic.client.Select(ic.config.MailboxIn, false)
	if err != nil {
		return fmt.Errorf("failed to select source mailbox: %v", err)
	}

	// Find the email by Message-Id
	criteria := imap.NewSearchCriteria()
	criteria.Header.Set("Message-Id", email.ID)

	uids, err := ic.client.Search(criteria)
	if err != nil {
		return fmt.Errorf("search failed: %v", err)
	}

	if len(uids) == 0 {
		return fmt.Errorf("email not found")
	}

	// Mark with AI-processed flag
	seqSet := new(imap.SeqSet)
	seqSet.AddNum(uids...)

	item := imap.FormatFlagsOp(imap.AddFlags, true)
	flags := []interface{}{AIProcessedFlag}
	if err := ic.client.Store(seqSet, item, flags, nil); err != nil {
		// If custom flags are not supported, fall back to marking as \Seen
		logger.WithFields(logrus.Fields{
			"error":   err,
			"subject": email.Subject,
		}).Warn("Failed to set custom AI-processed flag, falling back to \\Seen flag")

		flags = []interface{}{imap.SeenFlag}
		if err := ic.client.Store(seqSet, item, flags, nil); err != nil {
			return fmt.Errorf("failed to mark message with fallback flag: %v", err)
		}
		logger.WithField("subject", email.Subject).Info("Marked email as \\Seen (custom flag not supported)")
		return nil
	}

	logger.WithField("subject", email.Subject).Debug("Successfully marked email with AI-processed flag")
	return nil
}

func (ic *IMAPClient) Close() error {
	if ic.client != nil {
		ic.client.Logout()
	}
	return nil
}