ultisuite-backend/internal/api/mail/attachments.go

package mail

import (
	"context"
	"encoding/json"
	"errors"
	"io"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"

	"github.com/ultisuite/ulti-backend/internal/mail/limits"
	"github.com/ultisuite/ulti-backend/internal/mail/storage"
)

var (
	ErrAttachmentNotFound       = errors.New("attachment not found")
	ErrAttachmentTooLarge       = limits.ErrAttachmentTooLarge
	ErrTooManyAttachments       = limits.ErrTooManyAttachments
	ErrAttachmentsTotalTooLarge = limits.ErrAttachmentsTotalTooLarge
)

type draftAttachmentRef struct {
	ID          string `json:"id"`
	Filename    string `json:"filename"`
	ContentType string `json:"content_type"`
	Size        int64  `json:"size"`
	S3Bucket    string `json:"s3_bucket"`
	S3Key       string `json:"s3_key"`
	ContentID   string `json:"content_id,omitempty"`
	IsInline    bool   `json:"is_inline"`
}

func (s *Service) ListMessageAttachments(ctx context.Context, externalID, messageID string) ([]map[string]any, error) {
	if _, err := s.ensureMessageOwned(ctx, externalID, messageID); err != nil {
		return nil, err
	}

	rows, err := s.db.Query(ctx, `
		SELECT id, filename, content_type, size, content_id, is_inline
		FROM attachments WHERE message_id = $1
		ORDER BY created_at ASC
	`, messageID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := make([]map[string]any, 0)
	for rows.Next() {
		var id, filename, contentType, contentID string
		var size int64
		var isInline bool
		if err := rows.Scan(&id, &filename, &contentType, &size, &contentID, &isInline); err != nil {
			return nil, err
		}
		entry := map[string]any{
			"id": id, "filename": filename, "content_type": contentType,
			"size": size, "is_inline": isInline,
		}
		if contentID != "" {
			entry["content_id"] = contentID
		}
		out = append(out, entry)
	}
	return out, rows.Err()
}

func (s *Service) MessageAttachmentCIDMap(ctx context.Context, externalID, messageID string) (map[string]string, error) {
	if _, err := s.ensureMessageOwned(ctx, externalID, messageID); err != nil {
		return nil, err
	}

	rows, err := s.db.Query(ctx, `
		SELECT id, content_id FROM attachments
		WHERE message_id = $1 AND content_id <> ''
	`, messageID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	mapping := make(map[string]string)
	for rows.Next() {
		var id, contentID string
		if err := rows.Scan(&id, &contentID); err != nil {
			return nil, err
		}
		mapping[contentID] = id
	}
	return mapping, rows.Err()
}

func (s *Service) UploadMessageAttachment(
	ctx context.Context, externalID, messageID, filename, contentType, contentID string,
	isInline bool, reader io.Reader, size int64,
) (string, error) {
	if s.storage == nil {
		return "", errors.New("object storage unavailable")
	}
	if err := limits.ValidateAttachmentSize(size); err != nil {
		return "", err
	}
	userID, err := s.ensureMessageOwned(ctx, externalID, messageID)
	if err != nil {
		return "", err
	}

	var count int
	var totalSize int64
	err = s.db.QueryRow(ctx, `
		SELECT COUNT(*)::int, COALESCE(SUM(size), 0)::bigint
		FROM attachments WHERE message_id = $1
	`, messageID).Scan(&count, &totalSize)
	if err != nil {
		return "", err
	}
	if err := limits.ValidateAttachmentQuota(count, totalSize, size); err != nil {
		return "", err
	}

	objectKey := storage.MessageObjectKey(userID, messageID, filename)
	if err := s.storage.Put(ctx, objectKey, reader, size, contentType); err != nil {
		return "", err
	}

	var id string
	err = s.db.QueryRow(ctx, `
		INSERT INTO attachments (message_id, filename, content_type, size, s3_bucket, s3_key, content_id, is_inline)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
		RETURNING id
	`, messageID, filename, contentType, size, s.storageBucket(), objectKey, contentID, isInline).Scan(&id)
	if err != nil {
		_ = s.storage.Delete(ctx, objectKey)
		return "", err
	}

	_, err = s.db.Exec(ctx, `
		UPDATE messages SET has_attachments = true, updated_at = NOW() WHERE id = $1
	`, messageID)
	if err != nil {
		return "", err
	}
	return id, nil
}

func (s *Service) OpenAttachment(ctx context.Context, externalID, attachmentID string) (
	filename, contentType string, size int64, isInline bool, body io.ReadCloser, err error,
) {
	if s.storage == nil {
		return "", "", 0, false, nil, errors.New("object storage unavailable")
	}

	var s3Key string
	err = s.db.QueryRow(ctx, `
		SELECT a.filename, a.content_type, a.size, a.is_inline, a.s3_key
		FROM attachments a
		JOIN messages m ON a.message_id = m.id
		JOIN mail_accounts ma ON m.account_id = ma.id
		WHERE a.id = $1 AND ma.user_id = (SELECT id FROM users WHERE external_id = $2)
	`, attachmentID, externalID).Scan(&filename, &contentType, &size, &isInline, &s3Key)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return "", "", 0, false, nil, ErrAttachmentNotFound
		}
		return "", "", 0, false, nil, err
	}

	obj, err := s.storage.Get(ctx, s3Key)
	if err != nil {
		return "", "", 0, false, nil, err
	}
	return filename, contentType, size, isInline, obj, nil
}

