archivmail/internal/imapserver/server.go

// Package imapserver implements an embedded read-only IMAP4rev1 server for
// archivmail. It allows IMAP clients (Thunderbird, Outlook, etc.) to browse
// the mail archive without modifying it.
//
// This is NOT the IMAP importer (internal/imap/) which pulls mail from
// external servers. This package serves archived mail to clients.
package imapserver

import (
	"bufio"
	"context"
	"crypto/tls"
	"fmt"
	"log/slog"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/archivmail/config"
	"github.com/archivmail/internal/audit"
	"github.com/archivmail/internal/auth"
	"github.com/archivmail/internal/labelstore"
	"github.com/archivmail/internal/storage"
	"github.com/archivmail/internal/userstore"
	"github.com/archivmail/pkg/mailparser"
)

const (
	idleTimeout        = 30 * time.Minute
	maxConnsPerUser    = 5
	readBufferSize     = 8192
	maxLineLength      = 65536
)

// tenantIMAPModeGetter is satisfied by tenantstore.Store.
type tenantIMAPModeGetter interface {
	GetIMAPMode(ctx context.Context, tenantID int64) (string, error)
}

// Server is the embedded read-only IMAP archive server.
type Server struct {
	cfg         config.IMAPServerConfig
	mailStore   *storage.Store
	users       *userstore.Store
	labels      *labelstore.Store
	audit       *audit.Logger
	authMgr     *auth.Manager
	logger      *slog.Logger
	tenantStore tenantIMAPModeGetter
	listener    net.Listener

	mu        sync.Mutex
	running   bool
	wg        sync.WaitGroup
	done      chan struct{}

	// Per-user connection tracking
	connMu    sync.Mutex
	connCount map[string]*atomic.Int32
}

// New creates a new IMAP archive server.
func New(
	cfg config.IMAPServerConfig,
	mailStore *storage.Store,
	users *userstore.Store,
	labels *labelstore.Store,
	auditLog *audit.Logger,
	authMgr *auth.Manager,
	logger *slog.Logger,
	tenantStore tenantIMAPModeGetter,
) *Server {
	return &Server{
		cfg:         cfg,
		mailStore:   mailStore,
		users:       users,
		labels:      labels,
		audit:       auditLog,
		authMgr:     authMgr,
		logger:      logger,
		tenantStore: tenantStore,
		done:        make(chan struct{}),
		connCount:   make(map[string]*atomic.Int32),
	}
}

// Start launches the IMAP server in a background goroutine.
func (s *Server) Start() error {
	bind := s.cfg.Bind
	if bind == "" {
		if s.cfg.TLSCert != "" {
			bind = ":993"
		} else {
			bind = "127.0.0.1:1143"
		}
	}

	var ln net.Listener
	var err error
	if s.cfg.TLSCert != "" && s.cfg.TLSKey != "" {
		cert, err := tls.LoadX509KeyPair(s.cfg.TLSCert, s.cfg.TLSKey)
		if err != nil {
			return fmt.Errorf("imapserver: load TLS cert: %w", err)
		}
		tlsCfg := &tls.Config{
			Certificates: []tls.Certificate{cert},
			MinVersion:   tls.VersionTLS12,
		}
		ln, err = tls.Listen("tcp", bind, tlsCfg)
		if err != nil {
			return fmt.Errorf("imapserver: tls listen %s: %w", bind, err)
		}
		s.logger.Info("IMAP archive server TLS enabled", "addr", bind)
	} else {
		ln, err = net.Listen("tcp", bind)
		if err != nil {
			return fmt.Errorf("imapserver: listen %s: %w", bind, err)
		}
	}

	s.mu.Lock()
	s.listener = ln
	s.running = true
	s.mu.Unlock()

	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		s.logger.Info("IMAP archive server accepting connections", "addr", bind)
		for {
			conn, err := ln.Accept()
			if err != nil {
				select {
				case <-s.done:
					return
				default:
					s.logger.Error("imapserver: accept error", "err", err)
					continue
				}
			}
			s.wg.Add(1)
			go func() {
				defer s.wg.Done()
				s.handleConnection(conn)
			}()
		}
	}()

	return nil
}

// Stop gracefully shuts down the IMAP server.
func (s *Server) Stop() {
	s.mu.Lock()
	if !s.running {
		s.mu.Unlock()
		return
	}
	s.running = false
	s.mu.Unlock()

	close(s.done)
	if s.listener != nil {
		s.listener.Close()
	}
	s.wg.Wait()
}

// ── Per-user connection limiting ──────────────────────────────────────────

