new-api/controller/image_studio.go

package controller

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/service"
	"github.com/QuantumNous/new-api/setting"
	"github.com/gin-gonic/gin"
)

const (
	imageStudioGroup             = "image"
	imageStudioTaskStatusPending = "pending"
	imageStudioTaskStatusRunning = "running"
	imageStudioTaskStatusSuccess = "success"
	imageStudioTaskStatusFailure = "failure"
	imageStudioTaskCleanupEvery  = 5 * time.Minute
	imageStudioHTTPTimeout       = 10 * time.Minute
	imageStudioGenerateAction    = "generate"
	imageStudioEditAction        = "edit"
	imageStudioResultBaseDir     = "/data/image-studio-results"
	imageStudioResultURLPrefix   = "/image-studio-results"
)

type imageStudioGenerateRequest struct {
	TokenID int    `json:"token_id"`
	Model   string `json:"model"`
	Prompt  string `json:"prompt"`
	Size    string `json:"size"`
	Quality string `json:"quality"`
}

type imageStudioUploadedImage struct {
	Filename    string
	ContentType string
	Data        []byte
}

type imageStudioTask struct {
	ID        string          `json:"task_id"`
	UserID    int             `json:"-"`
	Mode      string          `json:"mode"`
	Model     string          `json:"model"`
	Size      string          `json:"size"`
	Quality   string          `json:"quality"`
	Status    string          `json:"status"`
	Error     string          `json:"error,omitempty"`
	Result    json.RawMessage `json:"result,omitempty"`
	ResultDir string          `json:"-"`
	CreatedAt int64           `json:"created_at"`
	UpdatedAt int64           `json:"updated_at"`
}

type imageStudioModelSetting struct {
	Model            string   `json:"model"`
	Enabled          bool     `json:"enabled"`
	AllowGenerate    bool     `json:"allow_generate"`
	AllowEdit        bool     `json:"allow_edit"`
	AllowRatio       bool     `json:"allow_ratio"`
	AllowResolution  bool     `json:"allow_resolution"`
	AllowQuality     bool     `json:"allow_quality"`
	AllowedSizes     []string `json:"allowed_sizes,omitempty"`
	AllowedQualities []string `json:"allowed_qualities,omitempty"`
	DefaultSize      string   `json:"default_size,omitempty"`
	DefaultQuality   string   `json:"default_quality,omitempty"`
}

type imageStudioSettingsPayload struct {
	Enabled           bool                      `json:"enabled"`
	DefaultModel      string                    `json:"default_model,omitempty"`
	AllowOpenOriginal bool                      `json:"allow_open_original"`
	RetentionMinutes  int                       `json:"retention_minutes"`
	AllowedModels     []string                  `json:"allowed_models,omitempty"`
	ModelSettings     []imageStudioModelSetting `json:"model_settings,omitempty"`
}

var imageStudioSizeRatios = map[string]string{
	"auto":      "auto",
	"1024x1024": "1:1",
	"1536x1024": "3:2",
	"1024x1536": "2:3",
	"2048x2048": "1:1",
	"2048x1152": "16:9",
	"1152x2048": "9:16",
	"3840x2160": "16:9",
	"2160x3840": "9:16",
}

var (
	imageStudioTasks = struct {
		sync.RWMutex
		items map[string]*imageStudioTask
	}{items: make(map[string]*imageStudioTask)}
	imageStudioCleanupOnce sync.Once
	imageStudioSettingsCache = struct {
		sync.RWMutex
		raw      string
		settings imageStudioSettingsPayload
	}{}
	imageStudioTTLCache = struct {
		sync.RWMutex
		raw string
		ttl time.Duration
	}{}
)

func getImageStudioTaskTTL() time.Duration {
	if raw, ok := common.OptionMap["ImageStudioRetentionMinutes"]; ok {
		normalized := strings.TrimSpace(raw)
		imageStudioTTLCache.RLock()
		if imageStudioTTLCache.raw == normalized && imageStudioTTLCache.ttl > 0 {
			ttl := imageStudioTTLCache.ttl
			imageStudioTTLCache.RUnlock()
			return ttl
		}
		imageStudioTTLCache.RUnlock()
		minutes := 30
		if parsed, err := strconv.Atoi(normalized); err == nil && parsed > 0 {
			minutes = parsed
		}
		ttl := time.Duration(minutes) * time.Minute
		imageStudioTTLCache.Lock()
		imageStudioTTLCache.raw = normalized
		imageStudioTTLCache.ttl = ttl
		imageStudioTTLCache.Unlock()
		return ttl
	}
	return 30 * time.Minute
}

func normalizeImageStudioSizeList(values []string) []string {
	if len(values) == 0 {
		return nil
	}
	seen := make(map[string]struct{}, len(values))
	result := make([]string, 0, len(values))
	for _, item := range values {
		size := normalizeImageStudioSize(item)
		if size == "" {
			continue
		}
		if _, ok := seen[size]; ok {
			continue
		}
		seen[size] = struct{}{}
		result = append(result, size)
	}
	if len(result) == 0 {
		return nil
	}
	return result
}

func normalizeImageStudioQualityList(values []string) []string {
	if len(values) == 0 {
		return nil
	}
	seen := make(map[string]struct{}, len(values))
	result := make([]string, 0, len(values))
	for _, item := range values {
		quality := normalizeImageStudioQuality(item)
		if quality == "" {
			continue
		}
		if _, ok := seen[quality]; ok {
			continue
		}
		seen[quality] = struct{}{}
		result = append(result, quality)
	}
	if len(result) == 0 {
		return nil
	}
	return result
}

