openmail/internal/core/message.go

package core

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"regexp"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/mjl-/mox/message"
)

type Message struct {
	ID            string    `json:"id"`
	InboxID       string    `json:"inbox_id"`
	ThreadID      *string   `json:"thread_id,omitempty"`
	MessageIDHdr  *string   `json:"message_id,omitempty"`
	InReplyTo     *string   `json:"in_reply_to,omitempty"`
	References    []string  `json:"references"`
	FromAddr      *string   `json:"from,omitempty"`
	ToAddrs       []string  `json:"to"`
	Cc            []string  `json:"cc"`
	Bcc           []string  `json:"bcc"`
	Subject       *string   `json:"subject,omitempty"`
	Preview       *string   `json:"preview,omitempty"`
	Text          *string   `json:"text,omitempty"`
	HTML          *string   `json:"html,omitempty"`
	ExtractedText *string   `json:"extracted_text,omitempty"`
	Labels        []string  `json:"labels"`
	SizeBytes     *int64    `json:"size_bytes,omitempty"`
	CreatedAt     time.Time `json:"created_at"`
}

const messageCols = `id::text, inbox_id::text, thread_id::text, message_id_hdr, in_reply_to, ` +
	`"references", from_addr, to_addrs, cc, bcc, subject, preview, text, html, extracted_text, ` +
	`labels, size_bytes, created_at`

func scanMessage(row pgx.Row) (Message, error) {
	var m Message
	err := row.Scan(&m.ID, &m.InboxID, &m.ThreadID, &m.MessageIDHdr, &m.InReplyTo,
		&m.References, &m.FromAddr, &m.ToAddrs, &m.Cc, &m.Bcc, &m.Subject, &m.Preview,
		&m.Text, &m.HTML, &m.ExtractedText, &m.Labels, &m.SizeBytes, &m.CreatedAt)
	return m, err
}

// GetMessage is scoped to the inbox: a message id that belongs to another inbox
// returns ErrNotFound, preventing cross-inbox access.
func (s *Service) GetMessage(ctx context.Context, inboxID, id string) (Message, error) {
	m, err := scanMessage(s.pool.QueryRow(ctx,
		`SELECT `+messageCols+` FROM messages WHERE id = $1 AND inbox_id = $2`, id, inboxID))
	if errors.Is(err, pgx.ErrNoRows) {
		return Message{}, ErrNotFound
	}
	return m, err
}