func (s *Server) acquireConn(username string) bool {
	s.connMu.Lock()
	defer s.connMu.Unlock()

	counter, ok := s.connCount[username]
	if !ok {
		counter = &atomic.Int32{}
		s.connCount[username] = counter
	}
	if counter.Load() >= int32(maxConnsPerUser) {
		return false
	}
	counter.Add(1)
	return true
}

func (s *Server) releaseConn(username string) {
	s.connMu.Lock()
	defer s.connMu.Unlock()

	if counter, ok := s.connCount[username]; ok {
		if counter.Add(-1) <= 0 {
			delete(s.connCount, username)
		}
	}
}

// ── Connection handler ────────────────────────────────────────────────────

func (s *Server) handleConnection(conn net.Conn) {
	defer conn.Close()

	remoteAddr := conn.RemoteAddr().String()
	s.logger.Debug("imapserver: new connection", "remote", remoteAddr)

	sess := &session{
		server:     s,
		conn:       conn,
		reader:     bufio.NewReaderSize(conn, readBufferSize),
		remoteAddr: remoteAddr,
		state:      stateNotAuth,
	}

	// Send greeting
	sess.writeResponse("* OK archivmail IMAP4rev1 Read-Only Archive Server ready")

	for {
		// Reset idle timeout
		conn.SetDeadline(time.Now().Add(idleTimeout))

		line, err := sess.readLine()
		if err != nil {
			if sess.username != "" {
				s.releaseConn(sess.username)
			}
			return
		}

		if len(line) == 0 {
			continue
		}

		sess.handleCommand(line)
		if sess.closed {
			if sess.username != "" {
				s.releaseConn(sess.username)
			}
			return
		}
	}
}

// ── IMAP session states ───────────────────────────────────────────────────

const (
	stateNotAuth    = 0
	stateAuth       = 1
	stateSelected   = 2
)

// session represents a single IMAP client connection.
type session struct {
	server     *Server
	conn       net.Conn
	reader     *bufio.Reader
	remoteAddr string

	state     int
	closed    bool
	username  string
	userID    int64
	userEmail string // for personal IMAP mode filtering
	tenantID  *int64

	// Selected mailbox state
	selectedMailbox string
	selectedMails   []mailEntry // ordered list of mails in the selected mailbox
}

// mailEntry represents a single mail in a mailbox listing.
type mailEntry struct {
	ID          string
	SeqNum      uint32
	UID         uint32
	From        string
	To          string
	Subject     string
	Date        time.Time
	Size        int64
	HasAttach   bool
}

func (sess *session) readLine() (string, error) {
	line, err := sess.reader.ReadString('\n')
	if err != nil {
		return "", err
	}
	return strings.TrimRight(line, "\r\n"), nil
}

func (sess *session) writeResponse(line string) {
	sess.conn.Write([]byte(line + "\r\n"))
}

func (sess *session) handleCommand(line string) {
	// Parse: TAG COMMAND [args...]
	parts := strings.SplitN(line, " ", 3)
	if len(parts) < 2 {
		sess.writeResponse("* BAD Invalid command")
		return
	}

	tag := parts[0]
	cmd := strings.ToUpper(parts[1])
	var args string
	if len(parts) > 2 {
		args = parts[2]
	}

	switch cmd {
	case "CAPABILITY":
		sess.cmdCapability(tag)
	case "NOOP":
		sess.writeResponse(tag + " OK NOOP completed")
	case "LOGOUT":
		sess.cmdLogout(tag)
	case "LOGIN":
		sess.cmdLogin(tag, args)
	case "LIST":
		sess.cmdList(tag, args)
	case "LSUB":
		// Treat LSUB same as LIST for simplicity
		sess.cmdList(tag, args)
	case "SELECT":
		sess.cmdSelect(tag, args)
	case "EXAMINE":
		// EXAMINE is SELECT but read-only (we are always read-only)
		sess.cmdSelect(tag, args)
	case "STATUS":
		sess.cmdStatus(tag, args)
	case "FETCH":
		sess.cmdFetch(tag, args)
	case "SEARCH":
		sess.cmdSearch(tag, args)
	case "UID":
		sess.cmdUID(tag, args)
	case "CLOSE":
		sess.cmdClose(tag)
	// Read-only enforcement: reject all mutating commands
	case "STORE", "DELETE", "COPY", "MOVE", "APPEND", "EXPUNGE", "CREATE", "RENAME":
		sess.writeResponse(tag + " NO [CANNOT] Read-only archive")
	case "SUBSCRIBE", "UNSUBSCRIBE":
		sess.writeResponse(tag + " NO [CANNOT] Read-only archive")
	case "IDLE":
		sess.cmdIdle(tag)
	default:
		sess.writeResponse(tag + " BAD Unknown command")
	}
}

// ── IMAP Commands ─────────────────────────────────────────────────────────