func normalizeImageStudioModelSetting(item imageStudioModelSetting) imageStudioModelSetting {
	item.Model = strings.TrimSpace(item.Model)
	item.AllowedSizes = normalizeImageStudioSizeList(item.AllowedSizes)
	item.AllowedQualities = normalizeImageStudioQualityList(item.AllowedQualities)
	item.DefaultSize = normalizeImageStudioSize(item.DefaultSize)
	item.DefaultQuality = normalizeImageStudioQuality(item.DefaultQuality)
	if item.DefaultSize == "" {
		item.DefaultSize = "auto"
	}
	if item.DefaultQuality == "" {
		item.DefaultQuality = "auto"
	}
	if !item.AllowGenerate && !item.AllowEdit {
		item.AllowGenerate = true
	}
	if len(item.AllowedSizes) == 0 {
		item.AllowedSizes = []string{"auto"}
	}
	if len(item.AllowedQualities) == 0 {
		item.AllowedQualities = []string{"auto", "low", "medium", "high"}
	}
	return item
}

func getDefaultImageStudioSettings() imageStudioSettingsPayload {
	return imageStudioSettingsPayload{
		Enabled:           true,
		DefaultModel:      "gpt-image-2",
		AllowOpenOriginal: true,
		RetentionMinutes:  30,
	}
}

func getImageStudioSettings() imageStudioSettingsPayload {
	raw, _ := common.OptionMap["ImageStudioSettings"]
	raw = strings.TrimSpace(raw)

	imageStudioSettingsCache.RLock()
	if imageStudioSettingsCache.raw == raw {
		cached := imageStudioSettingsCache.settings
		imageStudioSettingsCache.RUnlock()
		return cached
	}
	imageStudioSettingsCache.RUnlock()

	settings := getDefaultImageStudioSettings()
	if raw == "" {
		imageStudioSettingsCache.Lock()
		imageStudioSettingsCache.raw = raw
		imageStudioSettingsCache.settings = settings
		imageStudioSettingsCache.Unlock()
		return settings
	}
	var parsed imageStudioSettingsPayload
	if err := json.Unmarshal([]byte(raw), &parsed); err != nil {
		imageStudioSettingsCache.Lock()
		imageStudioSettingsCache.raw = raw
		imageStudioSettingsCache.settings = settings
		imageStudioSettingsCache.Unlock()
		return settings
	}
	settings.Enabled = parsed.Enabled
	if strings.TrimSpace(parsed.DefaultModel) != "" {
		settings.DefaultModel = strings.TrimSpace(parsed.DefaultModel)
	}
	settings.AllowOpenOriginal = parsed.AllowOpenOriginal
	if parsed.RetentionMinutes > 0 {
		settings.RetentionMinutes = parsed.RetentionMinutes
	}
	if len(parsed.AllowedModels) > 0 {
		seen := make(map[string]struct{}, len(parsed.AllowedModels))
		settings.AllowedModels = make([]string, 0, len(parsed.AllowedModels))
		for _, item := range parsed.AllowedModels {
			modelName := strings.TrimSpace(item)
			if modelName == "" {
				continue
			}
			if _, ok := seen[modelName]; ok {
				continue
			}
			seen[modelName] = struct{}{}
			settings.AllowedModels = append(settings.AllowedModels, modelName)
		}
	}
	if len(parsed.ModelSettings) > 0 {
		seen := make(map[string]struct{}, len(parsed.ModelSettings))
		settings.ModelSettings = make([]imageStudioModelSetting, 0, len(parsed.ModelSettings))
		for _, item := range parsed.ModelSettings {
			normalized := normalizeImageStudioModelSetting(item)
			if normalized.Model == "" {
				continue
			}
			if _, ok := seen[normalized.Model]; ok {
				continue
			}
			seen[normalized.Model] = struct{}{}
			settings.ModelSettings = append(settings.ModelSettings, normalized)
		}
	}
	imageStudioSettingsCache.Lock()
	imageStudioSettingsCache.raw = raw
	imageStudioSettingsCache.settings = settings
	imageStudioSettingsCache.Unlock()
	return settings
}

func getImageStudioModelSetting(modelName string) (imageStudioModelSetting, bool) {
	settings := getImageStudioSettings()
	for _, item := range settings.ModelSettings {
		if strings.EqualFold(strings.TrimSpace(item.Model), strings.TrimSpace(modelName)) {
			return item, true
		}
	}
	return imageStudioModelSetting{}, false
}

func shouldAllowImageStudioModel(modelName string) bool {
	settings := getImageStudioSettings()
	if len(settings.AllowedModels) == 0 {
		return true
	}
	for _, item := range settings.AllowedModels {
		if strings.EqualFold(strings.TrimSpace(item), strings.TrimSpace(modelName)) {
			return true
		}
	}
	return false
}

func getImageStudioAllowedSizesForModel(modelName string) []string {
	if setting, ok := getImageStudioModelSetting(modelName); ok && len(setting.AllowedSizes) > 0 {
		return setting.AllowedSizes
	}
	return nil
}

func getImageStudioAllowedQualitiesForModel(modelName string) []string {
	if setting, ok := getImageStudioModelSetting(modelName); ok && len(setting.AllowedQualities) > 0 {
		return setting.AllowedQualities
	}
	return nil
}

func buildImageStudioPublicSettings() gin.H {
	settings := getImageStudioSettings()
	pricingItems := model.GetPricing()
	allowedModels := make([]string, 0)
	modelSettings := make([]imageStudioModelSetting, 0)
	models := make([]gin.H, 0)
	seenModels := make(map[string]struct{})
	for _, item := range pricingItems {
		if !containsString(item.EnableGroup, imageStudioGroup) {
			continue
		}
		if !shouldAllowImageStudioModel(item.ModelName) {
			continue
		}
		if _, ok := seenModels[item.ModelName]; ok {
			continue
		}
		seenModels[item.ModelName] = struct{}{}
		allowedModels = append(allowedModels, item.ModelName)
		models = append(models, gin.H{
			"model_name":               item.ModelName,
			"enable_groups":            item.EnableGroup,
			"supported_endpoint_types": item.SupportedEndpointTypes,
		})
		if ms, ok := getImageStudioModelSetting(item.ModelName); ok {
			modelSettings = append(modelSettings, ms)
		}
	}
	sort.Strings(allowedModels)
	sort.SliceStable(models, func(i, j int) bool {
		return fmt.Sprintf("%v", models[i]["model_name"]) < fmt.Sprintf("%v", models[j]["model_name"])
	})
	sort.SliceStable(modelSettings, func(i, j int) bool {
		return modelSettings[i].Model < modelSettings[j].Model
	})
	return gin.H{
		"enabled":             settings.Enabled,
		"default_model":       settings.DefaultModel,
		"allow_open_original": settings.AllowOpenOriginal,
		"retention_minutes":   settings.RetentionMinutes,
		"allowed_models":      allowedModels,
		"models":              models,
		"model_settings":      modelSettings,
	}
}

