MetaLab/internal/service/avatar_service.go

package service

import (
	"bytes"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"metazone.cc/metalab/internal/model"

	"github.com/chai2010/webp"
	"golang.org/x/image/draw"
)

// 头像处理常量
const (
	avatarMaxFileSize = 5 * 1024 * 1024    // 5MB
	avatarMinWidth    = 128                 // 最小宽高，防止马赛克图被拉升
	avatarMinHeight   = 128
	avatarMaxWidth    = 3840
	avatarMaxHeight   = 2160
	avatarOutputSize  = 512
	avatarTargetKB    = 100
	avatarStorageDir  = "storage/uploads/avatars"
)

// userAvatarStore 头像服务所需的最小仓储接口
type userAvatarStore interface {
	FindByID(id uint) (*model.User, error)
	Update(user *model.User) error
	FindByIDUnscoped(userID uint) (*model.User, error)
}

// AvatarService 头像上传处理服务
type AvatarService struct {
	userRepo   userAvatarStore
	storageDir string
}

// NewAvatarService 构造函数
func NewAvatarService(userRepo userAvatarStore) *AvatarService {
	return &AvatarService{userRepo: userRepo, storageDir: avatarStorageDir}
}

// ProcessAvatar 处理头像上传流程（完整流程：处理图片 + 更新 User）
func (s *AvatarService) ProcessAvatar(userID uint, file io.Reader, contentType string, cropX, cropY, cropSize int) (string, error) {
	avatarURL, err := s.ProcessImage(userID, file, contentType, cropX, cropY, cropSize)
	if err != nil {
		return "", err
	}

	user, err := s.userRepo.FindByID(userID)
	if err != nil {
		return "", fmt.Errorf("用户不存在")
	}
	user.Avatar = avatarURL
	if err := s.userRepo.Update(user); err != nil {
		return "", fmt.Errorf("更新头像失败")
	}

	return avatarURL, nil
}

// ProcessImage 处理头像上传流程（仅处理图片 + 保存文件，不更新 User，审核场景用）
func (s *AvatarService) ProcessImage(userID uint, file io.Reader, contentType string, cropX, cropY, cropSize int) (string, error) {
	// 1. 检查文件类型（不支持 GIF：裁切后无法保持动画）
	if contentType != "image/jpeg" && contentType != "image/png" && contentType != "image/webp" {
		return "", fmt.Errorf("仅支持 JPEG、PNG、WebP 格式")
	}

	// 2. 读取全部字节并校验大小
	data, err := io.ReadAll(io.LimitReader(file, avatarMaxFileSize+1))
	if err != nil {
		return "", fmt.Errorf("读取文件失败")
	}
	if len(data) > avatarMaxFileSize {
		return "", fmt.Errorf("文件大小不能超过 5MB")
	}

	// 3. 解码图片（WebP 需显式调用 webp.Decode，Go 标准库不含 WebP 解码器）
	var img image.Image
	switch contentType {
	case "image/webp":
		img, err = webp.Decode(bytes.NewReader(data))
	default:
		img, _, err = image.Decode(bytes.NewReader(data))
	}
	if err != nil {
		return "", fmt.Errorf("图片解码失败：%v", err)
	}
	bounds := img.Bounds()
	w, h := bounds.Dx(), bounds.Dy()
	if w > avatarMaxWidth || h > avatarMaxHeight {
		return "", fmt.Errorf("图片分辨率不能超过 3840x2160")
	}
	if w < avatarMinWidth || h < avatarMinHeight {
		return "", fmt.Errorf("图片分辨率不能低于 128x128")
	}

	// 4. 裁切为正方形（优先使用自定义裁切参数）
	var cropped image.Image
	if cropSize > 0 && cropX >= 0 && cropY >= 0 &&
		cropX+cropSize <= w && cropY+cropSize <= h {
		cropped = subImage(img, cropX, cropY, cropSize)
	} else {
		cropped = centerCrop(img)
	}

	// 5. 缩放至 512x512（CatmullRom 高质量）
	resized := image.NewNRGBA(image.Rect(0, 0, avatarOutputSize, avatarOutputSize))
	draw.CatmullRom.Scale(resized, resized.Bounds(), cropped, cropped.Bounds(), draw.Over, nil)

	// 6. 编码为 WebP，目标 ≤100KB
	webpData, err := encodeWebP(resized, avatarTargetKB*1024)
	if err != nil {
		return "", fmt.Errorf("图片编码失败：%v", err)
	}

	// 7. 确保存储目录存在
	if err := os.MkdirAll(s.storageDir, 0755); err != nil {
		return "", fmt.Errorf("创建存储目录失败")
	}

	// 8. 删除该用户的旧头像，仅保留最新
	s.cleanOldAvatars(userID)

	// 9. 原子写入：先写临时文件，再 rename 到最终路径
	ts := time.Now().Unix()
	filename := fmt.Sprintf("%d_%d.webp", userID, ts)
	tmpPath := filepath.Join(s.storageDir, filename+".tmp")
	finalPath := filepath.Join(s.storageDir, filename)
	if err := os.WriteFile(tmpPath, webpData, 0644); err != nil {
		return "", fmt.Errorf("写入临时文件失败")
	}
	if err := os.Rename(tmpPath, finalPath); err != nil {
		os.Remove(tmpPath)
		return "", fmt.Errorf("保存文件失败")
	}

	return fmt.Sprintf("/uploads/avatars/%d_%d.webp", userID, ts), nil
}