func (sess *session) cmdCapability(tag string) {
	sess.writeResponse("* CAPABILITY IMAP4rev1 IDLE")
	sess.writeResponse(tag + " OK CAPABILITY completed")
}

func (sess *session) cmdLogout(tag string) {
	sess.writeResponse("* BYE archivmail IMAP server signing off")
	sess.writeResponse(tag + " OK LOGOUT completed")
	sess.closed = true
}

func (sess *session) cmdLogin(tag string, args string) {
	if sess.state != stateNotAuth {
		sess.writeResponse(tag + " BAD Already authenticated")
		return
	}

	username, password := parseLoginArgs(args)
	if username == "" || password == "" {
		sess.writeResponse(tag + " BAD Missing credentials")
		return
	}

	// Authenticate via userstore (direct bcrypt check, bypasses TOTP for IMAP)
	user, err := sess.server.users.VerifyPassword(username, password)
	if err != nil {
		sess.server.logger.Warn("imapserver: login failed", "user", username, "remote", sess.remoteAddr)
		sess.server.audit.Log(audit.Entry{
			EventType: "imap_login_failed",
			Username:  username,
			IPAddress: extractIP(sess.remoteAddr),
			Success:   false,
			Detail:    "IMAP login failed",
		})
		sess.writeResponse(tag + " NO Authentication failed")
		return
	}

	// Check connection limit
	if !sess.server.acquireConn(username) {
		sess.writeResponse(tag + " NO Too many connections")
		return
	}

	sess.username = user.Username
	sess.userID = user.ID
	sess.userEmail = user.Email
	sess.tenantID = user.TenantID
	sess.state = stateAuth

	sess.server.logger.Info("imapserver: login success", "user", username, "remote", sess.remoteAddr)
	sess.server.audit.Log(audit.Entry{
		EventType: "imap_login",
		Username:  username,
		IPAddress: extractIP(sess.remoteAddr),
		Success:   true,
		Detail:    "IMAP login successful",
	})

	sess.writeResponse(tag + " OK LOGIN completed")
}

func (sess *session) cmdList(tag string, args string) {
	if sess.state < stateAuth {
		sess.writeResponse(tag + " NO Not authenticated")
		return
	}

	// Parse LIST arguments: reference mailbox-pattern
	ref, pattern := parseListArgs(args)
	_ = ref

	// Build mailbox list: INBOX + label-based sub-folders
	mailboxes := []string{"INBOX"}

	if sess.tenantID != nil {
		labels, err := sess.server.labels.GetLabelsForUser(
			context.Background(), sess.userID, *sess.tenantID,
		)
		if err == nil {
			for _, l := range labels {
				mailboxes = append(mailboxes, "INBOX/"+l.Name)
			}
		}
	}

	for _, mbox := range mailboxes {
		if matchMailbox(pattern, mbox) {
			attrs := ""
			if mbox == "INBOX" {
				attrs = `\Noinferiors`
			}
			if strings.HasPrefix(mbox, "INBOX/") {
				attrs = `\Noinferiors`
			}
			sess.writeResponse(fmt.Sprintf(`* LIST (%s) "/" "%s"`, attrs, mbox))
		}
	}

	sess.writeResponse(tag + " OK LIST completed")
}

func (sess *session) cmdSelect(tag string, args string) {
	if sess.state < stateAuth {
		sess.writeResponse(tag + " NO Not authenticated")
		return
	}

	mailbox := stripQuotes(strings.TrimSpace(args))
	if mailbox == "" {
		sess.writeResponse(tag + " NO No mailbox specified")
		return
	}

	// Load mails for the selected mailbox
	mails, err := sess.loadMailsForMailbox(mailbox)
	if err != nil {
		sess.server.logger.Error("imapserver: select failed", "mailbox", mailbox, "err", err)
		sess.writeResponse(tag + " NO Mailbox not found")
		return
	}

	sess.selectedMailbox = mailbox
	sess.selectedMails = mails
	sess.state = stateSelected

	count := len(mails)
	sess.writeResponse(fmt.Sprintf("* %d EXISTS", count))
	sess.writeResponse("* 0 RECENT")
	sess.writeResponse("* OK [UIDVALIDITY 1]")
	if count > 0 {
		sess.writeResponse(fmt.Sprintf("* OK [UIDNEXT %d]", mails[count-1].UID+1))
	} else {
		sess.writeResponse("* OK [UIDNEXT 1]")
	}
	sess.writeResponse("* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)")
	sess.writeResponse("* OK [PERMANENTFLAGS ()] Read-only archive")
	sess.writeResponse(tag + " OK [READ-ONLY] SELECT completed")
}

