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) error {
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 exists: %v", err)
}
_, err = ic.client.Select(draftFolder, false)
if err != nil {
return fmt.Errorf("failed to select draft box: %v", err)
}
// 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"+
"
\r\n"+
"
\r\n"+
"
\r\n"+
"From: %s
\r\n"+
"Subject: %s
\r\n"+
"Date: %s
\r\n"+
"
\r\n"+
"
\r\n"+
"%s\r\n"+
"
",
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"}
if err := ic.client.Append(draftFolder, flags, time.Now(), literal); err != nil {
return fmt.Errorf("failed to append draft: %v", err)
}
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", "
\n")
content = strings.ReplaceAll(content, "\n", "
\n")
// Remove any remaining email signature markers
content = strings.Split(content, "\n-- ")[0]
content = strings.Split(content, "
-- ")[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", "
\n")
content = strings.ReplaceAll(content, "\n", "
\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", "
\n")
content = strings.ReplaceAll(content, "\n", "
\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) (string, error) {
// List all mailboxes to get the delimiter
mailboxes := make(chan *imap.MailboxInfo, 10)
done := make(chan error, 1)
go func() {
done <- ic.client.List("", "*", mailboxes)
}()
var delimiter string
for m := range mailboxes {
delimiter = m.Delimiter
break // We just need the first one to get the delimiter
}
if err := <-done; err != nil {
return "", fmt.Errorf("failed to list mailboxes: %v", err)
}
if delimiter == "" {
delimiter = "/" // fallback to common delimiter
}
// Replace any forward slashes with the server's delimiter
folderPath := strings.ReplaceAll(folderName, "/", delimiter)
// Try to create the folder if it doesn't exist
if err := ic.client.Create(folderPath); err != nil {
// Ignore errors if the folder already exists
logger.WithFields(logrus.Fields{
"folder": folderPath,
"error": err,
}).Debug("Folder creation failed (might already exist)")
}
return folderPath, nil
}
func (ic *IMAPClient) MarkAsProcessed(email Email) error {
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 exists: %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")
}
// 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
}