ultisuite-backend/internal/migration/gmail_attachments.go

package migration

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"mime"
	"strings"

	"github.com/jackc/pgx/v5/pgxpool"

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

type gmailAttachmentPart struct {
	Filename     string
	MimeType     string
	ContentID    string
	IsInline     bool
	Size         int
	BodyData     string
	AttachmentID string
}

func extractGmailAttachmentParts(payload gmailPayload) []gmailAttachmentPart {
	var out []gmailAttachmentPart
	var walk func(gmailPayload)
	walk = func(node gmailPayload) {
		headers := gmailHeaderMap(node.Headers)
		mimeType := strings.TrimSpace(node.MimeType)
		if strings.HasPrefix(strings.ToLower(mimeType), "multipart/") {
			for _, part := range node.Parts {
				walk(part)
			}
			return
		}
		if isGmailAttachmentPart(mimeType, headers, node.Body) {
			out = append(out, buildGmailAttachmentPart(mimeType, headers, node.Body))
		}
	}
	walk(payload)
	if len(out) > limits.MaxAttachmentsPerMessage {
		return out[:limits.MaxAttachmentsPerMessage]
	}
	return out
}

func gmailHeaderMap(headers []gmailHeader) map[string]string {
	out := make(map[string]string, len(headers))
	for _, h := range headers {
		key := strings.ToLower(strings.TrimSpace(h.Name))
		if key != "" && out[key] == "" {
			out[key] = h.Value
		}
	}
	return out
}

func isGmailAttachmentPart(mimeType string, headers map[string]string, body gmailBody) bool {
	if body.AttachmentID != "" {
		return true
	}
	disposition := strings.ToLower(headers["content-disposition"])
	if strings.HasPrefix(disposition, "attachment") {
		return true
	}
	contentID := normalizeGmailContentID(headers["content-id"])
	filename := gmailPartFilename(headers)
	if strings.HasPrefix(disposition, "inline") && (filename != "" || contentID != "") {
		return true
	}
	if contentID != "" && body.Data != "" {
		return true
	}
	lower := strings.ToLower(strings.TrimSpace(mimeType))
	if lower == "text/plain" || lower == "text/html" || lower == "" {
		return false
	}
	return filename != "" && body.Data != ""
}

func buildGmailAttachmentPart(mimeType string, headers map[string]string, body gmailBody) gmailAttachmentPart {
	disposition := strings.ToLower(headers["content-disposition"])
	contentID := normalizeGmailContentID(headers["content-id"])
	filename := gmailPartFilename(headers)
	isInline := strings.HasPrefix(disposition, "inline") || (contentID != "" && !strings.HasPrefix(disposition, "attachment"))
	if filename == "" {
		filename = inlineGmailAttachmentFilename(contentID, mimeType)
	}
	size := body.Size
	if size <= 0 && body.Data != "" {
		size = len(body.Data)
	}
	return gmailAttachmentPart{
		Filename:     filename,
		MimeType:     normalizeAttachmentMimeType(mimeType),
		ContentID:    contentID,
		IsInline:     isInline,
		Size:         size,
		BodyData:     body.Data,
		AttachmentID: body.AttachmentID,
	}
}

func gmailPartFilename(headers map[string]string) string {
	if name := strings.TrimSpace(headers["filename"]); name != "" {
		return name
	}
	disposition := headers["content-disposition"]
	if disposition == "" {
		return ""
	}
	_, params, err := mime.ParseMediaType(disposition)
	if err != nil {
		return ""
	}
	return strings.TrimSpace(params["filename"])
}

func normalizeGmailContentID(raw string) string {
	return strings.Trim(raw, "<> \t")
}

func inlineGmailAttachmentFilename(contentID, mimeType string) string {
	base := "inline"
	if contentID != "" {
		base = strings.Map(func(r rune) rune {
			switch r {
			case '<', '>', '/', '\\', ':', '"', '\'', '?', '*':
				return '_'
			default:
				return r
			}
		}, contentID)
	}
	ext := extensionFromMimeType(mimeType)
	if ext != "" && !strings.HasSuffix(strings.ToLower(base), ext) {
		return base + ext
	}
	return base
}

func extensionFromMimeType(mimeType string) string {
	switch strings.ToLower(strings.TrimSpace(strings.Split(mimeType, ";")[0])) {
	case "image/jpeg", "image/jpg":
		return ".jpg"
	case "image/png":
		return ".png"
	case "image/gif":
		return ".gif"
	case "image/webp":
		return ".webp"
	case "application/pdf":
		return ".pdf"
	default:
		return ""
	}
}