func (sess *session) cmdStatus(tag string, args string) {
	if sess.state < stateAuth {
		sess.writeResponse(tag + " NO Not authenticated")
		return
	}

	// Parse: "mailbox" (STATUS_ITEMS)
	mailbox, items := parseStatusArgs(args)
	if mailbox == "" {
		sess.writeResponse(tag + " NO Invalid arguments")
		return
	}

	mails, err := sess.loadMailsForMailbox(mailbox)
	if err != nil {
		sess.writeResponse(tag + " NO Mailbox not found")
		return
	}

	count := len(mails)
	var uidNext uint32 = 1
	if count > 0 {
		uidNext = mails[count-1].UID + 1
	}

	var parts []string
	for _, item := range items {
		switch strings.ToUpper(item) {
		case "MESSAGES":
			parts = append(parts, fmt.Sprintf("MESSAGES %d", count))
		case "RECENT":
			parts = append(parts, "RECENT 0")
		case "UIDNEXT":
			parts = append(parts, fmt.Sprintf("UIDNEXT %d", uidNext))
		case "UIDVALIDITY":
			parts = append(parts, "UIDVALIDITY 1")
		case "UNSEEN":
			parts = append(parts, fmt.Sprintf("UNSEEN %d", count)) // all "unseen" since we don't track flags
		}
	}

	sess.writeResponse(fmt.Sprintf(`* STATUS "%s" (%s)`, mailbox, strings.Join(parts, " ")))
	sess.writeResponse(tag + " OK STATUS completed")
}

func (sess *session) cmdFetch(tag string, args string) {
	if sess.state != stateSelected {
		sess.writeResponse(tag + " NO No mailbox selected")
		return
	}

	seqSet, items := parseFetchArgs(args)
	if seqSet == "" {
		sess.writeResponse(tag + " BAD Invalid arguments")
		return
	}

	sequences := parseSequenceSet(seqSet, uint32(len(sess.selectedMails)))
	fetchItems := parseFetchItems(items)

	for _, seqNum := range sequences {
		if seqNum < 1 || int(seqNum) > len(sess.selectedMails) {
			continue
		}
		entry := sess.selectedMails[seqNum-1]
		sess.fetchMail(seqNum, entry, fetchItems, false)
	}

	sess.writeResponse(tag + " OK FETCH completed")
}

func (sess *session) cmdSearch(tag string, args string) {
	if sess.state != stateSelected {
		sess.writeResponse(tag + " NO No mailbox selected")
		return
	}

	criteria := parseSearchCriteria(args)
	var results []uint32

	for _, entry := range sess.selectedMails {
		if matchesSearchCriteria(entry, criteria) {
			results = append(results, entry.SeqNum)
		}
	}

	var nums []string
	for _, n := range results {
		nums = append(nums, fmt.Sprintf("%d", n))
	}

	sess.writeResponse("* SEARCH " + strings.Join(nums, " "))
	sess.writeResponse(tag + " OK SEARCH completed")
}

func (sess *session) cmdUID(tag string, args string) {
	if sess.state != stateSelected {
		sess.writeResponse(tag + " NO No mailbox selected")
		return
	}

	parts := strings.SplitN(args, " ", 2)
	if len(parts) < 2 {
		sess.writeResponse(tag + " BAD Invalid UID command")
		return
	}

	subCmd := strings.ToUpper(parts[0])
	subArgs := parts[1]

	switch subCmd {
	case "FETCH":
		sess.cmdUIDFetch(tag, subArgs)
	case "SEARCH":
		sess.cmdUIDSearch(tag, subArgs)
	case "STORE", "COPY", "MOVE", "EXPUNGE":
		sess.writeResponse(tag + " NO [CANNOT] Read-only archive")
	default:
		sess.writeResponse(tag + " BAD Unknown UID subcommand")
	}
}

func (sess *session) cmdUIDFetch(tag string, args string) {
	uidSet, items := parseFetchArgs(args)
	if uidSet == "" {
		sess.writeResponse(tag + " BAD Invalid arguments")
		return
	}

	// Build a UID-to-entry map
	maxUID := uint32(0)
	for _, e := range sess.selectedMails {
		if e.UID > maxUID {
			maxUID = e.UID
		}
	}

	uids := parseSequenceSet(uidSet, maxUID)
	uidMap := make(map[uint32]*mailEntry, len(sess.selectedMails))
	for i := range sess.selectedMails {
		uidMap[sess.selectedMails[i].UID] = &sess.selectedMails[i]
	}

	fetchItems := parseFetchItems(items)

	for _, uid := range uids {
		entry, ok := uidMap[uid]
		if !ok {
			continue
		}
		sess.fetchMail(entry.SeqNum, *entry, fetchItems, true)
	}

	sess.writeResponse(tag + " OK UID FETCH completed")
}

