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
processed map[string]bool
}
// 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,
processed: make(map[string]bool),
}, 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
}
criteria := imap.NewSearchCriteria()
criteria.WithoutFlags = []string{"\\Seen", "Processed"}
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 already processed
if ic.processed[msg.Envelope.MessageId] {
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(),
})
ic.processed[msg.Envelope.MessageId] = true
}
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)
}
_, err := ic.client.Select(ic.config.DraftBox, 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
flags := []string{"\\Draft"}
if err := ic.client.Append(ic.config.DraftBox, 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")
if strings.HasPrefix(mediaType, "multipart/") {
boundary := params["boundary"]
if boundary == "" {
logger.WithField("mediaType", mediaType).Debug("No boundary found in multipart message, falling back to simple extraction")
return fallbackExtractContent(body)
}
logger.WithFields(logrus.Fields{
"mediaType": mediaType,
"boundary": boundary,
}).Debug("Processing multipart message")
reader := multipart.NewReader(msg.Body, boundary)
var textContent string
partIndex := 0
for {
part, err := reader.NextPart()
if err == io.EOF {
logger.Debug("Finished processing all multipart parts")
break
}
if err != nil {
logger.WithFields(logrus.Fields{
"error": err,
"partIndex": partIndex,
}).Debug("Error reading multipart part, falling back to simple extraction")
return fallbackExtractContent(body)
}
contentType := part.Header.Get("Content-Type")
logger.WithFields(logrus.Fields{
"partIndex": partIndex,
"contentType": contentType,
}).Debug("Processing message part")
if strings.HasPrefix(contentType, "text/plain") {
buf := new(bytes.Buffer)
_, err := buf.ReadFrom(part)
if err != nil {
logger.WithFields(logrus.Fields{
"error": err,
"partIndex": partIndex,
}).Debug("Failed to read text/plain part")
continue
}
textContent = buf.String()
logger.WithFields(logrus.Fields{
"partIndex": partIndex,
"contentLength": len(textContent),
}).Debug("Successfully extracted text/plain content")
textContent = strings.ReplaceAll(textContent, "\n", "
\n")
textContent = strings.ReplaceAll(textContent, "\r\n", "
\n")
break
}
partIndex++
}
if textContent != "" {
logger.WithField("contentLength", len(textContent)).Debug("Successfully extracted content from multipart message")
return textContent
}
logger.Debug("No text/plain content found in multipart message, trying to read body directly")
}
buf := new(bytes.Buffer)
_, err = buf.ReadFrom(msg.Body)
if err != nil {
logger.WithFields(logrus.Fields{
"error": err,
"mediaType": mediaType,
}).Debug("Failed to read message body, falling back to simple extraction")
return fallbackExtractContent(body)
}
content := buf.String()
content = strings.ReplaceAll(content, "\r\n", "
\n")
content = strings.ReplaceAll(content, "\n", "
\n")
logger.WithFields(logrus.Fields{
"contentLength": len(content),
"mediaType": mediaType,
}).Debug("Successfully extracted content from message body")
return 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
}
func (ic *IMAPClient) MarkAsProcessed(email Email) error {
if err := ic.ensureConnection(); err != nil {
return fmt.Errorf("failed to ensure IMAP connection: %v", err)
}
_, err := ic.client.Select(ic.config.MailboxIn, false)
if err != nil {
return err
}
criteria := imap.NewSearchCriteria()
criteria.Header.Set("Message-Id", email.ID)
uids, err := ic.client.Search(criteria)
if err != nil {
return err
}
if len(uids) == 0 {
return fmt.Errorf("email not found")
}
seqSet := new(imap.SeqSet)
seqSet.AddNum(uids...)
return ic.client.Store(seqSet, imap.FormatFlagsOp(imap.AddFlags, true), []interface{}{"Processed"}, nil)
}
func (ic *IMAPClient) Close() error {
if ic.client != nil {
ic.client.Logout()
}
return nil
}