func ensureImageStudioTaskCleanup() {
	imageStudioCleanupOnce.Do(func() {
		go func() {
			ticker := time.NewTicker(imageStudioTaskCleanupEvery)
			defer ticker.Stop()
			for range ticker.C {
				cutoff := time.Now().Add(-getImageStudioTaskTTL()).Unix()
				imageStudioTasks.Lock()
				for id, task := range imageStudioTasks.items {
					if task.UpdatedAt < cutoff {
						cleanupImageStudioTaskArtifacts(task)
						delete(imageStudioTasks.items, id)
					}
				}
				imageStudioTasks.Unlock()
			}
		}()
	})
}

func newImageStudioTask(userID int, mode string, req imageStudioGenerateRequest) (*imageStudioTask, error) {
	ensureImageStudioTaskCleanup()
	key, err := common.GenerateRandomCharsKey(24)
	if err != nil {
		return nil, err
	}
	now := time.Now().Unix()
	task := &imageStudioTask{
		ID:        "imgtask_" + key,
		UserID:    userID,
		Mode:      mode,
		Model:     strings.TrimSpace(req.Model),
		Size:      resolveImageStudioSize(req.Model, req.Size),
		Quality:   normalizeImageStudioQuality(req.Quality),
		Status:    imageStudioTaskStatusPending,
		CreatedAt: now,
		UpdatedAt: now,
	}
	imageStudioTasks.Lock()
	imageStudioTasks.items[task.ID] = task
	imageStudioTasks.Unlock()
	return task, nil
}

func getImageStudioTask(taskID string) (*imageStudioTask, bool) {
	imageStudioTasks.RLock()
	defer imageStudioTasks.RUnlock()
	task, ok := imageStudioTasks.items[taskID]
	return task, ok
}

func updateImageStudioTask(taskID string, updateFn func(task *imageStudioTask)) {
	imageStudioTasks.Lock()
	defer imageStudioTasks.Unlock()
	if task, ok := imageStudioTasks.items[taskID]; ok {
		updateFn(task)
		task.UpdatedAt = time.Now().Unix()
	}
}

func cleanupImageStudioTaskArtifacts(task *imageStudioTask) {
	if task == nil || strings.TrimSpace(task.ResultDir) == "" {
		return
	}
	baseDir := filepath.Clean(imageStudioResultBaseDir)
	resultDir := filepath.Clean(task.ResultDir)
	prefix := baseDir + string(os.PathSeparator)
	if resultDir == baseDir || !strings.HasPrefix(resultDir, prefix) {
		return
	}
	_ = os.RemoveAll(resultDir)
	task.ResultDir = ""
}

func normalizeImageStudioSize(v string) string {
	value := strings.TrimSpace(strings.ToLower(v))
	switch value {
	case "", "auto":
		return "auto"
	case "1024x1024", "1536x1024", "1024x1536", "2048x2048", "2048x1152", "1152x2048", "3840x2160", "2160x3840":
		return value
	default:
		return "auto"
	}
}

func normalizeImageStudioModel(v string) string {
	return strings.TrimSpace(strings.ToLower(v))
}

func isGPTImage2Model(modelName string) bool {
	return normalizeImageStudioModel(modelName) == "gpt-image-2"
}

func isValidGPTImage2Size(size string) bool {
	if size == "" || size == "auto" {
		return true
	}
	parts := strings.Split(size, "x")
	if len(parts) != 2 {
		return false
	}
	width, err := strconv.Atoi(parts[0])
	if err != nil {
		return false
	}
	height, err := strconv.Atoi(parts[1])
	if err != nil {
		return false
	}
	if width <= 0 || height <= 0 {
		return false
	}
	if width > 3840 || height > 3840 {
		return false
	}
	if width%16 != 0 || height%16 != 0 {
		return false
	}
	longEdge := width
	shortEdge := height
	if height > width {
		longEdge = height
		shortEdge = width
	}
	if longEdge > shortEdge*3 {
		return false
	}
	pixels := width * height
	if pixels < 655360 || pixels > 8294400 {
		return false
	}
	return true
}

func resolveImageStudioSize(modelName string, size string) string {
	normalizedSize := normalizeImageStudioSize(size)
	if isGPTImage2Model(modelName) {
		candidate := strings.TrimSpace(strings.ToLower(size))
		if isValidGPTImage2Size(candidate) {
			return candidate
		}
		if normalizedSize != "auto" {
			return normalizedSize
		}
	}
	return normalizedSize
}

func normalizeImageStudioQuality(v string) string {
	value := strings.TrimSpace(strings.ToLower(v))
	switch value {
	case "", "auto":
		return "auto"
	case "low", "medium", "high":
		return value
	default:
		return "auto"
	}
}

func normalizeImageStudioBearerToken(key string) string {
	trimmed := strings.TrimSpace(key)
	if strings.HasPrefix(trimmed, "sk-") {
		return trimmed
	}
	return "sk-" + trimmed
}

var imageStudioDisallowedKeywords = []string{
	"暴力",
	"血腥",
	"伤害",
	"虐待",
	"自残",
	"自杀",
	"恐怖",
	"战争",
	"武器",
	"尸体",
	"残肢",
	"恶心",
	"色情",
	"裸露",
	"性暗示",
	"卖淫",
	"约炮",
	"成人视频",
	"强奷",
	"sm",
	"性暴力",
	"性服务",
	"赌博",
	"博彩",
	"赌场",
	"下注",
	"洗钱",
	"毒品",
	"制毒",
	"吸毒",
	"贩毒",
	"违禁药物",
	"政治敏感",
	"反动",
	"分裂",
	"颠覆",
	"煽动",
	"国家领导人",
	"国旗",
	"国徽",
	"国歌",
	"恐怖主义",
	"暴恐",
}

