MailUI4Agents/gateway/internal/repo/attachments.go

package repo

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/agentmail/gateway/internal/db"
	"github.com/agentmail/gateway/internal/models"
	"github.com/google/uuid"
)

// ---------- 附件 ----------
//
// 元数据在库、内容在磁盘（internal/blob）。两者的一致性由调用顺序保证：
// 先落盘再入库 —— 反过来会出现「库里有记录但文件不存在」的下载 500。
// 落盘成功但入库失败时最多留下一个无引用的文件，由 GC 回收，不影响正确性。

var (
	// ErrAttachmentNotFound 附件不存在
	ErrAttachmentNotFound = errors.New("attachment not found")
	// ErrAttachmentNotOwned 附件不属于该上传者
	ErrAttachmentNotOwned = errors.New("attachment not owned by uploader")
	// ErrAttachmentAlreadyAttached 附件已挂到别的邮件上
	ErrAttachmentAlreadyAttached = errors.New("attachment already attached")
)

const attachmentCols = `attachment_id, mail_id, uploader, filename, content_type, size_bytes, sha256, created_at`

func scanAttachment(sc interface{ Scan(...any) error }) (*models.Attachment, error) {
	var a models.Attachment
	if err := sc.Scan(&a.ID, &a.MailID, &a.Uploader, &a.Filename,
		&a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt); err != nil {
		return nil, err
	}
	return &a, nil
}

// CreateAttachment 登记一条待挂载的附件（mail_id 为空）。
func CreateAttachment(ctx context.Context, uploader, filename, contentType string, size int64, sum string) (*models.Attachment, error) {
	a := &models.Attachment{
		Uploader:    uploader,
		Filename:    filename,
		ContentType: contentType,
		SizeBytes:   size,
		SHA256:      sum,
	}
	err := db.DB.QueryRowContext(ctx, `
		INSERT INTO attachments (uploader, filename, content_type, size_bytes, sha256)
		VALUES ($1, $2, $3, $4, $5)
		RETURNING attachment_id, created_at
	`, uploader, filename, contentType, size, sum).Scan(&a.ID, &a.CreatedAt)
	if err != nil {
		return nil, err
	}
	return a, nil
}

// GetAttachment 读取一条附件元数据。
func GetAttachment(ctx context.Context, id uuid.UUID) (*models.Attachment, error) {
	a, err := scanAttachment(db.DB.QueryRowContext(ctx,
		`SELECT `+attachmentCols+` FROM attachments WHERE attachment_id = $1`, id))
	if errors.Is(err, sql.ErrNoRows) {
		return nil, ErrAttachmentNotFound
	}
	return a, err
}

