moteva/server/internal/infrastructure/piagent/image_generation.go

package piagent

import (
	"bufio"
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net"
	"net/http"
	"net/textproto"
	"net/url"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"img_infinite_canvas/internal/domain/design"
	"img_infinite_canvas/internal/infrastructure/netutil"

	"github.com/zeromicro/go-zero/core/logx"
)

const (
	defaultImageGenerationTimeout               = 10 * time.Minute
	defaultImageGenerationResponseHeaderTimeout = 4 * time.Minute
	maxImageGenerationTimeout                   = 10 * time.Minute
	imageGenerationStreamMaxLineBytes           = 64 << 20
)

type ImageGenerator interface {
	Generate(ctx context.Context, prompt string, opts ImageRequestOptions) (GeneratedImage, error)
}

type GeneratedImage struct {
	URL  string
	URLs []string
}

type ImageRequestOptions struct {
	Model          string
	Size           string
	Quality        string
	Count          int
	Images         []string
	OutputFormat   string
	IdempotencyKey string
}

type ImageClientOptions struct {
	BaseURL                      string
	APIKey                       string
	InputImageTransport          string
	TimeoutSeconds               int
	ResponseHeaderTimeoutSeconds int
}

type ImageClient struct {
	baseURL             string
	apiKey              string
	inputImageTransport string
	client              *http.Client
}