func buildImageStudioSafetyPrompt() string {
	return setting.ImageStudioSafetyPrompt
}

func normalizeImageStudioPromptForSafety(prompt string) string {
	replacer := strings.NewReplacer(
		" ", "",
		"\n", "",
		"\r", "",
		"\t", "",
		"-", "",
		"_", "",
		"/", "",
		"\\", "",
		"|", "",
		"*", "",
		".", "",
		",", "",
		"，", "",
		"。", "",
		":", "",
		"：", "",
		";", "",
		"；", "",
		"(", "",
		")", "",
		"（", "",
		"）", "",
		"[", "",
		"]", "",
		"【", "",
		"】", "",
		"{", "",
		"}", "",
		"<", "",
		">", "",
		"《", "",
		"》", "",
		"'", "",
		"\"", "",
		"`", "",
		"!", "",
		"！", "",
		"?", "",
		"？", "",
	)
	return strings.ToLower(replacer.Replace(strings.TrimSpace(prompt)))
}

func containsImageStudioDisallowedKeyword(prompt string) (bool, string) {
	checkText := normalizeImageStudioPromptForSafety(prompt)
	for _, keyword := range setting.ImageStudioBannedWords {
		k := normalizeImageStudioPromptForSafety(keyword)
		if k == "" {
			continue
		}
		if strings.Contains(checkText, k) {
			return true, keyword
		}
	}
	return false, ""
}

func validateImageStudioPromptSafety(prompt string) error {
	trimmed := strings.TrimSpace(prompt)
	if trimmed == "" {
		return fmt.Errorf("请先填写提示词")
	}
	if ok, keyword := containsImageStudioDisallowedKeyword(trimmed); ok {
		return fmt.Errorf("提示词包含禁止内容：%s", keyword)
	}
	if ok, words := service.SensitiveWordContains(trimmed); ok {
		if len(words) > 0 {
			return fmt.Errorf("提示词包含敏感内容：%s", strings.Join(words, ", "))
		}
		return fmt.Errorf("提示词包含敏感内容，请修改后重试")
	}
	return nil
}

func imageStudioModeEndpointType(mode string) constant.EndpointType {
	if mode == imageStudioEditAction {
		return constant.EndpointTypeImageEdit
	}
	return constant.EndpointTypeImageGeneration
}

func imageStudioModePath(mode string) string {
	if mode == imageStudioEditAction {
		return "/v1/images/edits"
	}
	return "/v1/images/generations"
}

func resolveImageStudioGenerateRelayAction(images []imageStudioUploadedImage) string {
	if len(images) > 0 {
		return imageStudioEditAction
	}
	return imageStudioGenerateAction
}

func validateImageStudioTokenAndModel(userID int, req imageStudioGenerateRequest, mode string) (*model.Token, error) {
	token, err := model.GetTokenByIds(req.TokenID, userID)
	if err != nil {
		return nil, fmt.Errorf("未找到所选令牌，请刷新后重试")
	}
	if token.Status != common.TokenStatusEnabled {
		return nil, fmt.Errorf("当前令牌状态不可用，请去控制台-令牌管理检查该令牌")
	}
	if strings.TrimSpace(token.Group) != imageStudioGroup {
		return nil, fmt.Errorf("所选令牌不是 image 分组，请去控制台-令牌管理创建或选择 image 分组令牌")
	}
	modelName := strings.TrimSpace(req.Model)
	if modelName == "" {
		return nil, fmt.Errorf("请选择图片模型")
	}
	settings := getImageStudioSettings()
	if !settings.Enabled {
		return nil, fmt.Errorf("图片制作功能暂未开放")
	}
	if !shouldAllowImageStudioModel(modelName) {
		return nil, fmt.Errorf("当前模型暂未开放图片制作，请联系管理员调整设置")
	}
	if modelSetting, ok := getImageStudioModelSetting(modelName); ok {
		if !modelSetting.Enabled {
			return nil, fmt.Errorf("当前模型已被管理员禁用")
		}
		if mode == imageStudioEditAction && !modelSetting.AllowEdit {
			return nil, fmt.Errorf("当前模型未开放图片编辑")
		}
		if mode != imageStudioEditAction && !modelSetting.AllowGenerate {
			return nil, fmt.Errorf("当前模型未开放图片生成")
		}
	}
	if allowedSizes := getImageStudioAllowedSizesForModel(modelName); len(allowedSizes) > 0 {
		resolvedSize := resolveImageStudioSize(modelName, req.Size)
		allowed := false
		for _, size := range allowedSizes {
			if size == resolvedSize {
				allowed = true
				break
			}
		}
		if !allowed {
			return nil, fmt.Errorf("当前模型不允许使用所选尺寸")
		}
	}
	if allowedQualities := getImageStudioAllowedQualitiesForModel(modelName); len(allowedQualities) > 0 {
		resolvedQuality := normalizeImageStudioQuality(req.Quality)
		allowed := false
		for _, quality := range allowedQualities {
			if quality == resolvedQuality {
				allowed = true
				break
			}
		}
		if !allowed {
			return nil, fmt.Errorf("当前模型不允许使用所选质量")
		}
	}
	if isGPTImage2Model(modelName) && !isValidGPTImage2Size(strings.TrimSpace(strings.ToLower(req.Size))) {
		return nil, fmt.Errorf("gpt-image-2 尺寸不合法：最长边不能超过 3840，边长需为 16 的倍数，长宽比不能超过 3:1，总像素需在 655360 到 8294400 之间")
	}
	if token.ModelLimitsEnabled {
		limits := token.GetModelLimitsMap()
		if len(limits) == 0 || !limits[modelName] {
			return nil, fmt.Errorf("当前令牌没有放行该图片模型，请去控制台-令牌管理调整模型限制")
		}
	}

	pricingItems := model.GetPricing()
	var matched *model.Pricing
	for index := range pricingItems {
		item := pricingItems[index]
		if item.ModelName != modelName {
			continue
		}
		if !containsString(item.EnableGroup, imageStudioGroup) {
			continue
		}
		matched = &item
		break
	}
	if matched == nil {
		return nil, fmt.Errorf("当前 image 分组未配置该图片模型，请联系管理员检查 image 分组和模型配置")
	}
	if !containsEndpointType(matched.SupportedEndpointTypes, imageStudioModeEndpointType(mode)) {
		if mode == imageStudioEditAction {
			return nil, fmt.Errorf("当前模型不支持编辑图片，请联系管理员检查模型能力配置")
		}
		return nil, fmt.Errorf("当前模型不支持生成图片，请联系管理员检查模型能力配置")
	}
	return token, nil
}

