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"
	"time"
	"unicode/utf8"

	"img_infinite_canvas/internal/domain/design"

	"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
	Count          int
	Images         []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 (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)

	idempotencyKey := strings.TrimSpace(opts.IdempotencyKey)
	streamPlans := []bool{false}
	if len(images) == 0 {
		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, count, images, 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)
			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
		}
	}
	return GeneratedImage{}, lastErr
}

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

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

func jsonImageEditBody(model string, prompt string, size string, count int, images []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 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, count int, images []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 count > 1 {
		fields["n"] = strconv.Itoa(count)
	}
	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) (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 := collectGeneratedImageURLs(responsePayload)
	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(collectGeneratedImageURLs(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 = collectGeneratedImageURLs(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 {
	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 := generatedImageString(key, typed); 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 {
	value = strings.TrimSpace(value)
	if value == "" {
		return ""
	}
	if strings.HasPrefix(value, "data:image/") {
		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:image/png;base64," + strings.TrimPrefix(value, "data:image/png;base64,")
	}
	return ""
}

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

func isGeneratedImageBase64Key(key string) bool {
	switch key {
	case "b64_json", "b64json", "base64", "image_base64", "imagebase64", "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 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")
}