func NewImageClient(opts ImageClientOptions) *ImageClient {
	if strings.TrimSpace(opts.BaseURL) == "" || strings.TrimSpace(opts.APIKey) == "" {
		return nil
	}
	timeout := time.Duration(opts.TimeoutSeconds) * time.Second
	if timeout <= 0 {
		timeout = defaultImageGenerationTimeout
	}
	if timeout > maxImageGenerationTimeout {
		timeout = maxImageGenerationTimeout
	}
	responseHeaderTimeout := time.Duration(opts.ResponseHeaderTimeoutSeconds) * time.Second
	if responseHeaderTimeout <= 0 {
		responseHeaderTimeout = defaultImageGenerationResponseHeaderTimeout
	}
	if responseHeaderTimeout > timeout {
		responseHeaderTimeout = timeout
	}

	return &ImageClient{
		baseURL:             strings.TrimRight(opts.BaseURL, "/"),
		apiKey:              opts.APIKey,
		inputImageTransport: normalizeInputImageTransport(opts.InputImageTransport),
		client: &http.Client{
			Timeout: timeout,
			Transport: &http.Transport{
				Proxy:                 http.ProxyFromEnvironment,
				DialContext:           (&net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
				DisableKeepAlives:     true,
				ForceAttemptHTTP2:     false,
				ResponseHeaderTimeout: responseHeaderTimeout,
				TLSHandshakeTimeout:   10 * time.Second,
				ExpectContinueTimeout: time.Second,
			},
		},
	}
}

func normalizeInputImageTransport(value string) string {
	switch strings.ToLower(strings.TrimSpace(value)) {
	case "url", "image_url", "image-url":
		return "url"
	default:
		return "file"
	}
}

func normalizeImageOutputFormat(value string) string {
	switch strings.ToLower(strings.TrimSpace(value)) {
	case "psd", "png", "jpeg", "jpg", "webp":
		if strings.EqualFold(strings.TrimSpace(value), "jpg") {
			return "jpeg"
		}
		return strings.ToLower(strings.TrimSpace(value))
	default:
		return ""
	}
}

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

func (c *ImageClient) GenerateLayeredDocument(ctx context.Context, req design.LayeredDocumentRequest) (design.LayeredDocument, error) {
	imageURL := strings.TrimSpace(req.ImageURL)
	if imageURL == "" {
		return design.LayeredDocument{}, fmt.Errorf("layered document image_url is required")
	}
	outputFormat := normalizeImageOutputFormat(req.OutputFormat)
	if outputFormat == "" {
		outputFormat = "psd"
	}
	prompt := strings.TrimSpace(req.Prompt)
	if prompt == "" {
		prompt = "Create a layered PSD from the provided image. Preserve the exact canvas and visual design. Put a single clean background layer at the bottom with all text and small foreground elements removed. Put each small standalone foreground element on its own image layer. Keep text as editable text layers when possible. Do not merge text into the background. Return only the layered PSD file."
	}
	size := strings.TrimSpace(req.Size)
	if size == "" {
		size = "auto"
	}
	generated, err := c.Generate(ctx, prompt, ImageRequestOptions{
		Model:          req.Model,
		Size:           size,
		Count:          1,
		Images:         []string{imageURL},
		OutputFormat:   outputFormat,
		IdempotencyKey: strings.TrimSpace(req.IdempotencyKey),
	})
	if err != nil {
		return design.LayeredDocument{}, err
	}
	content := strings.TrimSpace(generated.URL)
	if content == "" && len(generated.URLs) > 0 {
		content = strings.TrimSpace(generated.URLs[0])
	}
	if content == "" {
		return design.LayeredDocument{}, fmt.Errorf("layered document generation returned empty file")
	}
	return design.LayeredDocument{
		Content:     content,
		ContentType: outputFormatContentType(outputFormat),
	}, nil
}

func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageRequestOptions) (GeneratedImage, error) {
	prompt = strings.TrimSpace(prompt)
	if prompt == "" {
		return GeneratedImage{}, fmt.Errorf("image prompt is required")
	}
	model := strings.TrimSpace(opts.Model)
	if model == "" {
		model = "gpt-image-2"
	}
	size := strings.TrimSpace(opts.Size)
	if size == "" {
		size = design.DefaultGeneratedImageSize
	}
	count := opts.Count
	if count < 1 {
		count = 1
	}
	if count > 10 {
		count = 10
	}
	images := normalizeImageInputs(opts.Images)
	quality := normalizeImageQuality(opts.Quality)
	outputFormat := normalizeImageOutputFormat(opts.OutputFormat)

	idempotencyKey := strings.TrimSpace(opts.IdempotencyKey)
	if count > 1 {
		logx.Infof(
			"gpt image generation splitting multi-image request model=%s size=%s count=%d prompt_runes=%d images=%d output_format=%s idempotency_key=%s",
			model,
			size,
			count,
			utf8.RuneCountInString(prompt),
			len(images),
			outputFormat,
			idempotencyKey,
		)
		return c.generateImageCountFallback(ctx, prompt, model, size, quality, count, images, outputFormat, idempotencyKey)
	}

	streamPlans := []bool{false}
	if len(images) == 0 && count == 1 && outputFormat == "" {
		streamPlans = []bool{true, false}
	}

	var lastErr error
	startedAt := time.Now()
	promptRunes := utf8.RuneCountInString(prompt)
	for _, stream := range streamPlans {
		body, contentType, err := c.imageRequestBody(ctx, model, prompt, size, quality, count, images, outputFormat, stream)
		if err != nil {
			return GeneratedImage{}, err
		}
		for attempt := 0; attempt < 3; attempt++ {
			image, err := c.generateOnce(ctx, body, contentType, idempotencyKey, len(images) > 0, stream, count, outputFormat)
			if err == nil {
				logx.Infof(
					"gpt image generation succeeded model=%s size=%s count=%d prompt_runes=%d attempt=%d stream=%t elapsed=%s timeout=%s images=%d idempotency_key=%s",
					model,
					size,
					count,
					promptRunes,
					attempt+1,
					stream,
					time.Since(startedAt).Round(time.Millisecond),
					c.client.Timeout,
					len(image.URLs),
					idempotencyKey,
				)
				return image, nil
			}
			if ctx.Err() != nil {
				return GeneratedImage{}, ctx.Err()
			}
			lastErr = err
			if stream && isImageGenerationStreamFallbackError(err) {
				logx.Errorf(
					"gpt image generation stream fallback model=%s size=%s count=%d prompt_runes=%d attempt=%d elapsed=%s idempotency_key=%s err=%v",
					model,
					size,
					count,
					promptRunes,
					attempt+1,
					time.Since(startedAt).Round(time.Millisecond),
					idempotencyKey,
					err,
				)
				break
			}
			timeoutRetryable := idempotencyKey != "" && isImageGenerationTimeoutError(err)
			retryable := isRetryableImageGenerationError(err) || timeoutRetryable
			terminal := isImageGenerationTerminalError(err) && !timeoutRetryable
			logx.Errorf(
				"gpt image generation failed model=%s size=%s count=%d prompt_runes=%d attempt=%d stream=%t elapsed=%s timeout=%s retryable=%t terminal=%t idempotency_key=%s err=%v",
				model,
				size,
				count,
				promptRunes,
				attempt+1,
				stream,
				time.Since(startedAt).Round(time.Millisecond),
				c.client.Timeout,
				retryable,
				terminal,
				idempotencyKey,
				err,
			)
			if !retryable || attempt == 2 {
				break
			}
			select {
			case <-ctx.Done():
				return GeneratedImage{}, ctx.Err()
			case <-time.After(time.Duration(attempt+1) * 1500 * time.Millisecond):
			}
		}
		if !(stream && shouldFallbackImageGenerationStream(lastErr, idempotencyKey)) {
			break
		}
	}
	if count > 1 && isImageGenerationCountFallbackError(lastErr) {
		logx.Errorf(
			"gpt image generation count fallback model=%s size=%s count=%d prompt_runes=%d elapsed=%s idempotency_key=%s err=%v",
			model,
			size,
			count,
			promptRunes,
			time.Since(startedAt).Round(time.Millisecond),
			idempotencyKey,
			lastErr,
		)
		return c.generateImageCountFallback(ctx, prompt, model, size, quality, count, images, outputFormat, idempotencyKey)
	}
	return GeneratedImage{}, lastErr
}