func (sess *session) cmdUIDSearch(tag string, args string) {
	criteria := parseSearchCriteria(args)
	var results []uint32

	for _, entry := range sess.selectedMails {
		if matchesSearchCriteria(entry, criteria) {
			results = append(results, entry.UID)
		}
	}

	var nums []string
	for _, n := range results {
		nums = append(nums, fmt.Sprintf("%d", n))
	}

	sess.writeResponse("* SEARCH " + strings.Join(nums, " "))
	sess.writeResponse(tag + " OK UID SEARCH completed")
}

func (sess *session) cmdClose(tag string) {
	if sess.state == stateSelected {
		sess.selectedMailbox = ""
		sess.selectedMails = nil
		sess.state = stateAuth
	}
	sess.writeResponse(tag + " OK CLOSE completed")
}

func (sess *session) cmdIdle(tag string) {
	if sess.state < stateAuth {
		sess.writeResponse(tag + " NO Not authenticated")
		return
	}

	sess.writeResponse("+ idling")

	// Wait for DONE from client (with idle timeout)
	sess.conn.SetDeadline(time.Now().Add(idleTimeout))
	line, err := sess.readLine()
	if err != nil {
		sess.closed = true
		return
	}

	if strings.ToUpper(strings.TrimSpace(line)) == "DONE" {
		sess.writeResponse(tag + " OK IDLE terminated")
	} else {
		sess.writeResponse(tag + " BAD Expected DONE")
	}
}

// ── Mail loading ──────────────────────────────────────────────────────────

func (sess *session) loadMailsForMailbox(mailbox string) ([]mailEntry, error) {
	ctx := context.Background()

	// Determine IMAP mode for this tenant
	mode := "personal"
	if sess.tenantID != nil && sess.server.tenantStore != nil {
		if m, err := sess.server.tenantStore.GetIMAPMode(ctx, *sess.tenantID); err == nil {
			mode = m
		}
	}

	// Load mails with stable UIDs depending on mode
	var rawMails []storage.MailWithUID
	var err error
	if mode == "shared" {
		rawMails, err = sess.server.mailStore.GetMailsWithUID(ctx, sess.tenantID)
	} else {
		rawMails, err = sess.server.mailStore.GetMailsByRecipient(ctx, sess.tenantID, sess.userEmail)
	}
	if err != nil {
		return nil, fmt.Errorf("load mails: %w", err)
	}

	// Label filter setup
	var filterLabelID *int64
	if strings.HasPrefix(mailbox, "INBOX/") {
		labelName := strings.TrimPrefix(mailbox, "INBOX/")
		if sess.tenantID != nil {
			labels, err := sess.server.labels.GetLabelsForUser(ctx, sess.userID, *sess.tenantID)
			if err == nil {
				for _, l := range labels {
					if l.Name == labelName {
						lid := l.ID
						filterLabelID = &lid
						break
					}
				}
			}
		}
		if filterLabelID == nil {
			return nil, fmt.Errorf("label not found: %s", mailbox)
		}
	} else if mailbox != "INBOX" {
		return nil, fmt.Errorf("unknown mailbox: %s", mailbox)
	}

	var labelEmailIDs map[string]bool
	if filterLabelID != nil {
		emailIDs, err := sess.server.labels.GetEmailIDsByLabel(ctx, *filterLabelID)
		if err != nil {
			return nil, fmt.Errorf("load label emails: %w", err)
		}
		labelEmailIDs = make(map[string]bool, len(emailIDs))
		for _, eid := range emailIDs {
			labelEmailIDs[eid] = true
		}
	}

	var entries []mailEntry
	var seqNum uint32 = 1
	for _, m := range rawMails {
		if labelEmailIDs != nil && !labelEmailIDs[m.ID] {
			continue
		}
		uid := uint32(m.UID)
		if uid == 0 {
			uid = seqNum // fallback if no UID in DB yet
		}
		entries = append(entries, mailEntry{
			ID:     m.ID,
			SeqNum: seqNum,
			UID:    uid,
		})
		seqNum++
	}

	return entries, nil
}

