slogan/server/styleagent/agent/avatar_tripo_client.go

package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"time"

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

// TripoClient 3D 化身客户端（图像转 3D：上传图片 → 提交 multiview 任务 → 轮询 → 下载 GLB）
type TripoClient struct {
	apiKey       string
	base         string
	version      string
	pollInterval time.Duration
	pollTimeout  time.Duration
}

func NewTripoClient(ctx context.Context) *TripoClient {
	return &TripoClient{
		apiKey:       g.Cfg().MustGet(ctx, "avatar.tripo_api_key", "").String(),
		base:         g.Cfg().MustGet(ctx, "avatar.tripo_base", "https://api.tripo3d.ai/v2/openapi").String(),
		version:      g.Cfg().MustGet(ctx, "avatar.tripo_model_version", "v2.5-20250123").String(),
		pollInterval: time.Duration(g.Cfg().MustGet(ctx, "avatar.poll_interval", 5).Int()) * time.Second,
		pollTimeout:  time.Duration(g.Cfg().MustGet(ctx, "avatar.poll_timeout", 900).Int()) * time.Second,
	}
}

// Enabled 是否已配置 API Key
func (c *TripoClient) Enabled() bool { return c.apiKey != "" }

// UploadImage 上传单张图片，返回 file_token
func (c *TripoClient) UploadImage(ctx context.Context, filePath string) (string, error) {
	body := &bytes.Buffer{}
	w := multipart.NewWriter(body)
	f, err := os.Open(filePath)
	if err != nil {
		return "", fmt.Errorf("打开图片失败: %w", err)
	}
	defer f.Close()
	fw, err := w.CreateFormFile("file", filepath.Base(filePath))
	if err != nil {
		return "", err
	}
	if _, err := io.Copy(fw, f); err != nil {
		return "", err
	}
	w.Close()

	req, err := http.NewRequestWithContext(ctx, "POST", c.base+"/upload/sts", body)
	if err != nil {
		return "", err
	}
	req.Header.Set("Authorization", "Bearer "+c.apiKey)
	req.Header.Set("Content-Type", w.FormDataContentType())

	data, err := c.do(req)
	if err != nil {
		return "", err
	}
	for _, key := range []string{"file_token", "image_token", "token"} {
		if v, ok := data[key].(string); ok && v != "" {
			return v, nil
		}
	}
	return "", fmt.Errorf("Tripo 上传响应缺少 file_token: %s", mustJSONStr(data))
}

// SubmitMultiview 提交多视角转 3D 任务（front 必填，left/back 可空），返回 task_id
func (c *TripoClient) SubmitMultiview(ctx context.Context, front, left, back string) (string, error) {
	files := make([]map[string]string, 0, 3)
	for _, t := range []string{front, left, back} {
		if t != "" {
			files = append(files, map[string]string{"type": "image", "file_token": t})
		}
	}
	body, err := json.Marshal(map[string]any{
		"type":          "multiview_to_model",
		"model_version": c.version,
		"files":         files,
		"texture":       true,
		"pbr":           true,
	})
	if err != nil {
		return "", err
	}
	req, err := http.NewRequestWithContext(ctx, "POST", c.base+"/task", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Authorization", "Bearer "+c.apiKey)
	req.Header.Set("Content-Type", "application/json")

	data, err := c.do(req)
	if err != nil {
		return "", err
	}
	taskID, _ := data["task_id"].(string)
	if taskID == "" {
		return "", fmt.Errorf("Tripo 提交任务响应缺少 task_id: %s", mustJSONStr(data))
	}
	return taskID, nil
}

// PollTask 轮询任务直到 success/failed，成功返回 GLB 下载地址
func (c *TripoClient) PollTask(ctx context.Context, taskID string) (string, error) {
	deadline := time.Now().Add(c.pollTimeout)
	for {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		default:
		}
		req, err := http.NewRequestWithContext(ctx, "GET", c.base+"/task/"+taskID, nil)
		if err != nil {
			return "", err
		}
		req.Header.Set("Authorization", "Bearer "+c.apiKey)
		data, err := c.do(req)
		if err != nil {
			return "", err
		}
		status, _ := data["status"].(string)
		if status == "" {
			return "", fmt.Errorf("Tripo 任务响应缺少 status: %s", mustJSONStr(data))
		}
		switch status {
		case "success":
			if output, ok := data["output"].(map[string]any); ok {
				if pbr, ok := output["pbr_model"].(map[string]any); ok {
					if url, ok := pbr["url"].(string); ok && url != "" {
						return url, nil
					}
				}
			}
			return "", fmt.Errorf("Tripo 任务成功但无模型下载地址")
		case "failed", "cancelled", "expired":
			msg, _ := data["error"].(string)
			if msg == "" {
				msg = mustJSONStr(data)
			}
			return "", fmt.Errorf("Tripo 任务%s: %s", status, msg)
		}
		if time.Now().After(deadline) {
			return "", fmt.Errorf("Tripo 任务超时（%s）", taskID)
		}
		time.Sleep(c.pollInterval)
	}
}

// DownloadGlb 下载 GLB 到 destPath（下载地址约 5 分钟过期，任务成功后应立即调用）
func (c *TripoClient) DownloadGlb(ctx context.Context, url, destPath string) error {
	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("下载 GLB 失败: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("下载 GLB 失败: http %d", resp.StatusCode)
	}
	if err := os.MkdirAll(filepath.Dir(destPath), 0o755); err != nil {
		return err
	}
	out, err := os.Create(destPath)
	if err != nil {
		return err
	}
	defer out.Close()
	if _, err := io.Copy(out, resp.Body); err != nil {
		return err
	}
	return nil
}

// do 统一请求：非 2xx 或 code != 0 时返回业务错误
func (c *TripoClient) do(req *http.Request) (map[string]any, error) {
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("Tripo 请求失败: %w", err)
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取 Tripo 响应失败: %w", err)
	}
	var r struct {
		Code    int            `json:"code"`
		Message string         `json:"message"`
		Data    map[string]any `json:"data"`
	}
	if err := json.Unmarshal(raw, &r); err != nil {
		return nil, fmt.Errorf("Tripo 响应解析失败: %s", string(raw))
	}
	if resp.StatusCode != http.StatusOK || r.Code != 0 {
		return nil, fmt.Errorf("Tripo 接口错误 code=%d msg=%s", r.Code, r.Message)
	}
	return r.Data, nil
}

func mustJSONStr(v any) string {
	b, err := json.Marshal(v)
	if err != nil {
		return fmt.Sprintf("%v", v)
	}
	return string(b)
}