func (c *ImageClient) generateImageCountFallback(ctx context.Context, prompt string, model string, size string, quality string, count int, images []string, outputFormat string, idempotencyKey string) (GeneratedImage, error) {
	fallbackCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	urls := make([]string, count)
	var wg sync.WaitGroup
	var mu sync.Mutex
	var firstErr error
	setErr := func(err error) {
		if err == nil {
			return
		}
		mu.Lock()
		defer mu.Unlock()
		if firstErr == nil {
			firstErr = err
			cancel()
		}
	}

	for index := 0; index < count; index++ {
		index := index
		wg.Add(1)
		go func() {
			defer wg.Done()
			fallbackOpts := ImageRequestOptions{
				Model:        model,
				Size:         size,
				Quality:      quality,
				Count:        1,
				Images:       append([]string(nil), images...),
				OutputFormat: outputFormat,
			}
			if idempotencyKey != "" {
				fallbackOpts.IdempotencyKey = fmt.Sprintf("%s-%d", idempotencyKey, index+1)
			}
			image, err := c.Generate(fallbackCtx, prompt, fallbackOpts)
			if err != nil {
				setErr(err)
				return
			}
			url := image.URL
			if url == "" && len(image.URLs) > 0 {
				url = image.URLs[0]
			}
			if url == "" {
				setErr(fmt.Errorf("image generation count fallback returned empty image"))
				return
			}
			urls[index] = url
		}()
	}
	wg.Wait()

	mu.Lock()
	err := firstErr
	mu.Unlock()
	if err != nil {
		return GeneratedImage{}, err
	}
	return GeneratedImage{URL: urls[0], URLs: urls}, nil
}

func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt string, size string, quality string, count int, images []string, outputFormat string, stream bool) ([]byte, string, error) {
	if len(images) > 0 {
		if c.inputImageTransport == "url" && imageInputsCanUseURLTransport(images) {
			return jsonImageEditBody(model, prompt, size, quality, count, images, outputFormat, stream)
		}
		return c.multipartImageEditBody(ctx, model, prompt, size, quality, count, images, outputFormat, stream)
	}

	requestPayload := map[string]any{
		"model":  model,
		"prompt": prompt,
	}
	if count > 1 {
		requestPayload["n"] = count
	}
	if size != "" {
		requestPayload["size"] = size
	}
	if quality != "" {
		requestPayload["quality"] = quality
	}
	if outputFormat != "" {
		requestPayload["output_format"] = outputFormat
	}
	if stream {
		requestPayload["stream"] = true
	}
	body, err := json.Marshal(requestPayload)
	return body, "application/json", err
}