func (s *Service) ListMessages(ctx context.Context, inboxID string, limit int) ([]Message, error) {
	if limit <= 0 || limit > 200 {
		limit = 50
	}
	rows, err := s.pool.Query(ctx,
		`SELECT `+messageCols+` FROM messages WHERE inbox_id = $1 ORDER BY created_at DESC LIMIT $2`,
		inboxID, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return collectMessages(rows)
}

func (s *Service) GetThreadMessages(ctx context.Context, inboxID, threadID string) ([]Message, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT `+messageCols+` FROM messages WHERE thread_id = $1 AND inbox_id = $2 ORDER BY created_at ASC`,
		threadID, inboxID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return collectMessages(rows)
}

func collectMessages(rows pgx.Rows) ([]Message, error) {
	out := []Message{}
	for rows.Next() {
		m, err := scanMessage(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

// Deliver implements mail.InboundSink: resolve the inbox by address and ingest.
func (s *Service) Deliver(ctx context.Context, inboxAddr string, raw []byte) error {
	ib, err := s.GetInboxByAddress(ctx, inboxAddr)
	if err != nil {
		return err
	}
	_, err = s.IngestRaw(ctx, ib.ID, raw)
	return err
}

// IngestRaw parses a raw RFC 5322 message with mox, resolves its thread, and
// stores it — the single inbound path shared by the ingest API and every
// MailBackend. Runs in one transaction.
func (s *Service) IngestRaw(ctx context.Context, inboxID string, raw []byte) (Message, error) {
	pr, err := parseRaw(raw)
	if err != nil {
		return Message{}, err
	}
	extracted := stripQuotes(pr.text)
	preview := makePreview(extracted, pr.text)
	headersJSON, _ := json.Marshal(pr.headers)

	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return Message{}, err
	}
	defer tx.Rollback(ctx)

	threadID, err := resolveThreadTx(ctx, tx, inboxID, pr)
	if err != nil {
		return Message{}, err
	}

	var msgID string
	err = tx.QueryRow(ctx,
		`INSERT INTO messages
		   (inbox_id, thread_id, message_id_hdr, in_reply_to, "references", from_addr,
		    to_addrs, cc, bcc, subject, preview, text, html, extracted_text,
		    size_bytes, headers, ts)
		 VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,
		         to_tsvector('english', coalesce($10,'') || ' ' || coalesce($12,'')))
		 RETURNING id::text`,
		inboxID, threadID, nullEmpty(pr.messageID), nullEmpty(pr.inReplyTo), pr.references,
		nullEmpty(pr.from), pr.to, pr.cc, pr.bcc, nullEmpty(pr.subject), nullEmpty(preview),
		nullEmpty(pr.text), nullEmpty(pr.html), nullEmpty(extracted),
		int64(len(raw)), headersJSON,
	).Scan(&msgID)
	if err != nil {
		return Message{}, err
	}

	if _, err := tx.Exec(ctx,
		`UPDATE threads
		    SET message_count = message_count + 1,
		        last_message_id = $1,
		        subject = COALESCE(subject, $2),
		        updated_at = now()
		  WHERE id = $3`,
		msgID, nullEmpty(pr.subject), threadID); err != nil {
		return Message{}, err
	}

	evt, _ := json.Marshal(map[string]string{"message_id": msgID, "thread_id": threadID})
	if _, err := tx.Exec(ctx,
		`INSERT INTO events (inbox_id, type, payload) VALUES ($1, 'message.received', $2)`,
		inboxID, evt); err != nil {
		return Message{}, err
	}

	if err := tx.Commit(ctx); err != nil {
		return Message{}, err
	}
	return s.GetMessage(ctx, inboxID, msgID)
}

// resolveThreadTx finds the thread for a message via In-Reply-To/References
// (the correct, false-merge-safe mechanism), creating a new thread otherwise.
// Subject-based fallback is intentionally deferred (see ARCHITECTURE.md §9).
func resolveThreadTx(ctx context.Context, tx pgx.Tx, inboxID string, pr parsed) (string, error) {
	cand := make([]string, 0, len(pr.references)+1)
	if pr.inReplyTo != "" {
		cand = append(cand, pr.inReplyTo)
	}
	cand = append(cand, pr.references...)

	if len(cand) > 0 {
		var tid string
		err := tx.QueryRow(ctx,
			`SELECT thread_id::text FROM messages
			  WHERE inbox_id = $1 AND message_id_hdr = ANY($2) AND thread_id IS NOT NULL
			  ORDER BY created_at DESC LIMIT 1`,
			inboxID, cand).Scan(&tid)
		if err == nil {
			return tid, nil
		}
		if !errors.Is(err, pgx.ErrNoRows) {
			return "", err
		}
	}

	var tid string
	err := tx.QueryRow(ctx,
		`INSERT INTO threads (inbox_id, subject) VALUES ($1, $2) RETURNING id::text`,
		inboxID, nullEmpty(pr.subject)).Scan(&tid)
	return tid, err
}

// --- parsing (mox/message) ---

// parsed is the normalized result of parsing a raw message.
type parsed struct {
	messageID  string
	inReplyTo  string
	references []string
	from       string
	to         []string
	cc         []string
	bcc        []string
	subject    string
	text       string
	html       string
	headers    map[string][]string
}

var discardLog = slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError}))

func parseRaw(raw []byte) (parsed, error) {
	var pr parsed
	// EnsurePart always returns a usable Part — building an octet-stream fallback
	// even when parsing hits a recoverable defect (bare CR/LF, bad Content-Type,
	// missing boundary, truncated DSN). That tolerance for messy real-world mail
	// is precisely why mox was chosen, so we proceed with the returned part and
	// do NOT treat the recoverable error as fatal.
	p, _ := message.EnsurePart(discardLog, false, bytes.NewReader(raw), int64(len(raw)))
	if p.Envelope != nil {
		e := p.Envelope
		pr.subject = e.Subject
		pr.messageID = e.MessageID
		pr.inReplyTo = e.InReplyTo
		pr.from = firstAddr(e.From)
		pr.to = addrList(e.To)
		pr.cc = addrList(e.CC)
		pr.bcc = addrList(e.BCC)
	}
	if hdr, herr := p.Header(); herr == nil {
		pr.references = strings.Fields(hdr.Get("References"))
		pr.headers = hdr
	}
	if pr.references == nil {
		pr.references = []string{}
	}
	pr.text, pr.html = extractBodies(&p)
	return pr, nil
}

// extractBodies walks the MIME tree and returns the first text/plain and
// text/html leaf bodies (coerced to valid UTF-8). It descends into embedded
// messages and skips attachment parts.
func extractBodies(p *message.Part) (text, html string) {
	// Embedded message (message/rfc822 or message/global): the sub-message lives
	// under p.Message, not p.Parts. Wire its reader, then recurse — otherwise
	// forwarded mail and DSN/bounce bodies are lost.
	if p.Message != nil {
		if err := p.SetMessageReaderAt(); err == nil {
			return extractBodies(p.Message)
		}
		return "", ""
	}
	if len(p.Parts) == 0 {
		if isAttachment(p) {
			return "", "" // an attachment is not the message body
		}
		body := readBody(p)
		switch {
		case p.MediaType == "TEXT" && p.MediaSubType == "HTML":
			return "", body
		case p.MediaType == "TEXT" || p.MediaType == "":
			return body, "" // PLAIN, or absent content-type → treat as plain
		default:
			return "", ""
		}
	}
	for i := range p.Parts {
		t, h := extractBodies(&p.Parts[i])
		if text == "" {
			text = t
		}
		if html == "" {
			html = h
		}
	}
	return text, html
}

// isAttachment reports whether a part is declared as an attachment (so it is not
// treated as the message body). Content-Disposition carries params, so we match
// the leading token.
func isAttachment(p *message.Part) bool {
	if p.ContentDisposition == nil {
		return false
	}
	return strings.HasPrefix(strings.ToLower(strings.TrimSpace(*p.ContentDisposition)), "attachment")
}

const maxBodyBytes = 2 << 20 // 2 MiB cap per body part for milestone 1

func readBody(p *message.Part) string {
	rd := p.ReaderUTF8OrBinary()
	if rd == nil {
		return ""
	}
	var b strings.Builder
	_, _ = io.Copy(&b, io.LimitReader(rd, maxBodyBytes))
	// Bodies may be non-UTF-8 (mox returns raw bytes for unknown/empty charsets)
	// and LimitReader can cut mid-rune; Postgres text/tsvector reject invalid
	// UTF-8 and would roll back the whole ingest. Coerce to valid UTF-8.
	return strings.ToValidUTF8(b.String(), "�")
}

// firstAddr returns the first address that has both a localpart and a host. mox
// appends empty-User/Host entries for addresses it cannot parse; emitting "@"
// for those would be wrong, so we skip them.
func firstAddr(as []message.Address) string {
	for _, a := range as {
		if a.User != "" && a.Host != "" {
			return a.User + "@" + a.Host
		}
	}
	return ""
}

func addrList(as []message.Address) []string {
	out := make([]string, 0, len(as))
	for _, a := range as {
		if a.User == "" || a.Host == "" {
			continue
		}
		out = append(out, a.User+"@"+a.Host)
	}
	return out
}

// --- text helpers ---

var onWroteRe = regexp.MustCompile(`(?i)^on .+wrote:$`)

// stripQuotes removes quoted history so an agent reads only the new content.
// Milestone-1 heuristic (talon-style port deferred): cut at the first quoted
// block or "On … wrote:" attribution line.
func stripQuotes(text string) string {
	if text == "" {
		return ""
	}
	lines := strings.Split(text, "\n")
	out := make([]string, 0, len(lines))
	for _, ln := range lines {
		t := strings.TrimSpace(ln)
		if strings.HasPrefix(t, ">") || onWroteRe.MatchString(t) {
			break
		}
		out = append(out, ln)
	}
	return strings.TrimSpace(strings.Join(out, "\n"))
}

func makePreview(extracted, full string) string {
	src := extracted
	if src == "" {
		src = full
	}
	src = strings.Join(strings.Fields(src), " ")
	const max = 200
	if len([]rune(src)) > max {
		src = string([]rune(src)[:max])
	}
	return src
}

// nullEmpty maps "" to a SQL NULL so optional text columns stay null, not blank.
func nullEmpty(s string) any {
	if s == "" {
		return nil
	}
	return s
}