func containsString(items []string, target string) bool {
	for _, item := range items {
		if strings.TrimSpace(item) == target {
			return true
		}
	}
	return false
}

func containsEndpointType(items []constant.EndpointType, target constant.EndpointType) bool {
	for _, item := range items {
		if item == target {
			return true
		}
	}
	return false
}

func readImageStudioUploads(c *gin.Context, required bool) ([]imageStudioUploadedImage, error) {
	form, err := c.MultipartForm()
	if err != nil {
		return nil, err
	}
	fileHeaders := form.File["image[]"]
	if len(fileHeaders) == 0 {
		fileHeaders = form.File["images[]"]
	}
	if len(fileHeaders) == 0 {
		fileHeaders = form.File["image"]
	}
	if len(fileHeaders) == 0 {
		if required {
			return nil, fmt.Errorf("编辑图片至少需要上传一张参考图")
		}
		return nil, nil
	}
	images := make([]imageStudioUploadedImage, 0, len(fileHeaders))
	for _, header := range fileHeaders {
		image, err := readImageStudioUpload(header)
		if err != nil {
			return nil, err
		}
		images = append(images, image)
	}
	return images, nil
}

func readImageStudioUpload(header *multipart.FileHeader) (imageStudioUploadedImage, error) {
	file, err := header.Open()
	if err != nil {
		return imageStudioUploadedImage{}, err
	}
	defer file.Close()
	data, err := io.ReadAll(file)
	if err != nil {
		return imageStudioUploadedImage{}, err
	}
	contentType := header.Header.Get("Content-Type")
	if contentType == "" {
		contentType = http.DetectContentType(data)
	}
	return imageStudioUploadedImage{
		Filename:    header.Filename,
		ContentType: contentType,
		Data:        data,
	}, nil
}

func buildImageStudioInternalURL(mode string) string {
	return fmt.Sprintf("http://127.0.0.1:%d%s", *common.Port, imageStudioModePath(mode))
}

func decodeImageStudioError(body []byte, fallback string) string {
	if len(body) == 0 {
		return fallback
	}
	var payload map[string]any
	if err := json.Unmarshal(body, &payload); err == nil {
		if errorPart, ok := payload["error"].(map[string]any); ok {
			if message, ok := errorPart["message"].(string); ok && strings.TrimSpace(message) != "" {
				return message
			}
		}
		if message, ok := payload["message"].(string); ok && strings.TrimSpace(message) != "" {
			return message
		}
	}
	return fallback
}

func getImageStudioResultExtension(outputFormat string) string {
	switch strings.ToLower(strings.TrimSpace(outputFormat)) {
	case "jpg", "jpeg":
		return "jpg"
	case "webp":
		return "webp"
	default:
		return "png"
	}
}

func signImageStudioResult(relativePath string, expiresAt int64) string {
	mac := hmac.New(sha256.New, []byte(common.CryptoSecret))
	_, _ = mac.Write([]byte(relativePath))
	_, _ = mac.Write([]byte("\n"))
	_, _ = mac.Write([]byte(strconv.FormatInt(expiresAt, 10)))
	return hex.EncodeToString(mac.Sum(nil))
}

func buildImageStudioResultFileURL(taskID string, index int, outputFormat string) string {
	relativePath := fmt.Sprintf("%s/%d.%s", taskID, index, getImageStudioResultExtension(outputFormat))
	expiresAt := time.Now().Add(getImageStudioTaskTTL()).Unix()
	signature := signImageStudioResult(relativePath, expiresAt)
	return fmt.Sprintf("%s/%s?expires=%d&sig=%s", imageStudioResultURLPrefix, relativePath, expiresAt, signature)
}

func persistImageStudioTaskResult(taskID string, respBody []byte) (json.RawMessage, string, error) {
	var payload map[string]any
	if err := json.Unmarshal(respBody, &payload); err != nil {
		return nil, "", err
	}
	data, ok := payload["data"].([]any)
	if !ok || len(data) == 0 {
		return json.RawMessage(respBody), "", nil
	}
	resultDir := filepath.Join(imageStudioResultBaseDir, taskID)
	wroteFiles := false
	for index, rawItem := range data {
		item, ok := rawItem.(map[string]any)
		if !ok {
			continue
		}
		b64, _ := item["b64_json"].(string)
		b64 = strings.TrimSpace(b64)
		if b64 == "" {
			continue
		}
		decoded, err := base64.StdEncoding.DecodeString(b64)
		if err != nil {
			_ = os.RemoveAll(resultDir)
			return nil, "", fmt.Errorf("图片结果解码失败")
		}
		if !wroteFiles {
			_ = os.RemoveAll(resultDir)
			if err := os.MkdirAll(resultDir, 0o755); err != nil {
				return nil, "", fmt.Errorf("图片结果目录创建失败")
			}
		}
		outputFormat, _ := item["output_format"].(string)
		ext := getImageStudioResultExtension(outputFormat)
		filename := fmt.Sprintf("%d.%s", index, ext)
		fullPath := filepath.Join(resultDir, filename)
		if err := os.WriteFile(fullPath, decoded, 0o644); err != nil {
			_ = os.RemoveAll(resultDir)
			return nil, "", fmt.Errorf("图片结果写入失败")
		}
		delete(item, "b64_json")
		item["url"] = buildImageStudioResultFileURL(taskID, index, outputFormat)
		wroteFiles = true
	}
	if !wroteFiles {
		return json.RawMessage(respBody), "", nil
	}
	encoded, err := json.Marshal(payload)
	if err != nil {
		_ = os.RemoveAll(resultDir)
		return nil, "", fmt.Errorf("图片结果序列化失败")
	}
	return json.RawMessage(encoded), resultDir, nil
}