func imageInputsCanUseURLTransport(images []string) bool {
	for _, image := range images {
		if imageInputRequiresFileUpload(image) {
			return false
		}
	}
	return true
}

func imageInputRequiresFileUpload(value string) bool {
	value = strings.TrimSpace(value)
	if value == "" {
		return false
	}
	if strings.HasPrefix(value, "data:image/") || strings.HasPrefix(value, "/") {
		return true
	}
	parsed, err := url.Parse(value)
	if err != nil {
		return true
	}
	scheme := strings.ToLower(parsed.Scheme)
	if scheme != "http" && scheme != "https" {
		return true
	}
	host := strings.ToLower(parsed.Hostname())
	if host == "" {
		return true
	}
	return netutil.IsLocalOrPrivateHost(host)
}

func jsonImageEditBody(model string, prompt string, size string, quality string, count int, images []string, outputFormat string, stream bool) ([]byte, string, error) {
	requestPayload := map[string]any{
		"model":  model,
		"prompt": prompt,
		"images": imageReferencePayload(images),
	}
	if count > 1 {
		requestPayload["n"] = count
	}
	if size != "" {
		requestPayload["size"] = size
	}
	if quality != "" {
		requestPayload["quality"] = quality
	}
	if outputFormat != "" {
		requestPayload["output_format"] = outputFormat
	}
	if stream {
		requestPayload["stream"] = true
	}
	body, err := json.Marshal(requestPayload)
	return body, "application/json", err
}

func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string, prompt string, size string, quality string, count int, images []string, outputFormat string, stream bool) ([]byte, string, error) {
	var body bytes.Buffer
	writer := multipart.NewWriter(&body)
	fields := map[string]string{
		"model":  model,
		"prompt": prompt,
	}
	if size != "" {
		fields["size"] = size
	}
	if quality != "" {
		fields["quality"] = quality
	}
	if count > 1 {
		fields["n"] = strconv.Itoa(count)
	}
	if outputFormat != "" {
		fields["output_format"] = outputFormat
	}
	if stream {
		fields["stream"] = "true"
	}
	for key, value := range fields {
		if err := writer.WriteField(key, value); err != nil {
			return nil, "", err
		}
	}

	fieldName := "image"
	if len(images) > 1 {
		fieldName = "image[]"
	}
	for index, image := range images {
		input, err := c.imageInputFile(ctx, image, index)
		if err != nil {
			_ = writer.Close()
			return nil, "", err
		}
		part, err := createImageFormFile(writer, fieldName, input)
		if err != nil {
			_ = writer.Close()
			return nil, "", err
		}
		if _, err := part.Write(input.Data); err != nil {
			_ = writer.Close()
			return nil, "", err
		}
	}
	if err := writer.Close(); err != nil {
		return nil, "", err
	}
	return body.Bytes(), writer.FormDataContentType(), nil
}

func createImageFormFile(writer *multipart.Writer, fieldName string, input imageInputFile) (io.Writer, error) {
	header := make(textproto.MIMEHeader)
	header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, escapeMultipartValue(fieldName), escapeMultipartValue(input.FileName)))
	contentType := strings.TrimSpace(input.ContentType)
	if contentType == "" {
		contentType = "application/octet-stream"
	}
	header.Set("Content-Type", contentType)
	return writer.CreatePart(header)
}

func escapeMultipartValue(value string) string {
	return strings.NewReplacer("\\", "\\\\", `"`, "\\\"").Replace(value)
}

type imageInputFile struct {
	FileName    string
	ContentType string
	Data        []byte
}

