slogan/scripts/gen_outfit_plan/main.go

package main

// 为 wenwu901 真实生成一套 AI 穿搭方案（真实调用 imagegen 与 LLM，非 mock）：
//   go run scripts/gen_outfit_plan/main.go
// 步骤：补衣橱（8 件单品，imagegen 生成服装图）→ 调 OutfitService.Generate →
// 轮询任务到 done → 选主方案（触发效果图异步生成）→ 等 3 张效果图完成。
// 前置：config.yml 已配置 geo.amap_key + weather.qweather_key（天气硬依赖）。

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	_ "github.com/gogf/gf/contrib/drivers/sqlite/v2"
	"github.com/gogf/gf/v2/frame/g"

	"slogan-agent/styleagent/agent"
	"slogan-agent/styleagent/consts"
	"slogan-agent/styleagent/dao"
	"slogan-agent/styleagent/model/dto"
	"slogan-agent/styleagent/model/entity"
	"slogan-agent/styleagent/service"
)

const username = "wenwu901"

type garment struct {
	name     string
	category string
	style    string
	color    string
	prompt   string
}

var garments = []garment{
	{"白色长袖衬衫", "上衣", "休闲", "白色", "纯白背景的白色长袖衬衫商品图，正面展示，高清，电商风格"},
	{"灰色圆领T恤", "上衣", "休闲", "灰色", "纯白背景的灰色圆领T恤商品图，正面展示，高清，电商风格"},
	{"深蓝夹克外套", "上衣", "外套", "深蓝", "纯白背景的深蓝色夹克外套商品图，正面展示，高清，电商风格"},
	{"深灰休闲长裤", "下装", "休闲", "深灰", "纯白背景的深灰色休闲长裤商品图，正面展示，高清，电商风格"},
	{"蓝色牛仔裤", "下装", "休闲", "蓝色", "纯白背景的蓝色牛仔裤商品图，正面展示，高清，电商风格"},
	{"白色运动鞋", "鞋", "休闲", "白色", "纯白背景的白色运动鞋商品图，侧面展示，高清，电商风格"},
	{"棕色皮鞋", "鞋", "商务", "棕色", "纯白背景的棕色皮鞋商品图，侧面展示，高清，电商风格"},
	{"黑色双肩背包", "配饰", "休闲", "黑色", "纯白背景的黑色双肩背包商品图，正面展示，高清，电商风格"},
}

func main() {
	ctx := context.Background()

	var user entity.User
	if err := g.DB().Model(consts.TableNameUser).Ctx(ctx).
		Where("username", username).Scan(&user); err != nil || user.Id == 0 {
		panic(fmt.Sprintf("用户 %s 不存在: %v", username, err))
	}
	fmt.Printf("用户: %s (id=%d)\n", username, user.Id)

	if err := ensureWardrobe(ctx, user.Id); err != nil {
		panic(err)
	}

	// 幂等：已有方案则只重试效果图（主方案 → select-main → 等 3 张 done）
	var mainPlan *entity.OutfitPlan
	existing, _ := dao.OutfitPlan.ListByUser(ctx, user.Id)
	for _, p := range existing {
		if p.MainFlag == 1 {
			mainPlan = p
		}
	}
	if len(existing) == 0 {
		taskId, err := service.OutfitService.Generate(ctx, user.Id, &dto.OutfitGenerateReq{
			StartDate: "2026-08-01", EndDate: "2026-08-07", Location: "上海", Occasion: "通勤",
		})
		if err != nil {
			if strings.Contains(err.Error(), "未配置") {
				panic(fmt.Sprintf("%v\n请先在 config.yml 配置 geo.amap_key / weather.qweather_key 后重跑", err))
			}
			panic(fmt.Sprintf("发起方案生成失败: %v", err))
		}
		fmt.Printf("生成任务已提交: task_id=%d，轮询中...\n", taskId)

		waitTaskDone(ctx, taskId, user.Id)
		plans, err := dao.OutfitPlan.ListByTask(ctx, taskId)
		if err != nil || len(plans) == 0 {
			panic(fmt.Sprintf("任务完成但无方案: %v", err))
		}
		mainPlan = plans[0]
		for _, p := range plans {
			if p.MainFlag == 1 {
				mainPlan = p
			}
		}
	} else if mainPlan == nil {
		panic("已有方案但无主方案，请先选主方案")
	}
	fmt.Printf("主方案: id=%d %s 评分=%d\n", mainPlan.Id, mainPlan.Title, mainPlan.Score)
	if err := service.OutfitPlanService.SelectMain(ctx, user.Id, mainPlan.Id); err != nil {
		panic(fmt.Sprintf("选主方案失败: %v", err))
	}

	waitEffects(ctx, mainPlan.Id)

	all, _ := dao.OutfitPlan.ListByUser(ctx, user.Id)
	fmt.Printf("完成！wenwu901 现有 %d 套方案：\n", len(all))
	for _, p := range all {
		fmt.Printf("  - plan %d: %s（评分 %d，%s）\n", p.Id, p.Title, p.Score, p.Source)
	}
}