func (s *Service) UploadDraftAttachment(
	ctx context.Context, externalID, draftID, filename, contentType, contentID string,
	isInline bool, reader io.Reader, size int64,
) (string, error) {
	if s.storage == nil {
		return "", errors.New("object storage unavailable")
	}
	if err := limits.ValidateAttachmentSize(size); err != nil {
		return "", err
	}

	userID, err := s.ResolveUserID(ctx, externalID)
	if err != nil {
		return "", err
	}

	var attachmentsJSON []byte
	err = s.db.QueryRow(ctx, `
		SELECT attachments FROM outbox
		WHERE id = $1 AND user_id = $2 AND status = 'draft'
	`, draftID, userID).Scan(&attachmentsJSON)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return "", ErrNotFound
		}
		return "", err
	}

	objectKey := storage.DraftObjectKey(userID, draftID, filename)
	if err := s.storage.Put(ctx, objectKey, reader, size, contentType); err != nil {
		return "", err
	}

	refs := make([]draftAttachmentRef, 0)
	if len(attachmentsJSON) > 0 && string(attachmentsJSON) != "[]" {
		_ = json.Unmarshal(attachmentsJSON, &refs)
	}

	var totalSize int64
	for _, ref := range refs {
		totalSize += ref.Size
	}
	if err := limits.ValidateAttachmentQuota(len(refs), totalSize, size); err != nil {
		return "", err
	}

	attID := uuid.NewString()
	refs = append(refs, draftAttachmentRef{
		ID: attID, Filename: filename, ContentType: contentType, Size: size,
		S3Bucket: s.storageBucket(), S3Key: objectKey,
		ContentID: contentID, IsInline: isInline,
	})

	updated, _ := json.Marshal(refs)
	result, err := s.db.Exec(ctx, `
		UPDATE outbox SET attachments = $1, updated_at = NOW()
		WHERE id = $2 AND user_id = $3 AND status = 'draft'
	`, updated, draftID, userID)
	if err != nil {
		_ = s.storage.Delete(ctx, objectKey)
		return "", err
	}
	if result.RowsAffected() == 0 {
		_ = s.storage.Delete(ctx, objectKey)
		return "", ErrNotFound
	}
	return attID, nil
}

func (s *Service) OpenDraftAttachment(ctx context.Context, externalID, draftID, attachmentID string) (
	filename, contentType string, body io.ReadCloser, err error,
) {
	if s.storage == nil {
		return "", "", nil, errors.New("object storage unavailable")
	}

	userID, err := s.ResolveUserID(ctx, externalID)
	if err != nil {
		return "", "", nil, err
	}

	var attachmentsJSON []byte
	err = s.db.QueryRow(ctx, `
		SELECT attachments FROM outbox
		WHERE id = $1 AND user_id = $2 AND status = 'draft'
	`, draftID, userID).Scan(&attachmentsJSON)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return "", "", nil, ErrNotFound
		}
		return "", "", nil, err
	}

	var refs []draftAttachmentRef
	if err := json.Unmarshal(attachmentsJSON, &refs); err != nil {
		return "", "", nil, err
	}
	for _, ref := range refs {
		if ref.ID != attachmentID {
			continue
		}
		obj, err := s.storage.Get(ctx, ref.S3Key)
		if err != nil {
			return "", "", nil, err
		}
		return ref.Filename, ref.ContentType, obj, nil
	}
	return "", "", nil, ErrAttachmentNotFound
}

func (s *Service) ensureMessageOwned(ctx context.Context, externalID, messageID string) (userID string, err error) {
	err = s.db.QueryRow(ctx, `
		SELECT u.id FROM messages m
		JOIN mail_accounts ma ON m.account_id = ma.id
		JOIN users u ON ma.user_id = u.id
		WHERE m.id = $1 AND u.external_id = $2
	`, messageID, externalID).Scan(&userID)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return "", ErrNotFound
		}
		return "", err
	}
	return userID, nil
}

func (s *Service) storageBucket() string {
	if s.attachmentsBucket != "" {
		return s.attachmentsBucket
	}
	return "mail-attachments"
}