func (c *ImageClient) imageInputFile(ctx context.Context, value string, index int) (imageInputFile, error) {
	value = strings.TrimSpace(value)
	if strings.HasPrefix(value, "data:image/") {
		header, payload, ok := strings.Cut(value, ",")
		if !ok {
			return imageInputFile{}, fmt.Errorf("invalid image data url")
		}
		contentType := strings.TrimPrefix(strings.Split(header, ";")[0], "data:")
		data, err := base64.StdEncoding.DecodeString(payload)
		if err != nil {
			return imageInputFile{}, err
		}
		return imageInputFile{
			FileName:    fmt.Sprintf("reference-%d%s", index+1, imageExtension(contentType, "")),
			ContentType: contentType,
			Data:        data,
		}, nil
	}
	if !strings.HasPrefix(value, "http://") && !strings.HasPrefix(value, "https://") {
		return imageInputFile{}, fmt.Errorf("unsupported image input %q", value)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, value, nil)
	if err != nil {
		return imageInputFile{}, err
	}
	req.Header.Set("User-Agent", "img-infinite-canvas/1.0")
	resp, err := c.client.Do(req)
	if err != nil {
		return imageInputFile{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return imageInputFile{}, fmt.Errorf("download reference image failed: %s", resp.Status)
	}
	contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
	data, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20))
	if err != nil {
		return imageInputFile{}, err
	}
	if len(data) == 0 {
		return imageInputFile{}, fmt.Errorf("download reference image returned empty body")
	}
	if contentType == "" || !strings.HasPrefix(contentType, "image/") {
		contentType = http.DetectContentType(data)
	}
	if !strings.HasPrefix(contentType, "image/") {
		return imageInputFile{}, fmt.Errorf("reference input is not an image: %s", contentType)
	}
	return imageInputFile{
		FileName:    imageFileName(value, contentType, index),
		ContentType: contentType,
		Data:        data,
	}, nil
}

func normalizeImageInputs(images []string) []string {
	if len(images) == 0 {
		return nil
	}
	normalized := make([]string, 0, len(images))
	seen := make(map[string]struct{})
	for _, image := range images {
		image = strings.TrimSpace(image)
		if image == "" {
			continue
		}
		if _, ok := seen[image]; ok {
			continue
		}
		seen[image] = struct{}{}
		normalized = append(normalized, image)
		if len(normalized) >= 16 {
			break
		}
	}
	return normalized
}

func imageReferencePayload(images []string) []map[string]string {
	items := make([]map[string]string, 0, len(images))
	for _, image := range images {
		items = append(items, map[string]string{"image_url": image})
	}
	return items
}

func imageFileName(value string, contentType string, index int) string {
	name := ""
	if parsed, err := url.Parse(value); err == nil {
		name = filepath.Base(parsed.Path)
	}
	if name == "" || name == "." || name == "/" {
		name = fmt.Sprintf("reference-%d%s", index+1, imageExtension(contentType, ""))
	}
	if filepath.Ext(name) == "" {
		name += imageExtension(contentType, ".png")
	}
	return name
}

func imageExtension(contentType string, fallback string) string {
	switch strings.ToLower(strings.TrimSpace(contentType)) {
	case "image/png":
		return ".png"
	case "image/jpeg", "image/jpg":
		return ".jpg"
	case "image/webp":
		return ".webp"
	case "image/gif":
		return ".gif"
	default:
		if fallback != "" {
			return fallback
		}
		return ".png"
	}
}