// ListAttachmentsFor 列出某封邮件的附件。
func ListAttachmentsFor(ctx context.Context, mailID uuid.UUID) ([]models.Attachment, error) {
	rows, err := db.DB.QueryContext(ctx,
		`SELECT `+attachmentCols+` FROM attachments WHERE mail_id = $1 ORDER BY created_at`, mailID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := []models.Attachment{}
	for rows.Next() {
		a, err := scanAttachment(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, *a)
	}
	return out, rows.Err()
}

// AttachToMail 把一批待挂载附件绑到某封邮件上。
//
// 每条都要求：存在、属于该上传者、且尚未挂载。
// 用 WHERE mail_id IS NULL AND uploader = ? 一条 UPDATE 完成判断与写入，
// 避免「先查后改」在并发下把同一个附件挂到两封邮件上。
func AttachToMail(ctx context.Context, mailID uuid.UUID, ids []uuid.UUID, uploader string) error {
	for _, id := range ids {
		tag, err := db.DB.ExecContext(ctx, `
			UPDATE attachments SET mail_id = $1
			WHERE attachment_id = $2 AND uploader = $3 AND mail_id IS NULL
		`, mailID, id, uploader)
		if err != nil {
			return err
		}
		if n, _ := tag.RowsAffected(); n > 0 {
			continue
		}

		// 没改到：查明原因，给调用方一个能照着修的错误
		a, gErr := GetAttachment(ctx, id)
		if gErr != nil {
			return gErr
		}
		if a.Uploader != uploader {
			return ErrAttachmentNotOwned
		}
		return ErrAttachmentAlreadyAttached
	}
	return nil
}

// CopyAttachmentsTo 把源邮件的附件复制到目标邮件（转发时用）。
//
// 内容寻址下「复制」只是新增一条指向同一 sha256 的元数据，不拷磁盘文件。
// uploader 记为转发人：附件随新邮件重新分发，其可见范围由新邮件的参与方决定，
// 而不是沿用原上传者。返回复制的数量。
func CopyAttachmentsTo(ctx context.Context, srcMailID, dstMailID uuid.UUID, forwarder string) (int, error) {
	src, err := ListAttachmentsFor(ctx, srcMailID)
	if err != nil {
		return 0, err
	}
	for _, a := range src {
		_, err := db.DB.ExecContext(ctx, `
			INSERT INTO attachments (mail_id, uploader, filename, content_type, size_bytes, sha256)
			VALUES ($1, $2, $3, $4, $5, $6)
		`, dstMailID, forwarder, a.Filename, a.ContentType, a.SizeBytes, a.SHA256)
		if err != nil {
			return 0, err
		}
	}
	return len(src), nil
}

// DeleteAttachment 删除一条附件元数据，返回它的 sha256 以及该内容是否已无人引用。
// 内容寻址下多条记录可能共享同一个文件，只有最后一条引用消失才能删磁盘文件。
func DeleteAttachment(ctx context.Context, id uuid.UUID) (sum string, orphaned bool, err error) {
	a, err := GetAttachment(ctx, id)
	if err != nil {
		return "", false, err
	}
	if _, err = db.DB.ExecContext(ctx,
		`DELETE FROM attachments WHERE attachment_id = $1`, id); err != nil {
		return "", false, err
	}

	var refs int
	if err = db.DB.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM attachments WHERE sha256 = $1`, a.SHA256).Scan(&refs); err != nil {
		return "", false, err
	}
	return a.SHA256, refs == 0, nil
}

// SweepOrphanAttachments 清理超过 age 仍未挂载到邮件的附件记录，
// 返回可以从磁盘删除的 sha256 列表（已确认无任何记录引用）。
//
// 上传后没走完发信流程（用户取消、Agent 崩溃）会留下这类记录，
// 不清理的话磁盘只会单调增长。
func SweepOrphanAttachments(ctx context.Context, age time.Duration) ([]string, error) {
	cutoff := time.Now().Add(-age)

	rows, err := db.DB.QueryContext(ctx,
		`SELECT attachment_id, sha256 FROM attachments
		 WHERE mail_id IS NULL AND created_at < $1`, cutoff)
	if err != nil {
		return nil, err
	}
	type orphan struct {
		id  uuid.UUID
		sum string
	}
	var found []orphan
	for rows.Next() {
		var o orphan
		if err := rows.Scan(&o.id, &o.sum); err != nil {
			rows.Close()
			return nil, err
		}
		found = append(found, o)
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return nil, err
	}

	var removable []string
	for _, o := range found {
		if _, err := db.DB.ExecContext(ctx,
			`DELETE FROM attachments WHERE attachment_id = $1`, o.id); err != nil {
			return nil, err
		}
		var refs int
		if err := db.DB.QueryRowContext(ctx,
			`SELECT COUNT(*) FROM attachments WHERE sha256 = $1`, o.sum).Scan(&refs); err != nil {
			return nil, err
		}
		if refs == 0 {
			removable = append(removable, o.sum)
		}
	}
	return removable, nil
}

// AttachmentAccessible 判断某人是否有权读取某附件：
// 已挂载的看邮件所属会话的参与关系，未挂载的只有上传者本人能看。
func AttachmentAccessible(ctx context.Context, a *models.Attachment, name string) (bool, error) {
	if a.MailID == nil {
		return a.Uploader == name, nil
	}
	var n int
	err := db.DB.QueryRowContext(ctx, `
		SELECT COUNT(*) FROM mails m
		WHERE m.mail_id = $1
		  AND (m.from_name = $2 OR m.to_name = $2 OR `+db.CCHas("m.cc_list", 2)+`)
	`, *a.MailID, name).Scan(&n)
	return n > 0, err
}

// ListAttachmentsForMails 批量取多封邮件的附件，返回 mail_id → 附件列表。
//
// 为什么要批量：会话线程与收发件箱都是「一批邮件」，逐封调 ListAttachmentsFor
// 就是 N+1 —— 一个 200 封的会话打开一次要打 200 次库。
// 用 IN (...) 一次取回后在内存里分组。
//
// 占位符手工拼而非用数组参数：SQLite 驱动不支持 PG 的 = ANY($1)，
// 而这里的元素是已解析的 uuid.UUID，不存在注入面。
func ListAttachmentsForMails(ctx context.Context, mailIDs []uuid.UUID) (map[uuid.UUID][]models.Attachment, error) {
	out := map[uuid.UUID][]models.Attachment{}
	if len(mailIDs) == 0 {
		return out, nil
	}

	ph := make([]string, len(mailIDs))
	args := make([]any, len(mailIDs))
	for i, id := range mailIDs {
		ph[i] = fmt.Sprintf("$%d", i+1)
		args[i] = id
	}

	rows, err := db.DB.QueryContext(ctx,
		`SELECT `+attachmentCols+` FROM attachments
		 WHERE mail_id IN (`+strings.Join(ph, ",")+`)
		 ORDER BY created_at`, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		a, err := scanAttachment(rows)
		if err != nil {
			return nil, err
		}
		if a.MailID == nil {
			continue // WHERE 已排除，只是防御
		}
		out[*a.MailID] = append(out[*a.MailID], *a)
	}
	return out, rows.Err()
}