JuYou/API/internal/provider/apimart/seedance2.go

package apimart

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"

	"juhe-factory/api/internal/provider"
)

type Seedance2 struct {
	Gateway Client
}

type ReferenceCapability struct {
	Supported          bool  `json:"supported"`
	MaxCount           int   `json:"max_count"`
	MinTotalDurationMS int64 `json:"min_total_duration_ms,omitempty"`
	MaxTotalDurationMS int64 `json:"max_total_duration_ms,omitempty"`
	MinShortSidePixels int   `json:"min_short_side_pixels,omitempty"`
	MaxShortSidePixels int   `json:"max_short_side_pixels,omitempty"`
}

type VideoModelCapabilities struct {
	// DurationMinSeconds / DurationMaxSeconds 表示视频模型支持的生成时长范围（秒），供前端可视化选择与后端校验。
	DurationMinSeconds int `json:"duration_min_seconds"`
	DurationMaxSeconds int `json:"duration_max_seconds"`
	ReferenceImages    ReferenceCapability `json:"reference_images"`
	ReferenceAudios    ReferenceCapability `json:"reference_audios"`
	ReferenceVideos    ReferenceCapability `json:"reference_videos"`
}

var seedance2Capabilities = map[string]VideoModelCapabilities{
	"doubao-seedance-2.0":      seedance2ReferenceCapabilities(),
	"doubao-seedance-2.0-fast": seedance2ReferenceCapabilities(),
	"doubao-seedance-2.0-mini": seedance2ReferenceCapabilities(),
}

func seedance2ReferenceCapabilities() VideoModelCapabilities {
	return VideoModelCapabilities{
		DurationMinSeconds: 5,
		DurationMaxSeconds: 15,
		ReferenceImages:    ReferenceCapability{Supported: true, MaxCount: 9},
		ReferenceAudios:    ReferenceCapability{Supported: true, MaxCount: 3, MaxTotalDurationMS: 15_000},
		ReferenceVideos:    ReferenceCapability{Supported: true, MaxCount: 3, MinTotalDurationMS: 1_800, MaxTotalDurationMS: 15_200, MinShortSidePixels: 480, MaxShortSidePixels: 720},
	}
}

func VideoCapabilities(modelName string) (VideoModelCapabilities, bool) {
	capabilities, ok := seedance2Capabilities[strings.TrimSpace(modelName)]
	return capabilities, ok
}

func NewSeedance2(client *http.Client) Seedance2 {
	return Seedance2{Gateway: NewClient(client)}
}

func BuildSeedance2Payload(modelName string, inputData json.RawMessage) (map[string]any, error) {
	modelName = strings.TrimSpace(modelName)
	capabilities, supported := VideoCapabilities(modelName)
	if !supported {
		return nil, fmt.Errorf("暂不支持视频模型 %q", modelName)
	}
	input := map[string]any{}
	if err := json.Unmarshal(inputData, &input); err != nil {
		return nil, err
	}
	prompt := strings.TrimSpace(fmt.Sprint(input["prompt"]))
	if prompt == "" || prompt == "<nil>" {
		return nil, errors.New("生成提示词不能为空")
	}
	size := strings.TrimSpace(fmt.Sprint(input["size"]))
	if size != "16:9" && size != "9:16" {
		return nil, errors.New("项目画面比例不受 Seedance 2 支持")
	}
	duration, err := strconv.Atoi(fmt.Sprint(input["duration"]))
	if err != nil || duration < 5 || duration > 15 {
		return nil, errors.New("Seedance 2 视频时长必须为 5 到 15 秒")
	}
	images, err := seedanceURLList(input["image_urls"])
	if err != nil {
		return nil, fmt.Errorf("Seedance 2 参考图无效: %w", err)
	}
	if len(images) > capabilities.ReferenceImages.MaxCount {
		return nil, fmt.Errorf("Seedance 2 最多支持 %d 张参考图", capabilities.ReferenceImages.MaxCount)
	}
	audios, err := seedanceURLList(input["audio_urls"])
	if err != nil {
		return nil, fmt.Errorf("Seedance 2 参考音频无效: %w", err)
	}
	if len(audios) > capabilities.ReferenceAudios.MaxCount {
		return nil, fmt.Errorf("Seedance 2 最多支持 %d 个参考音频", capabilities.ReferenceAudios.MaxCount)
	}
	videos, err := seedanceURLList(input["video_urls"])
	if err != nil {
		return nil, fmt.Errorf("Seedance 2 参考视频无效: %w", err)
	}
	if len(videos) > capabilities.ReferenceVideos.MaxCount {
		return nil, fmt.Errorf("Seedance 2 最多支持 %d 个参考视频", capabilities.ReferenceVideos.MaxCount)
	}
	if len(audios) > 0 && len(images) == 0 && len(videos) == 0 {
		return nil, errors.New("Seedance 2 参考音频必须与参考图片或参考视频同时使用")
	}
	payload := map[string]any{
		"model": modelName, "prompt": prompt, "size": size, "duration": duration,
		"resolution": seedanceString(input["resolution"], "480p"),
	}
	if len(images) > 0 {
		payload["image_urls"] = images
	}
	if len(audios) > 0 {
		payload["audio_urls"] = audios
	}
	if len(videos) > 0 {
		payload["video_urls"] = videos
	}
	if value, ok := input["generate_audio"].(bool); ok {
		payload["generate_audio"] = value
	}
	if value, ok := input["return_last_frame"].(bool); ok {
		payload["return_last_frame"] = value
	}
	return payload, nil
}