// fetchMail handles FETCH for a single mail entry.
func (sess *session) fetchMail(seqNum uint32, entry mailEntry, items []string, includeUID bool) {
	var dataParts []string

	if includeUID {
		dataParts = append(dataParts, fmt.Sprintf("UID %d", entry.UID))
	}

	var rawBody []byte
	var parsed *mailparser.ParsedMail
	needBody := false
	needEnvelope := false
	needFlags := false
	needSize := false
	needInternalDate := false

	for _, item := range items {
		upper := strings.ToUpper(item)
		switch {
		case upper == "FLAGS":
			needFlags = true
		case upper == "UID":
			if !includeUID {
				dataParts = append(dataParts, fmt.Sprintf("UID %d", entry.UID))
			}
		case upper == "RFC822.SIZE":
			needSize = true
			needBody = true
		case upper == "INTERNALDATE":
			needInternalDate = true
			needBody = true
		case upper == "ENVELOPE":
			needEnvelope = true
			needBody = true
		case strings.HasPrefix(upper, "BODY") || upper == "RFC822" || upper == "RFC822.HEADER" || upper == "RFC822.TEXT":
			needBody = true
		}
	}

	// Load the actual mail if needed
	if needBody || needEnvelope || needSize || needInternalDate {
		raw, err := sess.server.mailStore.Load(entry.ID)
		if err != nil {
			sess.server.logger.Warn("imapserver: fetch load failed", "id", entry.ID, "err", err)
			return
		}
		rawBody = raw
		pm, err := mailparser.Parse(raw)
		if err == nil {
			parsed = pm
		}
	}

	for _, item := range items {
		upper := strings.ToUpper(item)
		switch {
		case upper == "FLAGS":
			// intentionally handled below after the loop
		case upper == "UID":
			// Already added above
		case upper == "RFC822.SIZE":
			dataParts = append(dataParts, fmt.Sprintf("RFC822.SIZE %d", len(rawBody)))
		case upper == "INTERNALDATE":
			date := time.Now()
			if parsed != nil && !parsed.Date.IsZero() {
				date = parsed.Date
			}
			dataParts = append(dataParts, fmt.Sprintf(`INTERNALDATE "%s"`, date.Format("02-Jan-2006 15:04:05 -0700")))
		case upper == "ENVELOPE":
			dataParts = append(dataParts, "ENVELOPE "+buildEnvelope(parsed))
		case upper == "BODY[]" || upper == "BODY.PEEK[]" || upper == "RFC822":
			dataParts = append(dataParts, fmt.Sprintf("BODY[] {%d}\r\n%s", len(rawBody), string(rawBody)))
		case upper == "BODY[HEADER]" || upper == "BODY.PEEK[HEADER]" || upper == "RFC822.HEADER":
			header := extractHeader(rawBody)
			dataParts = append(dataParts, fmt.Sprintf("BODY[HEADER] {%d}\r\n%s", len(header), header))
		case upper == "BODY[TEXT]" || upper == "BODY.PEEK[TEXT]" || upper == "RFC822.TEXT":
			body := extractBody(rawBody)
			dataParts = append(dataParts, fmt.Sprintf("BODY[TEXT] {%d}\r\n%s", len(body), body))
		case upper == "BODYSTRUCTURE":
			dataParts = append(dataParts, "BODYSTRUCTURE "+buildBodyStructure(parsed, rawBody))
		}
	}

	if needFlags {
		dataParts = append([]string{"FLAGS (\\Seen)"}, dataParts...)
	}

	sess.writeResponse(fmt.Sprintf("* %d FETCH (%s)", seqNum, strings.Join(dataParts, " ")))
}

// ── Parsing helpers ───────────────────────────────────────────────────────

// parseLoginArgs extracts username and password from LOGIN args.
// Handles both quoted and unquoted forms.
func parseLoginArgs(args string) (string, string) {
	args = strings.TrimSpace(args)
	if args == "" {
		return "", ""
	}

	var username, password string
	if args[0] == '"' {
		end := strings.Index(args[1:], "\"")
		if end < 0 {
			return "", ""
		}
		username = args[1 : end+1]
		rest := strings.TrimSpace(args[end+2:])
		password = stripQuotes(rest)
	} else {
		parts := strings.SplitN(args, " ", 2)
		if len(parts) < 2 {
			return "", ""
		}
		username = parts[0]
		password = stripQuotes(strings.TrimSpace(parts[1]))
	}
	return username, password
}

func parseListArgs(args string) (string, string) {
	args = strings.TrimSpace(args)
	// LIST reference pattern
	var ref, pattern string

	parts := splitIMAPArgs(args)
	if len(parts) >= 2 {
		ref = stripQuotes(parts[0])
		pattern = stripQuotes(parts[1])
	} else if len(parts) == 1 {
		pattern = stripQuotes(parts[0])
	}

	// Default pattern
	if pattern == "" {
		pattern = "*"
	}
	return ref, pattern
}

