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"
	"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"
	imageStudioTaskTTL            = 30 * time.Minute
	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"`
}

var (
	imageStudioTasks = struct {
		sync.RWMutex
		items map[string]*imageStudioTask
	}{items: make(map[string]*imageStudioTask)}
	imageStudioCleanupOnce sync.Once
)

func ensureImageStudioTaskCleanup() {
	imageStudioCleanupOnce.Do(func() {
		go func() {
			ticker := time.NewTicker(imageStudioTaskCleanupEvery)
			defer ticker.Stop()
			for range ticker.C {
				cutoff := time.Now().Add(-imageStudioTaskTTL).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("请选择图片模型")
	}
	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(imageStudioTaskTTL).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 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
		}
		if b64, _ := item["b64_json"].(string); strings.TrimSpace(b64) != "" {
			delete(item, "b64_json")
			item["url"] = buildImageStudioTaskResultURL(task.ID, index)
			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 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)
}