func normalizeAttachmentMimeType(mimeType string) string {
	mimeType = strings.TrimSpace(mimeType)
	if mimeType == "" {
		return "application/octet-stream"
	}
	mediaType, _, err := mime.ParseMediaType(mimeType)
	if err != nil || mediaType == "" {
		return mimeType
	}
	return mediaType
}

func (g *GmailImporter) storeGmailAttachments(
	ctx context.Context,
	userID, messageID, gmailID, accessToken string,
	payload gmailPayload,
	messageExisted bool,
) error {
	if g.storage == nil {
		return nil
	}
	parts := extractGmailAttachmentParts(payload)
	if len(parts) == 0 {
		return nil
	}

	var existingCount int
	var existingTotal int64
	if err := g.db.QueryRow(ctx, `
		SELECT COUNT(*)::int, COALESCE(SUM(size), 0)::bigint
		FROM attachments WHERE message_id = $1::uuid
	`, messageID).Scan(&existingCount, &existingTotal); err != nil {
		return err
	}

	stored := 0
	for _, part := range parts {
		if messageExisted && gmailAttachmentExists(ctx, g.db, messageID, part) {
			continue
		}
		data, err := g.loadGmailAttachmentData(ctx, accessToken, gmailID, part)
		if err != nil {
			return err
		}
		if len(data) == 0 {
			continue
		}
		if err := limits.ValidateAttachmentSize(int64(len(data))); err != nil {
			continue
		}
		if err := limits.ValidateAttachmentQuota(existingCount, existingTotal, int64(len(data))); err != nil {
			break
		}

		objectKey := storage.MessageObjectKey(userID, messageID, part.Filename)
		if err := g.storage.Put(ctx, objectKey, bytes.NewReader(data), int64(len(data)), part.MimeType); err != nil {
			return err
		}
		bucket := g.attachBucket
		if bucket == "" {
			bucket = "mail-attachments"
		}
		_, err = g.db.Exec(ctx, `
			INSERT INTO attachments (message_id, filename, content_type, size, s3_bucket, s3_key, content_id, is_inline, virus_scan_status)
			VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8, 'skipped')
		`, messageID, part.Filename, part.MimeType, len(data), bucket, objectKey, part.ContentID, part.IsInline)
		if err != nil {
			_ = g.storage.Delete(ctx, objectKey)
			return err
		}
		existingCount++
		existingTotal += int64(len(data))
		stored++
	}

	if stored > 0 {
		_, err := g.db.Exec(ctx, `UPDATE messages SET has_attachments = true, updated_at = NOW() WHERE id = $1::uuid`, messageID)
		return err
	}
	return nil
}

func (g *GmailImporter) loadGmailAttachmentData(ctx context.Context, accessToken, gmailID string, part gmailAttachmentPart) ([]byte, error) {
	if part.AttachmentID != "" {
		url := fmt.Sprintf(
			"https://gmail.googleapis.com/gmail/v1/users/me/messages/%s/attachments/%s",
			gmailID,
			part.AttachmentID,
		)
		raw, err := g.apiGet(ctx, url, accessToken)
		if err != nil {
			return nil, err
		}
		var parsed struct {
			Data string `json:"data"`
		}
		if err := json.Unmarshal(raw, &parsed); err != nil {
			return nil, err
		}
		if parsed.Data == "" {
			return nil, nil
		}
		return []byte(decodeGmailBody(parsed.Data)), nil
	}
	if part.BodyData == "" {
		return nil, nil
	}
	return []byte(decodeGmailBody(part.BodyData)), nil
}

func gmailAttachmentExists(ctx context.Context, db *pgxpool.Pool, messageID string, part gmailAttachmentPart) bool {
	var count int
	if part.ContentID != "" {
		_ = db.QueryRow(ctx, `
			SELECT COUNT(*) FROM attachments
			WHERE message_id = $1::uuid AND (content_id = $2 OR filename = $3)
		`, messageID, part.ContentID, part.Filename).Scan(&count)
		return count > 0
	}
	_ = db.QueryRow(ctx, `
		SELECT COUNT(*) FROM attachments WHERE message_id = $1::uuid AND filename = $2
	`, messageID, part.Filename).Scan(&count)
	return count > 0
}