func parseStatusArgs(args string) (string, []string) {
	args = strings.TrimSpace(args)
	// STATUS "mailbox" (MESSAGES RECENT UIDNEXT UIDVALIDITY UNSEEN)
	var mailbox string
	var rest string

	if args == "" {
		return "", nil
	}

	if args[0] == '"' {
		end := strings.Index(args[1:], "\"")
		if end < 0 {
			return "", nil
		}
		mailbox = args[1 : end+1]
		rest = strings.TrimSpace(args[end+2:])
	} else {
		parts := strings.SplitN(args, " ", 2)
		mailbox = parts[0]
		if len(parts) > 1 {
			rest = strings.TrimSpace(parts[1])
		}
	}

	// Parse (ITEM1 ITEM2 ...)
	rest = strings.TrimPrefix(rest, "(")
	rest = strings.TrimSuffix(rest, ")")
	items := strings.Fields(rest)

	return mailbox, items
}

func parseFetchArgs(args string) (string, string) {
	args = strings.TrimSpace(args)
	parts := strings.SplitN(args, " ", 2)
	if len(parts) < 2 {
		return parts[0], ""
	}
	return parts[0], parts[1]
}

func parseFetchItems(items string) []string {
	items = strings.TrimSpace(items)
	if items == "" {
		return nil
	}

	// Handle macro shortcuts
	upper := strings.ToUpper(items)
	switch upper {
	case "ALL":
		return []string{"FLAGS", "INTERNALDATE", "RFC822.SIZE", "ENVELOPE"}
	case "FAST":
		return []string{"FLAGS", "INTERNALDATE", "RFC822.SIZE"}
	case "FULL":
		return []string{"FLAGS", "INTERNALDATE", "RFC822.SIZE", "ENVELOPE", "BODY"}
	}

	// Parse parenthesized list
	items = strings.TrimPrefix(items, "(")
	items = strings.TrimSuffix(items, ")")

	var result []string
	current := ""
	bracketDepth := 0
	for _, ch := range items {
		if ch == '[' {
			bracketDepth++
			current += string(ch)
		} else if ch == ']' {
			bracketDepth--
			current += string(ch)
		} else if ch == ' ' && bracketDepth == 0 {
			if current != "" {
				result = append(result, current)
				current = ""
			}
		} else {
			current += string(ch)
		}
	}
	if current != "" {
		result = append(result, current)
	}

	return result
}

// parseSequenceSet parses an IMAP sequence set (e.g. "1:*", "1,3:5", "1")
// and returns the expanded list of numbers.
func parseSequenceSet(set string, maxVal uint32) []uint32 {
	if maxVal == 0 {
		return nil
	}

	var result []uint32
	seen := make(map[uint32]bool)

	for _, part := range strings.Split(set, ",") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}

		if strings.Contains(part, ":") {
			bounds := strings.SplitN(part, ":", 2)
			start := parseSeqNum(bounds[0], maxVal)
			end := parseSeqNum(bounds[1], maxVal)
			if start > end {
				start, end = end, start
			}
			for i := start; i <= end; i++ {
				if !seen[i] {
					result = append(result, i)
					seen[i] = true
				}
			}
		} else {
			n := parseSeqNum(part, maxVal)
			if !seen[n] {
				result = append(result, n)
				seen[n] = true
			}
		}
	}
	return result
}

func parseSeqNum(s string, maxVal uint32) uint32 {
	s = strings.TrimSpace(s)
	if s == "*" {
		return maxVal
	}
	var n uint32
	for _, ch := range s {
		if ch >= '0' && ch <= '9' {
			n = n*10 + uint32(ch-'0')
		}
	}
	if n == 0 {
		n = 1
	}
	if n > maxVal {
		n = maxVal
	}
	return n
}

// ── Search criteria ───────────────────────────────────────────────────────

type searchCriteria struct {
	all     bool
	from    string
	to      string
	subject string
	since   *time.Time
	before  *time.Time
}

func parseSearchCriteria(args string) searchCriteria {
	c := searchCriteria{}
	tokens := strings.Fields(args)

	if len(tokens) == 0 {
		c.all = true
		return c
	}

	i := 0
	for i < len(tokens) {
		key := strings.ToUpper(tokens[i])
		switch key {
		case "ALL":
			c.all = true
			i++
		case "FROM":
			if i+1 < len(tokens) {
				c.from = stripQuotes(tokens[i+1])
				i += 2
			} else {
				i++
			}
		case "TO":
			if i+1 < len(tokens) {
				c.to = stripQuotes(tokens[i+1])
				i += 2
			} else {
				i++
			}
		case "SUBJECT":
			if i+1 < len(tokens) {
				c.subject = stripQuotes(tokens[i+1])
				i += 2
			} else {
				i++
			}
		case "SINCE":
			if i+1 < len(tokens) {
				t, err := time.Parse("2-Jan-2006", stripQuotes(tokens[i+1]))
				if err == nil {
					c.since = &t
				}
				i += 2
			} else {
				i++
			}
		case "BEFORE":
			if i+1 < len(tokens) {
				t, err := time.Parse("2-Jan-2006", stripQuotes(tokens[i+1]))
				if err == nil {
					c.before = &t
				}
				i += 2
			} else {
				i++
			}
		default:
			i++
		}
	}

	// If no specific criteria set, treat as ALL
	if c.from == "" && c.to == "" && c.subject == "" && c.since == nil && c.before == nil {
		c.all = true
	}

	return c
}

