ai-mailer/internal/imap/imap.go

package imap

import (
	"fmt"
	"io"
	"paraclub-ai-mailer/config"
	"strings"
	"time"

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

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
}

func (m *MessageLiteral) Read(p []byte) (n int, err error) {
	return copy(p, m.content), io.EOF
}

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) FetchUnprocessedEmails() ([]Email, error) {
	_, err := ic.client.Select(ic.config.MailboxIn, false)
	if err != nil {
		return nil, err
	}

	criteria := imap.NewSearchCriteria()
	criteria.WithoutFlags = []string{"\\Seen", "Processed"}

	uids, err := ic.client.Search(criteria)
	if err != nil {
		return nil, err
	}

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

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

	messages := make(chan *imap.Message, 10)
	section := &imap.BodySectionName{Peek: true}

	done := make(chan error, 1)
	go func() {
		done <- ic.client.Fetch(seqSet, []imap.FetchItem{section.FetchItem()}, messages)
	}()

	var emails []Email
	for msg := range messages {
		if ic.processed[msg.Envelope.MessageId] {
			continue
		}

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

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

		emails = append(emails, Email{
			ID:      msg.Envelope.MessageId,
			Subject: msg.Envelope.Subject,
			From:    msg.Envelope.From[0].Address(),
			Body:    bodyBuilder.String(),
		})

		ic.processed[msg.Envelope.MessageId] = true
	}

	if err := <-done; err != nil {
		return nil, err
	}

	return emails, nil
}

func (ic *IMAPClient) SaveDraft(email Email, response string) error {
	_, err := ic.client.Select(ic.config.DraftBox, false)
	if err != nil {
		return err
	}

	// Format the draft message
	draft := fmt.Sprintf("From: %s\r\n"+
		"To: %s\r\n"+
		"Subject: Re: %s\r\n"+
		"Content-Type: text/plain; charset=UTF-8\r\n"+
		"\r\n%s",
		ic.config.Username,
		email.From,
		email.Subject,
		response)

	// Create literal message
	literal := &MessageLiteral{content: []byte(draft)}

	// Save the draft
	flags := []string{"\\Draft"}
	if err := ic.client.Append(ic.config.DraftBox, flags, time.Now(), literal); err != nil {
		return err
	}

	return nil
}

func (ic *IMAPClient) MarkAsProcessed(email Email) error {
	_, 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
}