ai-mailer/internal/imap/imap.go

package imap

import (
	"bytes"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"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"
)

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
	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
	criteria := imap.NewSearchCriteria()
	criteria.WithoutFlags = []string{"\\Seen"}

	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,
			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,
		time.Now().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")
	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")
		return fallbackExtractContent(body)
	}

	contentTypeHeader := msg.Header.Get("Content-Type")
	logger.WithField("contentTypeHeader", contentTypeHeader).Debug("Got 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")
		return fallbackExtractContent(body)
	}

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

	var content string
	if strings.HasPrefix(mediaType, "multipart/") {
		content = handleMultipartMessage(msg.Body, params["boundary"])
	} else {
		content = handleSinglePartMessage(msg.Body)
	}

	if content == "" {
		logger.Debug("No content extracted, falling back to simple extraction")
		return fallbackExtractContent(body)
	}

	// Clean up the content
	content = cleanMessageContent(content)

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

func handleMultipartMessage(reader io.Reader, boundary string) string {
	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")
		if strings.HasPrefix(contentType, "text/plain") {
			buf := new(bytes.Buffer)
			if _, err := buf.ReadFrom(part); err != nil {
				continue
			}
			textContent = buf.String()
			break
		}
		partIndex++
	}

	return textContent
}

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

func cleanMessageContent(content string) string {
	// 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
		if !headerSection {
			cleanLines = append(cleanLines, line)
		}
	}

	content = strings.Join(cleanLines, "\n")

	// Convert newlines to HTML breaks for display
	content = strings.ReplaceAll(content, "\r\n", "<br>\n")
	content = strings.ReplaceAll(content, "\n", "<br>\n")

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

	return strings.TrimSpace(content)
}

// fallbackExtractContent is the previous implementation used as fallback
func fallbackExtractContent(body string) string {
	logger.WithField("bodyLength", len(body)).Debug("Using fallback content extraction method")
	parts := strings.Split(body, "\r\n\r\n")
	if len(parts) > 1 {
		content := strings.Join(parts[1:], "\r\n\r\n")
		content = strings.ReplaceAll(content, "\r\n", "<br>\n")
		content = strings.ReplaceAll(content, "\n", "<br>\n")
		logger.WithFields(logrus.Fields{
			"contentLength": len(content),
			"partsCount":    len(parts),
		}).Debug("Successfully extracted content using fallback method")
		return content
	}
	content := body
	content = strings.ReplaceAll(content, "\r\n", "<br>\n")
	content = strings.ReplaceAll(content, "\n", "<br>\n")
	logger.WithFields(logrus.Fields{
		"contentLength": len(content),
		"fullBody":      true,
	}).Debug("Using full body as content in fallback method")
	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
}

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