func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType string, idempotencyKey string, edit bool, stream bool, expectedCount int, outputFormat string) (GeneratedImage, error) {
	endpoint := "/v1/images/generations"
	if edit {
		endpoint = "/v1/images/edits"
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+endpoint, bytes.NewReader(body))
	if err != nil {
		return GeneratedImage{}, err
	}
	req.Header.Set("Authorization", "Bearer "+c.apiKey)
	req.Header.Set("Content-Type", contentType)
	req.Header.Set("Connection", "close")
	if stream {
		req.Header.Set("Accept", "text/event-stream")
	}
	if idempotencyKey != "" {
		req.Header.Set("Idempotency-Key", idempotencyKey)
	}
	req.Close = true

	resp, err := c.client.Do(req)
	if err != nil {
		return GeneratedImage{}, err
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		respBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
		if readErr != nil {
			return GeneratedImage{}, readErr
		}
		detail := imageErrorBody(respBody)
		if detail != "" {
			return GeneratedImage{}, fmt.Errorf("image generation failed: %s: %s", resp.Status, detail)
		}
		return GeneratedImage{}, fmt.Errorf("image generation failed: %s", resp.Status)
	}

	if stream {
		return decodeImageGenerationStreamOrJSON(resp.Body, expectedCount)
	}

	var responsePayload any
	if err := json.NewDecoder(resp.Body).Decode(&responsePayload); err != nil {
		return GeneratedImage{}, err
	}
	urls := collectGeneratedAssetURLs(responsePayload, outputFormat)
	if len(urls) == 0 {
		return GeneratedImage{}, fmt.Errorf("image generation returned empty image")
	}

	return GeneratedImage{URL: urls[0], URLs: urls}, nil
}

type imageStreamEvent struct {
	name      string
	dataLines []string
	hasData   bool
}

func decodeImageGenerationStreamOrJSON(body io.Reader, expectedCount int) (GeneratedImage, error) {
	if expectedCount <= 0 {
		expectedCount = 1
	}
	scanner := bufio.NewScanner(body)
	scanner.Buffer(make([]byte, 0, 64*1024), imageGenerationStreamMaxLineBytes)
	var raw strings.Builder
	event := imageStreamEvent{name: "message"}
	sawEvents := false
	urls := make([]string, 0, expectedCount)
	seen := make(map[string]struct{})
	appendURLs := func(values []string) {
		for _, value := range values {
			if value == "" {
				continue
			}
			if _, ok := seen[value]; ok {
				continue
			}
			seen[value] = struct{}{}
			urls = append(urls, value)
		}
	}
	flush := func() (bool, error) {
		if !event.hasData {
			event = imageStreamEvent{name: "message"}
			return false, nil
		}
		sawEvents = true
		name := event.name
		dataText := strings.TrimSpace(strings.Join(event.dataLines, "\n"))
		event = imageStreamEvent{name: "message"}
		if dataText == "" {
			return false, nil
		}
		if dataText == "[DONE]" {
			return true, nil
		}
		if strings.Contains(strings.ToLower(name), "partial_image") {
			return false, nil
		}
		var payload any
		if err := json.Unmarshal([]byte(dataText), &payload); err != nil {
			return false, nil
		}
		if isPartialImagePayload(payload) {
			return false, nil
		}
		if message := imageStreamFailureMessage(payload); message != "" {
			return true, fmt.Errorf("image generation stream failed: %s", message)
		}
		appendURLs(collectGeneratedAssetURLs(payload, ""))
		return len(urls) >= expectedCount, nil
	}

	for scanner.Scan() {
		line := strings.TrimSuffix(scanner.Text(), "\r")
		trimmed := strings.TrimSpace(line)
		switch {
		case line == "":
			done, err := flush()
			if err != nil {
				return GeneratedImage{}, err
			}
			if done && len(urls) > 0 {
				return GeneratedImage{URL: urls[0], URLs: urls}, nil
			}
		case strings.HasPrefix(line, ":"):
			continue
		case strings.HasPrefix(line, "event:"):
			event.name = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
		case strings.HasPrefix(line, "data:"):
			event.hasData = true
			event.dataLines = append(event.dataLines, strings.TrimSpace(strings.TrimPrefix(line, "data:")))
		default:
			if !sawEvents && trimmed != "" {
				raw.WriteString(line)
				raw.WriteByte('\n')
			}
		}
	}
	if err := scanner.Err(); err != nil {
		return GeneratedImage{}, err
	}
	done, err := flush()
	if err != nil {
		return GeneratedImage{}, err
	}
	if len(urls) > 0 {
		return GeneratedImage{URL: urls[0], URLs: urls}, nil
	}
	if sawEvents || done {
		return GeneratedImage{}, fmt.Errorf("image generation stream returned no complete image")
	}

	var responsePayload any
	if err := json.Unmarshal([]byte(raw.String()), &responsePayload); err != nil {
		return GeneratedImage{}, err
	}
	urls = collectGeneratedAssetURLs(responsePayload, "")
	if len(urls) == 0 {
		return GeneratedImage{}, fmt.Errorf("image generation returned empty image")
	}
	return GeneratedImage{URL: urls[0], URLs: urls}, nil
}