// cleanOldAvatars 删除指定用户的所有旧头像文件（仅保留最新一次上传）
func (s *AvatarService) cleanOldAvatars(userID uint) {
	prefix := strconv.FormatUint(uint64(userID), 10) + "_"
	entries, err := os.ReadDir(s.storageDir)
	if err != nil {
		return
	}
	for _, e := range entries {
		if !e.IsDir() && strings.HasPrefix(e.Name(), prefix) && strings.HasSuffix(e.Name(), ".webp") {
			os.Remove(filepath.Join(s.storageDir, e.Name()))
		}
	}
}

// subImage 从原图中裁切指定正方形区域
func subImage(img image.Image, x, y, size int) image.Image {
	cropped := image.NewNRGBA(image.Rect(0, 0, size, size))
	draw.Draw(cropped, cropped.Bounds(), img, image.Point{x, y}, draw.Src)
	return cropped
}

// centerCrop 中心裁切为正方形（取 min(width, height)）
func centerCrop(img image.Image) image.Image {
	b := img.Bounds()
	w, h := b.Dx(), b.Dy()
	size := w
	if h < size {
		size = h
	}
	ox := (w - size) / 2
	oy := (h - size) / 2
	cropped := image.NewNRGBA(image.Rect(0, 0, size, size))
	draw.Draw(cropped, cropped.Bounds(), img, image.Point{ox, oy}, draw.Src)
	return cropped
}

// encodeWebP 编码为 WebP，通过二分查找质量参数使文件大小接近 targetBytes
func encodeWebP(img image.Image, targetBytes int) ([]byte, error) {
	buf := &bytes.Buffer{}

	// 如果图片非常简单（如纯色），即使 quality=1 也可能很小
	if err := webp.Encode(buf, img, &webp.Options{Quality: 80}); err != nil {
		return nil, err
	}
	if buf.Len() <= targetBytes {
		return buf.Bytes(), nil
	}

	// 二分查找合适的质量参数
	low, high := 1, 80
	var best []byte
	for low <= high {
		mid := (low + high) / 2
		buf.Reset()
		if err := webp.Encode(buf, img, &webp.Options{Quality: float32(mid)}); err != nil {
			return nil, err
		}
		if buf.Len() <= targetBytes {
			best = make([]byte, buf.Len())
			copy(best, buf.Bytes())
			low = mid + 1 // 尝试更高画质
		} else {
			high = mid - 1
		}
	}
	if best != nil {
		return best, nil
	}
	// 降到底仍超限，返回最低质量
	buf.Reset()
	if err := webp.Encode(buf, img, &webp.Options{Quality: 1}); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}