func executeImageStudioGenerateTask(taskID string, tokenKey string, req imageStudioGenerateRequest, images []imageStudioUploadedImage) {
	updateImageStudioTask(taskID, func(task *imageStudioTask) {
		task.Status = imageStudioTaskStatusRunning
		task.Error = ""
	})

	relayAction := resolveImageStudioGenerateRelayAction(images)
	var body bytes.Buffer
	contentType := "application/json"
	if len(images) > 0 {
		writer := multipart.NewWriter(&body)
		_ = writer.WriteField("model", strings.TrimSpace(req.Model))
		_ = writer.WriteField("prompt", buildImageStudioSafetyPrompt()+"\n\n"+strings.TrimSpace(req.Prompt))
		_ = writer.WriteField("size", resolveImageStudioSize(req.Model, req.Size))
		_ = writer.WriteField("quality", normalizeImageStudioQuality(req.Quality))
		_ = writer.WriteField("response_format", "b64_json")
		_ = writer.WriteField("output_format", "png")

		for _, image := range images {
			headers := make(textproto.MIMEHeader)
			headers.Set("Content-Disposition", fmt.Sprintf(`form-data; name="image[]"; filename="%s"`, image.Filename))
			headers.Set("Content-Type", image.ContentType)

			part, err := writer.CreatePart(headers)
			if err != nil {
				updateImageStudioTask(taskID, func(task *imageStudioTask) {
					task.Status = imageStudioTaskStatusFailure
					task.Error = "图片生成参考图封装失败"
				})
				_ = writer.Close()
				return
			}
			if _, err = part.Write(image.Data); err != nil {
				updateImageStudioTask(taskID, func(task *imageStudioTask) {
					task.Status = imageStudioTaskStatusFailure
					task.Error = "图片生成参考图写入失败"
				})
				_ = writer.Close()
				return
			}
		}
		if err := writer.Close(); err != nil {
			updateImageStudioTask(taskID, func(task *imageStudioTask) {
				task.Status = imageStudioTaskStatusFailure
				task.Error = "图片生成请求构建失败"
			})
			return
		}
		contentType = writer.FormDataContentType()
	} else {
		payload, err := json.Marshal(gin.H{
			"model":           strings.TrimSpace(req.Model),
			"prompt":          buildImageStudioSafetyPrompt() + "\n\n" + strings.TrimSpace(req.Prompt),
			"size":            resolveImageStudioSize(req.Model, req.Size),
			"quality":         normalizeImageStudioQuality(req.Quality),
			"response_format": "b64_json",
			"output_format":   "png",
		})
		if err != nil {
			updateImageStudioTask(taskID, func(task *imageStudioTask) {
				task.Status = imageStudioTaskStatusFailure
				task.Error = "图片生成请求构建失败"
			})
			return
		}
		body.Write(payload)
	}

	httpReq, err := http.NewRequest(http.MethodPost, buildImageStudioInternalURL(relayAction), bytes.NewReader(body.Bytes()))
	if err != nil {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = "图片生成请求创建失败"
		})
		return
	}
	httpReq.Header.Set("Authorization", "Bearer "+normalizeImageStudioBearerToken(tokenKey))
	httpReq.Header.Set("Content-Type", contentType)

	client := &http.Client{Timeout: imageStudioHTTPTimeout}
	resp, err := client.Do(httpReq)
	if err != nil {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = "图片生成请求失败：" + err.Error()
		})
		return
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = "读取图片生成结果失败"
		})
		return
	}
	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = decodeImageStudioError(respBody, "图片生成失败")
		})
		return
	}

	sanitizedResult, resultDir, err := persistImageStudioTaskResult(taskID, respBody)
	if err != nil {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = err.Error()
		})
		return
	}

	updateImageStudioTask(taskID, func(task *imageStudioTask) {
		task.Status = imageStudioTaskStatusSuccess
		cleanupImageStudioTaskArtifacts(task)
		task.Result = sanitizedResult
		task.ResultDir = resultDir
		task.Error = ""
	})
}