func imageErrorBody(body []byte) string {
	body = bytes.TrimSpace(body)
	if len(body) == 0 {
		return ""
	}
	var detail struct {
		Error any `json:"error"`
	}
	if err := json.Unmarshal(body, &detail); err == nil && detail.Error != nil {
		return fmt.Sprint(detail.Error)
	}
	return trimErrorBody(body)
}

func isPartialImagePayload(payload any) bool {
	item, ok := payload.(map[string]any)
	if !ok {
		return false
	}
	if value, ok := item["type"].(string); ok && strings.Contains(strings.ToLower(value), "partial_image") {
		return true
	}
	return false
}

func imageStreamFailureMessage(payload any) string {
	item, ok := payload.(map[string]any)
	if !ok {
		return ""
	}
	if message := imageErrorMessage(item["error"]); message != "" {
		return message
	}
	typeValue, _ := item["type"].(string)
	statusValue, _ := item["status"].(string)
	if strings.Contains(strings.ToLower(typeValue), "failed") || strings.EqualFold(statusValue, "failed") {
		if message, ok := item["message"].(string); ok && strings.TrimSpace(message) != "" {
			return strings.TrimSpace(message)
		}
		return "upstream reported failed status"
	}
	return ""
}

func imageErrorMessage(value any) string {
	switch typed := value.(type) {
	case nil:
		return ""
	case string:
		return strings.TrimSpace(typed)
	case map[string]any:
		for _, key := range []string{"message", "detail", "type", "code"} {
			if message, ok := typed[key].(string); ok && strings.TrimSpace(message) != "" {
				return strings.TrimSpace(message)
			}
		}
		return fmt.Sprint(typed)
	default:
		return fmt.Sprint(typed)
	}
}

func collectGeneratedImageURLs(payload any) []string {
	return collectGeneratedAssetURLs(payload, "")
}

func collectGeneratedAssetURLs(payload any, outputFormat string) []string {
	urls := make([]string, 0)
	seen := make(map[string]struct{})
	var walk func(any, string)
	walk = func(value any, key string) {
		switch typed := value.(type) {
		case map[string]any:
			for childKey, childValue := range typed {
				walk(childValue, childKey)
			}
		case []any:
			for _, childValue := range typed {
				walk(childValue, key)
			}
		case string:
			if imageURL := generatedAssetString(key, typed, outputFormat); imageURL != "" {
				if _, ok := seen[imageURL]; ok {
					return
				}
				seen[imageURL] = struct{}{}
				urls = append(urls, imageURL)
			}
		}
	}
	walk(payload, "")
	return urls
}

func generatedImageString(key string, value string) string {
	return generatedAssetString(key, value, "")
}

func generatedAssetString(key string, value string, outputFormat string) string {
	value = strings.TrimSpace(value)
	if value == "" {
		return ""
	}
	if strings.HasPrefix(value, "data:") {
		return value
	}
	normalizedKey := strings.ToLower(strings.TrimSpace(key))
	if isGeneratedImageURLKey(normalizedKey) && (strings.HasPrefix(value, "http://") || strings.HasPrefix(value, "https://")) {
		return value
	}
	if isGeneratedImageBase64Key(normalizedKey) {
		return "data:" + outputFormatContentType(outputFormat) + ";base64," + stripDataURLPrefix(value)
	}
	return ""
}