// ensureWardrobe 为指定用户补齐 8 件衣橱单品（同 Category 已有则跳过该分类），服装图用 imagegen 生成
func ensureWardrobe(ctx context.Context, userId int64) error {
	existing, err := dao.WardrobeItem.ListAllByUser(ctx, userId)
	if err != nil {
		return err
	}
	have := map[string]bool{}
	for _, it := range existing {
		have[it.Category] = true
	}
	need := make([]garment, 0, len(garments))
	for _, ga := range garments {
		if !have[ga.category] {
			need = append(need, ga)
		}
	}
	if len(need) == 0 {
		fmt.Println("衣橱 4 类已齐，跳过补衣橱")
		return nil
	}

	client, err := agent.NewClient(g.Cfg().MustGet(ctx, "imagegen.supplier", "wanx").String())
	if err != nil {
		return err
	}
	dir := filepath.Join("workspace", fmt.Sprintf("user_%d", userId), "wardrobe")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	for _, ga := range need {
		fmt.Printf("生成服装图: %s...\n", ga.name)
		url, err := client.Generate(ctx, &agent.GenerateReq{
			Prompt: ga.prompt, Seed: time.Now().UnixNano() % 1_000_000,
		})
		if err != nil {
			return fmt.Errorf("生成 %s 服装图失败: %w", ga.name, err)
		}
		path := filepath.Join(dir, fmt.Sprintf("%d_%s.png", time.Now().UnixNano(), ga.name))
		if err := download(url, path); err != nil {
			return fmt.Errorf("保存 %s 失败: %w", ga.name, err)
		}
		if _, err := dao.WardrobeItem.Insert(ctx, &entity.WardrobeItem{
			UserId: userId, PhotoUrl: "/" + filepath.ToSlash(path),
			Name: ga.name, Category: ga.category, Season: "四季", StyleTags: ga.style,
			ColorInfo: ga.color, Status: 1,
		}); err != nil {
			return fmt.Errorf("入库 %s 失败: %w", ga.name, err)
		}
		fmt.Printf("%s 完成: %s\n", ga.name, path)
	}
	fmt.Println("衣橱补齐完毕")
	return nil
}

func waitTaskDone(ctx context.Context, taskId, userId int64) {
	for i := 0; i < 30; i++ {
		task, err := dao.OutfitGenTask.GetOne(ctx, taskId, userId)
		if err != nil || task == nil {
			panic(fmt.Sprintf("读取任务失败: %v", err))
		}
		switch task.Status {
		case consts.TaskStatusDone:
			fmt.Println("方案生成完成")
			return
		case consts.TaskStatusFailed:
			panic(fmt.Sprintf("方案生成失败: %s", task.Error))
		}
		time.Sleep(10 * time.Second)
	}
	panic("方案生成超时（5 分钟）")
}

// waitEffects 等主方案的 3 张效果图（正面/侧面/背面）全部 done
func waitEffects(ctx context.Context, planId int64) {
	for i := 0; i < 20; i++ {
		images, _ := dao.PlanEffectImage.ListByPlan(ctx, planId)
		done := 0
		for _, im := range images {
			if im.Status == consts.EffectStatusDone {
				done++
			}
		}
		if done >= 3 {
			fmt.Printf("效果图 3 张完成\n")
			return
		}
		if i == 19 {
			fmt.Printf("警告: 效果图超时（完成 %d/3），可稍后查看\n", done)
			return
		}
		time.Sleep(15 * time.Second)
	}
}

func download(url, dest string) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("下载失败: http %d", resp.StatusCode)
	}
	out, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, resp.Body)
	return err
}