func matchesSearchCriteria(entry mailEntry, c searchCriteria) bool {
	if c.all && c.from == "" && c.to == "" && c.subject == "" && c.since == nil && c.before == nil {
		return true
	}

	if c.from != "" && !containsCI(entry.From, c.from) {
		return false
	}
	if c.to != "" && !containsCI(entry.To, c.to) {
		return false
	}
	if c.subject != "" && !containsCI(entry.Subject, c.subject) {
		return false
	}
	if c.since != nil && entry.Date.Before(*c.since) {
		return false
	}
	if c.before != nil && !entry.Date.Before(*c.before) {
		return false
	}

	return true
}

// ── Helper functions ──────────────────────────────────────────────────────

func stripQuotes(s string) string {
	s = strings.TrimSpace(s)
	if len(s) >= 2 && s[0] == '"' && s[len(s)-1] == '"' {
		return s[1 : len(s)-1]
	}
	return s
}

func splitIMAPArgs(s string) []string {
	var parts []string
	current := ""
	inQuote := false
	for _, ch := range s {
		if ch == '"' {
			inQuote = !inQuote
			current += string(ch)
		} else if ch == ' ' && !inQuote {
			if current != "" {
				parts = append(parts, current)
				current = ""
			}
		} else {
			current += string(ch)
		}
	}
	if current != "" {
		parts = append(parts, current)
	}
	return parts
}

func containsCI(haystack, needle string) bool {
	return strings.Contains(strings.ToLower(haystack), strings.ToLower(needle))
}

func matchMailbox(pattern, mailbox string) bool {
	if pattern == "*" || pattern == "%" {
		return true
	}
	// Simple glob: * matches everything, % matches one level
	pattern = strings.ReplaceAll(pattern, "*", ".*")
	pattern = strings.ReplaceAll(pattern, "%", "[^/]*")
	return containsCI(mailbox, strings.ReplaceAll(pattern, ".*", ""))
}

func extractIP(addr string) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return addr
	}
	return host
}

func extractHeader(raw []byte) string {
	s := string(raw)
	idx := strings.Index(s, "\r\n\r\n")
	if idx >= 0 {
		return s[:idx+4]
	}
	idx = strings.Index(s, "\n\n")
	if idx >= 0 {
		return s[:idx+2]
	}
	return s
}

func extractBody(raw []byte) string {
	s := string(raw)
	idx := strings.Index(s, "\r\n\r\n")
	if idx >= 0 {
		return s[idx+4:]
	}
	idx = strings.Index(s, "\n\n")
	if idx >= 0 {
		return s[idx+2:]
	}
	return ""
}

func buildEnvelope(pm *mailparser.ParsedMail) string {
	if pm == nil {
		return "NIL"
	}
	date := ""
	if !pm.Date.IsZero() {
		date = pm.Date.Format(time.RFC822Z)
	}
	subj := quoteString(pm.Subject)
	from := quoteString(pm.From)
	to := quoteString(strings.Join(pm.To, ", "))
	msgID := quoteString(pm.MessageID)

	return fmt.Sprintf("(%s %s ((%s NIL NIL NIL)) NIL NIL ((%s NIL NIL NIL)) NIL NIL NIL %s)",
		quoteString(date), subj, from, to, msgID)
}

func buildBodyStructure(pm *mailparser.ParsedMail, raw []byte) string {
	if pm == nil {
		return fmt.Sprintf(`("TEXT" "PLAIN" NIL NIL NIL "7BIT" %d NIL NIL NIL NIL)`, len(raw))
	}
	size := len(raw)
	if pm.TextBody != "" {
		return fmt.Sprintf(`("TEXT" "PLAIN" ("CHARSET" "UTF-8") NIL NIL "8BIT" %d NIL NIL NIL NIL)`, size)
	}
	return fmt.Sprintf(`("TEXT" "PLAIN" NIL NIL NIL "7BIT" %d NIL NIL NIL NIL)`, size)
}

func quoteString(s string) string {
	if s == "" {
		return "NIL"
	}
	s = strings.ReplaceAll(s, "\\", "\\\\")
	s = strings.ReplaceAll(s, "\"", "\\\"")
	return "\"" + s + "\""
}