func executeImageStudioEditTask(taskID string, tokenKey string, req imageStudioGenerateRequest, images []imageStudioUploadedImage) {
	updateImageStudioTask(taskID, func(task *imageStudioTask) {
		task.Status = imageStudioTaskStatusRunning
		task.Error = ""
	})

	var body bytes.Buffer
	writer := multipart.NewWriter(&body)
	_ = writer.WriteField("model", strings.TrimSpace(req.Model))
	_ = writer.WriteField("prompt", buildImageStudioSafetyPrompt()+"\n\n"+strings.TrimSpace(req.Prompt))
	_ = writer.WriteField("size", resolveImageStudioSize(req.Model, req.Size))
	_ = writer.WriteField("quality", normalizeImageStudioQuality(req.Quality))
	_ = writer.WriteField("response_format", "b64_json")
	_ = writer.WriteField("output_format", "png")

	for _, image := range images {
		headers := make(textproto.MIMEHeader)
		headers.Set("Content-Disposition", fmt.Sprintf(`form-data; name="image[]"; filename="%s"`, image.Filename))
		headers.Set("Content-Type", image.ContentType)

		part, err := writer.CreatePart(headers)
		if err != nil {
			updateImageStudioTask(taskID, func(task *imageStudioTask) {
				task.Status = imageStudioTaskStatusFailure
				task.Error = "编辑图片参考图封装失败"
			})
			_ = writer.Close()
			return
		}
		if _, err = part.Write(image.Data); err != nil {
			updateImageStudioTask(taskID, func(task *imageStudioTask) {
				task.Status = imageStudioTaskStatusFailure
				task.Error = "编辑图片参考图写入失败"
			})
			_ = writer.Close()
			return
		}
	}
	if err := writer.Close(); err != nil {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = "编辑图片请求构建失败"
		})
		return
	}

	httpReq, err := http.NewRequest(http.MethodPost, buildImageStudioInternalURL(imageStudioEditAction), bytes.NewReader(body.Bytes()))
	if err != nil {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = "编辑图片请求创建失败"
		})
		return
	}
	httpReq.Header.Set("Authorization", "Bearer "+normalizeImageStudioBearerToken(tokenKey))
	httpReq.Header.Set("Content-Type", writer.FormDataContentType())

	client := &http.Client{Timeout: imageStudioHTTPTimeout}
	resp, err := client.Do(httpReq)
	if err != nil {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = "编辑图片请求失败：" + err.Error()
		})
		return
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = "读取编辑图片结果失败"
		})
		return
	}
	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = decodeImageStudioError(respBody, "图片编辑失败")
		})
		return
	}

	sanitizedResult, resultDir, err := persistImageStudioTaskResult(taskID, respBody)
	if err != nil {
		updateImageStudioTask(taskID, func(task *imageStudioTask) {
			task.Status = imageStudioTaskStatusFailure
			task.Error = err.Error()
		})
		return
	}

	updateImageStudioTask(taskID, func(task *imageStudioTask) {
		task.Status = imageStudioTaskStatusSuccess
		cleanupImageStudioTaskArtifacts(task)
		task.Result = sanitizedResult
		task.ResultDir = resultDir
		task.Error = ""
	})
}

func SubmitImageStudioGenerate(c *gin.Context) {
	userID := c.GetInt("id")
	var req imageStudioGenerateRequest
	var images []imageStudioUploadedImage
	var err error

	if strings.Contains(strings.ToLower(c.GetHeader("Content-Type")), "multipart/form-data") {
		req = imageStudioGenerateRequest{
			Model:   c.PostForm("model"),
			Prompt:  strings.TrimSpace(c.PostForm("prompt")),
			Size:    c.PostForm("size"),
			Quality: c.PostForm("quality"),
		}
		if tokenID := strings.TrimSpace(c.PostForm("token_id")); tokenID != "" {
			_, _ = fmt.Sscanf(tokenID, "%d", &req.TokenID)
		}
		images, err = readImageStudioUploads(c, false)
		if err != nil {
			common.ApiErrorMsg(c, "图片生成参考图读取失败，请重新上传后重试")
			return
		}
	} else if err = c.ShouldBindJSON(&req); err != nil {
		common.ApiErrorMsg(c, "图片生成参数格式不正确，请刷新后重试")
		return
	}

	req.Prompt = strings.TrimSpace(req.Prompt)
	if err := validateImageStudioPromptSafety(req.Prompt); err != nil {
		common.ApiErrorMsg(c, err.Error())
		return
	}
	validationMode := imageStudioGenerateAction
	if len(images) > 0 {
		validationMode = imageStudioEditAction
	}
	token, err := validateImageStudioTokenAndModel(userID, req, validationMode)
	if err != nil {
		common.ApiErrorMsg(c, err.Error())
		return
	}
	task, err := newImageStudioTask(userID, imageStudioGenerateAction, req)
	if err != nil {
		common.ApiErrorMsg(c, "创建图片任务失败，请稍后重试")
		return
	}
	go executeImageStudioGenerateTask(task.ID, token.GetFullKey(), req, images)
	common.ApiSuccess(c, task)
}

func SubmitImageStudioEdit(c *gin.Context) {
	userID := c.GetInt("id")
	req := imageStudioGenerateRequest{
		Model:   c.PostForm("model"),
		Prompt:  strings.TrimSpace(c.PostForm("prompt")),
		Size:    c.PostForm("size"),
		Quality: c.PostForm("quality"),
	}
	if tokenID := strings.TrimSpace(c.PostForm("token_id")); tokenID != "" {
		_, _ = fmt.Sscanf(tokenID, "%d", &req.TokenID)
	}
	if err := validateImageStudioPromptSafety(req.Prompt); err != nil {
		common.ApiErrorMsg(c, err.Error())
		return
	}
	images, err := readImageStudioUploads(c, true)
	if err != nil {
		common.ApiErrorMsg(c, err.Error())
		return
	}
	token, err := validateImageStudioTokenAndModel(userID, req, imageStudioEditAction)
	if err != nil {
		common.ApiErrorMsg(c, err.Error())
		return
	}
	task, err := newImageStudioTask(userID, imageStudioEditAction, req)
	if err != nil {
		common.ApiErrorMsg(c, "创建图片任务失败，请稍后重试")
		return
	}
	go executeImageStudioEditTask(task.ID, token.GetFullKey(), req, images)
	common.ApiSuccess(c, task)
}

func GetImageStudioSettings(c *gin.Context) {
	common.ApiSuccess(c, buildImageStudioPublicSettings())
}

func buildImageStudioTaskResultURL(taskID string, index int) string {
	return fmt.Sprintf("/api/image-studio/task/%s/result/%d", taskID, index)
}

func getImageStudioResultContentType(outputFormat string) string {
	switch strings.ToLower(strings.TrimSpace(outputFormat)) {
	case "jpg", "jpeg":
		return "image/jpeg"
	case "webp":
		return "image/webp"
	default:
		return "image/png"
	}
}

