slogan/styleagent/service/effect_image_service.go

package service

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strings"

	"slogan-agent/styleagent/consts"
	"slogan-agent/styleagent/dao"
	"slogan-agent/styleagent/imagegen"
	"slogan-agent/styleagent/model/entity"

	"github.com/gogf/gf/v2/frame/g"
)

type effectImageService struct{}

var EffectImageService = new(effectImageService)

var effectAngles = []string{"正面", "侧面", "背面"}

// GenerateForPlan 选定主方案后异步生成 3 视角效果图
func (s *effectImageService) GenerateForPlan(ctx context.Context, planId, userId int64) {
	go s.run(ctx, planId, userId)
}

func (s *effectImageService) run(ctx context.Context, planId, userId int64) {
	plan, err := dao.OutfitPlan.GetOne(ctx, planId, userId)
	if err != nil || plan == nil {
		g.Log().Errorf(ctx, "效果图任务: 方案不存在 planId=%d", planId)
		return
	}
	// 每日免费次数校验
	limit := dailyEffectLimit(ctx)
	if limit > 0 {
		used, _ := dao.PlanEffectImage.CountByUserToday(ctx, userId)
		if used >= limit {
			g.Log().Warningf(ctx, "效果图任务: 用户 %d 当日次数已用尽（%d/%d）", userId, used, limit)
			return
		}
	}

	_ = dao.OutfitGenTask.UpdateStatus(ctx, plan.TaskId, consts.TaskStatusRendering, "")
	defer dao.OutfitGenTask.UpdateStatus(ctx, plan.TaskId, consts.TaskStatusDone, "")

	items, _ := dao.PlanOutfitItem.ListByPlan(ctx, planId)
	planDesc := planTitleDesc(plan.Title, items)

	// 用户全身正面照作 base image
	baseImageURL := ""
	if photos, err := dao.UserPhoto.ListByUser(ctx, userId, 0); err == nil {
		for _, p := range photos {
			if p.Type == consts.PhotoTypeFullFront {
				baseImageURL = p.Url
				break
			}
		}
	}

	client := imagegen.NewClient(g.Cfg().MustGet(ctx, "imagegen.supplier", "mock").String())
	for i, angle := range effectAngles {
		cacheKey := effectCacheKey(plan, angle)
		if url, ok := imagegen.CacheGet(cacheKey); ok {
			_, _ = dao.PlanEffectImage.Insert(ctx, &entity.PlanEffectImage{
				PlanId: planId, Angle: angle, Url: url, Status: consts.EffectStatusDone,
				PromptSnapshot: planDesc,
			})
			continue
		}
		recId, err := dao.PlanEffectImage.Insert(ctx, &entity.PlanEffectImage{
			PlanId: planId, Angle: angle, Status: consts.EffectStatusRendering,
			PromptSnapshot: planDesc,
		})
		if err != nil {
			continue
		}
		url, err := client.Generate(ctx, &imagegen.GenerateReq{
			BaseImageURL: baseImageURL, Prompt: planDesc, Angle: angle, Seed: plan.Id*100 + int64(i),
		})
		if err != nil {
			g.Log().Warningf(ctx, "效果图生成失败 plan=%d angle=%s: %v", planId, angle, err)
			_ = dao.PlanEffectImage.UpdateStatus(ctx, recId, consts.EffectStatusFailed, "")
			continue
		}
		imagegen.CacheSet(cacheKey, url)
		_ = dao.PlanEffectImage.UpdateStatus(ctx, recId, consts.EffectStatusDone, url)
	}
	g.Log().Infof(ctx, "方案 %d 效果图生成完成", planId)
}

func planTitleDesc(title string, items []*entity.PlanOutfitItem) string {
	var sb strings.Builder
	sb.WriteString("方案：")
	sb.WriteString(title)
	sb.WriteString("；")
	for _, it := range items {
		sb.WriteString(it.Slot)
		sb.WriteString("：")
		sb.WriteString(it.Name)
		sb.WriteString("；")
	}
	return strings.TrimSuffix(sb.String(), "；")
}

func effectCacheKey(plan *entity.OutfitPlan, angle string) string {
	sum := md5.Sum([]byte(fmt.Sprintf("%d:%s:%s", plan.Id, plan.Title, angle)))
	return "plan:" + hex.EncodeToString(sum[:])
}

func dailyEffectLimit(ctx context.Context) int {
	limit := consts.DefaultDailyEffectLimit
	rules, err := dao.ScoringRule.ListEnabled(ctx)
	if err != nil {
		return limit
	}
	for _, r := range rules {
		if r.Dimension == "effect_limit" {
			var v struct {
				Daily int `json:"daily"`
			}
			if json.Unmarshal([]byte(r.RulesJson), &v) == nil && v.Daily > 0 {
				return v.Daily
			}
		}
	}
	return limit
}