refactor AI response generation and enhance IMAP connection handling

internal/ai/ai.go

@@ -48,14 +48,18 @@ func (a *AI) GenerateReply(emailContent string, contextContent map[string]string
 	}
 
 	// Prepare the system message and user message
-	systemMsg := "You are a helpful assistant. Use the provided context to help answer the email professionally and accurately."
+	systemMsg := "You are a helpful assistant responding to emails. Analyze the language of the incoming email and translate the response to the same language. Format your response in HTML. Do not include any explanations or extra text - just write the email response directly in HTML format. Use appropriate HTML tags for formatting."
-	userMsg := fmt.Sprintf("Using the following context:\n%s\n\nPlease generate a reply for this email:\n%s", context, emailContent)
+	userMsg := fmt.Sprintf("Using the following context:\n%s\n\nPlease analyze the language of and generate a reply in the same language for this email:\n%s", context, emailContent)
 
 	messages := []Message{
 		{Role: "system", Content: systemMsg},
 		{Role: "user", Content: userMsg},
 	}
 
+	const maxRetries = 3
+	var lastErr error
+
+	for attempt := 0; attempt < maxRetries; attempt++ {
 		reqBody := OpenRouterRequest{
 			Model:       a.config.Model,
 			Messages:    messages,
@@ -78,22 +82,29 @@ func (a *AI) GenerateReply(emailContent string, contextContent map[string]string
 
 		resp, err := a.client.Do(req)
 		if err != nil {
-		return "", err
+			lastErr = err
+			continue
 		}
 		defer resp.Body.Close()
 
 		if resp.StatusCode != http.StatusOK {
-		return "", fmt.Errorf("OpenRouter API returned status code: %d", resp.StatusCode)
+			lastErr = fmt.Errorf("OpenRouter API returned status code: %d", resp.StatusCode)
+			continue
 		}
 
 		var result OpenRouterResponse
 		if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
-		return "", err
+			lastErr = err
+			continue
 		}
 
-	if len(result.Choices) == 0 {
+		if len(result.Choices) == 0 || result.Choices[0].Message.Content == "" {
-		return "", fmt.Errorf("no response generated")
+			lastErr = fmt.Errorf("empty response received")
+			continue
 		}
 
 		return result.Choices[0].Message.Content, nil
 	}
+
+	return "", fmt.Errorf("failed after %d attempts, last error: %v", maxRetries, lastErr)
+}

internal/imap/imap.go

@@ -1,8 +1,12 @@
 package imap
 
 import (
+	"bytes"
 	"fmt"
 	"io"
+	"mime"
+	"mime/multipart"
+	"net/mail"
 	"paraclub-ai-mailer/config"
 	"strings"
 	"time"
@@ -20,10 +24,19 @@ type IMAPClient struct {
 // MessageLiteral implements imap.Literal for draft messages
 type MessageLiteral struct {
 	content []byte
+	pos     int
 }
 
 func (m *MessageLiteral) Read(p []byte) (n int, err error) {
-	return copy(p, m.content), io.EOF
+	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 {
@@ -63,10 +76,63 @@ type Email struct {
 	Body    string
 }
 
-func (ic *IMAPClient) FetchUnprocessedEmails() ([]Email, error) {
+func (ic *IMAPClient) ensureConnection() error {
-	_, err := ic.client.Select(ic.config.MailboxIn, false)
+	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 nil, err
+		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()
@@ -74,7 +140,7 @@ func (ic *IMAPClient) FetchUnprocessedEmails() ([]Email, error) {
 
 	uids, err := ic.client.Search(criteria)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("search failed: %v", err)
 	}
 
 	if len(uids) == 0 {
@@ -84,35 +150,54 @@ func (ic *IMAPClient) FetchUnprocessedEmails() ([]Email, error) {
 	seqSet := new(imap.SeqSet)
 	seqSet.AddNum(uids...)
 
-	messages := make(chan *imap.Message, 10)
+	// 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, []imap.FetchItem{section.FetchItem()}, messages)
+		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
-		_, err := io.Copy(&bodyBuilder, r)
+		if _, err := io.Copy(&bodyBuilder, r); err != nil {
-		if 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:    msg.Envelope.From[0].Address(),
+			From:    from,
 			Body:    bodyBuilder.String(),
 		})
 
@@ -120,42 +205,137 @@ func (ic *IMAPClient) FetchUnprocessedEmails() ([]Email, error) {
 	}
 
 	if err := <-done; err != nil {
-		return nil, err
+		return nil, fmt.Errorf("fetch failed: %v", err)
 	}
 
 	return emails, nil
 }
 
 func (ic *IMAPClient) SaveDraft(email Email, response string) error {
-	_, err := ic.client.Select(ic.config.DraftBox, false)
+	if err := ic.ensureConnection(); err != nil {
-	if err != nil {
+		return fmt.Errorf("failed to ensure IMAP connection: %v", err)
-		return err
 	}
 
-	// Format the draft message
+	_, 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"+
-		"Content-Type: text/plain; charset=UTF-8\r\n"+
+		"MIME-Version: 1.0\r\n"+
-		"\r\n%s",
+		"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)
+		response,
+		email.From,
+		email.Subject,
+		time.Now().Format("Mon, 02 Jan 2006 15:04:05 -0700"),
+		extractMessageContent(email.Body))
 
-	// Create literal message
+	literal := &MessageLiteral{
-	literal := &MessageLiteral{content: []byte(draft)}
+		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 err
+		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 {
+	msg, err := mail.ReadMessage(strings.NewReader(body))
+	if err != nil {
+		// Fallback to previous method if parsing fails
+		return fallbackExtractContent(body)
+	}
+
+	mediaType, params, err := mime.ParseMediaType(msg.Header.Get("Content-Type"))
+	if err != nil {
+		return fallbackExtractContent(body)
+	}
+
+	if strings.HasPrefix(mediaType, "multipart/") {
+		boundary := params["boundary"]
+		if boundary == "" {
+			return fallbackExtractContent(body)
+		}
+
+		reader := multipart.NewReader(msg.Body, boundary)
+		var textContent string
+
+		for {
+			part, err := reader.NextPart()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				return fallbackExtractContent(body)
+			}
+
+			// Check Content-Type of the part
+			contentType := part.Header.Get("Content-Type")
+			if strings.HasPrefix(contentType, "text/plain") {
+				buf := new(bytes.Buffer)
+				_, err := buf.ReadFrom(part)
+				if err != nil {
+					continue
+				}
+				textContent = buf.String()
+				break
+			}
+		}
+
+		if textContent != "" {
+			return strings.TrimSpace(textContent)
+		}
+	}
+
+	// For non-multipart messages, just read the body
+	buf := new(bytes.Buffer)
+	_, err = buf.ReadFrom(msg.Body)
+	if err != nil {
+		return fallbackExtractContent(body)
+	}
+
+	return strings.TrimSpace(buf.String())
+}
+
+// fallbackExtractContent is the previous implementation used as fallback
+func fallbackExtractContent(body string) string {
+	parts := strings.Split(body, "\r\n\r\n")
+	if len(parts) > 1 {
+		return strings.TrimSpace(strings.Join(parts[1:], "\r\n\r\n"))
+	}
+	return strings.TrimSpace(body)
+}
+
 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