func sanitizeImageStudioTaskForResponse(task *imageStudioTask) *imageStudioTask {
	if task == nil {
		return nil
	}
	clone := *task
	if clone.Status != imageStudioTaskStatusSuccess || len(clone.Result) == 0 {
		return &clone
	}
	var payload map[string]any
	if err := json.Unmarshal(clone.Result, &payload); err != nil {
		return &clone
	}
	data, ok := payload["data"].([]any)
	if !ok || len(data) == 0 {
		return &clone
	}
	changed := false
	for index, rawItem := range data {
		item, ok := rawItem.(map[string]any)
		if !ok {
			continue
		}
		taskResultURL := buildImageStudioTaskResultURL(task.ID, index)
		if b64, _ := item["b64_json"].(string); strings.TrimSpace(b64) != "" {
			delete(item, "b64_json")
			changed = true
		}
		if url, _ := item["url"].(string); strings.TrimSpace(url) != "" && strings.TrimSpace(url) != taskResultURL {
			item["raw_url"] = url
			changed = true
		}
		item["url"] = taskResultURL
		item["task_result_url"] = taskResultURL
		changed = true
	}
	if !changed {
		return &clone
	}
	if encoded, err := json.Marshal(payload); err == nil {
		clone.Result = json.RawMessage(encoded)
	}
	return &clone
}

func GetImageStudioTask(c *gin.Context) {
	userID := c.GetInt("id")
	taskID := strings.TrimSpace(c.Param("id"))
	if taskID == "" {
		common.ApiErrorMsg(c, "任务不存在")
		return
	}
	task, ok := getImageStudioTask(taskID)
	if !ok || task == nil || task.UserID != userID {
		common.ApiErrorMsg(c, "任务不存在或已过期，请重新提交")
		return
	}
	common.ApiSuccess(c, sanitizeImageStudioTaskForResponse(task))
}

func GetImageStudioTaskResult(c *gin.Context) {
	userID := c.GetInt("id")
	taskID := strings.TrimSpace(c.Param("id"))
	if taskID == "" {
		common.ApiErrorMsg(c, "任务不存在")
		return
	}
	index, err := strconv.Atoi(strings.TrimSpace(c.Param("index")))
	if err != nil || index < 0 {
		common.ApiErrorMsg(c, "图片结果不存在")
		return
	}
	task, ok := getImageStudioTask(taskID)
	if !ok || task == nil || task.UserID != userID {
		common.ApiErrorMsg(c, "任务不存在或已过期，请重新提交")
		return
	}
	if task.Status != imageStudioTaskStatusSuccess || len(task.Result) == 0 {
		common.ApiErrorMsg(c, "图片结果尚未生成完成")
		return
	}
	var payload map[string]any
	if err := json.Unmarshal(task.Result, &payload); err != nil {
		common.ApiErrorMsg(c, "图片结果解析失败")
		return
	}
	data, ok := payload["data"].([]any)
	if !ok || index >= len(data) {
		common.ApiErrorMsg(c, "图片结果不存在")
		return
	}
	item, ok := data[index].(map[string]any)
	if !ok {
		common.ApiErrorMsg(c, "图片结果不存在")
		return
	}

	if strings.TrimSpace(task.ResultDir) != "" {
		outputFormat, _ := item["output_format"].(string)
		filename := fmt.Sprintf("%d.%s", index, getImageStudioResultExtension(outputFormat))
		fullPath := filepath.Clean(filepath.Join(task.ResultDir, filename))
		resultDir := filepath.Clean(task.ResultDir)
		prefix := resultDir + string(os.PathSeparator)
		if fullPath == resultDir || !strings.HasPrefix(fullPath, prefix) {
			common.ApiErrorMsg(c, "图片结果不存在")
			return
		}
		info, statErr := os.Stat(fullPath)
		if statErr == nil && !info.IsDir() {
			c.Header("Cache-Control", "private, no-store, no-cache, max-age=0")
			c.Header("Pragma", "no-cache")
			c.Header("Content-Type", getImageStudioResultContentType(outputFormat))
			c.File(fullPath)
			return
		}
	}

	if url, _ := item["url"].(string); strings.TrimSpace(url) != "" {
		c.Redirect(http.StatusTemporaryRedirect, url)
		return
	}
	b64, _ := item["b64_json"].(string)
	if strings.TrimSpace(b64) == "" {
		common.ApiErrorMsg(c, "图片结果不存在")
		return
	}
	decoded, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		common.ApiErrorMsg(c, "图片结果解码失败")
		return
	}
	contentType := getImageStudioResultContentType(fmt.Sprintf("%v", item["output_format"]))
	c.Header("Cache-Control", "private, no-store, max-age=0")
	c.Data(http.StatusOK, contentType, decoded)
}

func GetImageStudioResultFile(c *gin.Context) {
	relativePath := strings.TrimPrefix(c.Param("filepath"), "/")
	relativePath = strings.TrimSpace(relativePath)
	if relativePath == "" {
		c.AbortWithStatus(http.StatusNotFound)
		return
	}
	expiresAt, err := strconv.ParseInt(strings.TrimSpace(c.Query("expires")), 10, 64)
	if err != nil || expiresAt <= 0 || time.Now().Unix() > expiresAt {
		c.AbortWithStatus(http.StatusForbidden)
		return
	}
	sig := strings.TrimSpace(c.Query("sig"))
	if sig == "" {
		c.AbortWithStatus(http.StatusForbidden)
		return
	}
	expected := signImageStudioResult(relativePath, expiresAt)
	if !hmac.Equal([]byte(strings.ToLower(sig)), []byte(expected)) {
		c.AbortWithStatus(http.StatusForbidden)
		return
	}
	baseDir := filepath.Clean(imageStudioResultBaseDir)
	fullPath := filepath.Clean(filepath.Join(baseDir, filepath.Clean(relativePath)))
	prefix := baseDir + string(os.PathSeparator)
	if fullPath == baseDir || !strings.HasPrefix(fullPath, prefix) {
		c.AbortWithStatus(http.StatusNotFound)
		return
	}
	info, err := os.Stat(fullPath)
	if err != nil || info.IsDir() {
		c.AbortWithStatus(http.StatusNotFound)
		return
	}
	c.Header("Cache-Control", "private, no-store, no-cache, max-age=0")
	c.Header("Pragma", "no-cache")
	c.File(fullPath)
}