func (s Seedance2) Submit(ctx context.Context, baseURL, apiKey, requestID string, payload map[string]any) (provider.SubmitResult, error) {
	data, err := s.Gateway.requestJSON(ctx, http.MethodPost, baseURL, "/videos/generations", apiKey, requestID, payload)
	if err != nil {
		return provider.SubmitResult{}, err
	}
	var decoded any
	if err := json.Unmarshal(data, &decoded); err != nil {
		return provider.SubmitResult{}, errors.New("Seedance 2 提交响应不是有效 JSON")
	}
	result := provider.SubmitResult{TaskID: findString(decoded, "task_id", "id"), URL: findSeedanceVideoURL(decoded), Raw: append(json.RawMessage(nil), data...)}
	if result.TaskID == "" && result.URL == "" {
		return provider.SubmitResult{}, errors.New("Seedance 2 提交响应缺少任务 ID 或视频地址")
	}
	return result, nil
}

func (s Seedance2) Poll(ctx context.Context, baseURL, apiKey, taskID string) (provider.PollResult, error) {
	data, err := s.Gateway.requestJSON(ctx, http.MethodGet, baseURL, "/tasks/"+taskID, apiKey, "", nil)
	if err != nil {
		return provider.PollResult{}, err
	}
	var decoded any
	if err := json.Unmarshal(data, &decoded); err != nil {
		return provider.PollResult{}, errors.New("Seedance 2 轮询响应不是有效 JSON")
	}
	status := strings.ToLower(findString(decoded, "task_status", "status", "state"))
	if status == "" {
		status = "processing"
	}
	return provider.PollResult{
		Status: status, URL: findSeedanceVideoURL(decoded),
		Error: findString(decoded, "error_message", "message", "detail"), Raw: append(json.RawMessage(nil), data...),
	}, nil
}

func seedanceURLList(value any) ([]string, error) {
	if value == nil {
		return nil, nil
	}
	raw, ok := value.([]any)
	if !ok {
		return nil, errors.New("必须是 URL 数组")
	}
	result := make([]string, 0, len(raw))
	for _, item := range raw {
		url := strings.TrimSpace(fmt.Sprint(item))
		if !strings.HasPrefix(url, "https://") && !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "asset://") {
			return nil, errors.New("包含不可访问的 URL")
		}
		result = append(result, url)
	}
	return result, nil
}

func seedanceString(value any, fallback string) string {
	text := strings.TrimSpace(fmt.Sprint(value))
	if text == "" || text == "<nil>" {
		return fallback
	}
	return text
}

func findSeedanceVideoURL(value any) string {
	if url := findURLInSeedanceVideoField(value); url != "" {
		return url
	}
	return findURL(value)
}

func findURLInSeedanceVideoField(value any) string {
	switch data := value.(type) {
	case map[string]any:
		for key, child := range data {
			if strings.Contains(strings.ToLower(key), "video") {
				if text, ok := child.(string); ok && (strings.HasPrefix(text, "https://") || strings.HasPrefix(text, "http://")) {
					return text
				}
				if url := findURL(child); url != "" {
					return url
				}
			}
		}
		for _, child := range data {
			if url := findURLInSeedanceVideoField(child); url != "" {
				return url
			}
		}
	case []any:
		for _, child := range data {
			if url := findURLInSeedanceVideoField(child); url != "" {
				return url
			}
		}
	}
	return ""
}