func stripDataURLPrefix(value string) string {
	if _, payload, ok := strings.Cut(value, ","); ok && strings.HasPrefix(strings.TrimSpace(value), "data:") {
		return payload
	}
	return value
}

func outputFormatContentType(format string) string {
	switch normalizeImageOutputFormat(format) {
	case "psd":
		return "image/vnd.adobe.photoshop"
	case "jpeg":
		return "image/jpeg"
	case "webp":
		return "image/webp"
	default:
		return "image/png"
	}
}

func isGeneratedImageURLKey(key string) bool {
	switch key {
	case "url", "image_url", "imageurl", "image", "output_url", "public_url", "publicurl", "file_url", "fileurl", "download_url", "downloadurl":
		return true
	default:
		return false
	}
}

func isGeneratedImageBase64Key(key string) bool {
	switch key {
	case "b64_json", "b64json", "base64", "image_base64", "imagebase64", "file_base64", "filebase64", "file_b64", "fileb64", "result":
		return true
	default:
		return false
	}
}

func isImageGenerationStreamFallbackError(err error) bool {
	if err == nil {
		return false
	}
	message := strings.ToLower(err.Error())
	return strings.Contains(message, "unknown parameter") && strings.Contains(message, "stream") ||
		strings.Contains(message, "unsupported") && strings.Contains(message, "stream") ||
		strings.Contains(message, "not support") && strings.Contains(message, "stream") ||
		strings.Contains(message, "stream returned no complete image") ||
		strings.Contains(message, "incomplete stream") ||
		strings.Contains(message, "text/event-stream")
}

func isImageGenerationCountFallbackError(err error) bool {
	if err == nil {
		return false
	}
	message := strings.ToLower(err.Error())
	return strings.Contains(message, "n > 1") && strings.Contains(message, "stream=false") ||
		strings.Contains(message, "n > 1") && strings.Contains(message, "non-streaming") ||
		strings.Contains(message, "unknown parameter") && (strings.Contains(message, "tools[0].n") ||
			strings.Contains(message, "param:n") ||
			strings.Contains(message, "param: n") ||
			strings.Contains(message, "parameter: 'n'") ||
			strings.Contains(message, `parameter: "n"`))
}

func shouldFallbackImageGenerationStream(err error, idempotencyKey string) bool {
	if isImageGenerationStreamFallbackError(err) {
		return true
	}
	if strings.TrimSpace(idempotencyKey) == "" {
		return false
	}
	return isRetryableImageGenerationError(err) || isImageGenerationTimeoutError(err)
}

func isRetryableImageGenerationError(err error) bool {
	if err == nil {
		return false
	}
	if isImageGenerationTerminalError(err) {
		return false
	}
	message := strings.ToLower(err.Error())
	return strings.Contains(message, "eof") ||
		strings.Contains(message, "connection reset") ||
		strings.Contains(message, "unexpected end") ||
		strings.Contains(message, "temporarily unavailable") ||
		strings.Contains(message, "too many requests") ||
		strings.Contains(message, "rate limit") ||
		strings.Contains(message, "429") ||
		strings.Contains(message, "500") ||
		strings.Contains(message, "502") ||
		strings.Contains(message, "503") ||
		strings.Contains(message, "504")
}

func isImageGenerationTerminalError(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return true
	}
	return isImageGenerationTimeoutError(err)
}

func isImageGenerationTimeoutError(err error) bool {
	if err == nil {
		return false
	}
	var netErr net.Error
	if errors.As(err, &netErr) && netErr.Timeout() {
		return true
	}
	message := strings.ToLower(err.Error())
	return strings.Contains(message, "client.timeout exceeded") ||
		strings.Contains(message, "context deadline exceeded") ||
		strings.Contains(message, "deadline exceeded") ||
		strings.Contains(message, "request canceled") ||
		strings.Contains(message, "request cancelled") ||
		strings.Contains(message, "timeout awaiting response headers") ||
		strings.Contains(message, "timeout while awaiting headers")
}