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

package piagent

import (
	"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
	maxImageGenerationTimeout     = 10 * time.Minute
)

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
}

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
	}

	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,
				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)
	body, contentType, err := c.imageRequestBody(ctx, model, prompt, size, count, images)
	if err != nil {
		return GeneratedImage{}, err
	}

	var lastErr error
	startedAt := time.Now()
	promptRunes := utf8.RuneCountInString(prompt)
	for attempt := 0; attempt < 3; attempt++ {
		image, err := c.generateOnce(ctx, body, contentType, idempotencyKey, len(images) > 0)
		if err == nil {
			logx.Infof(
				"gpt image generation succeeded model=%s size=%s count=%d prompt_runes=%d attempt=%d elapsed=%s timeout=%s images=%d idempotency_key=%s",
				model,
				size,
				count,
				promptRunes,
				attempt+1,
				time.Since(startedAt).Round(time.Millisecond),
				c.client.Timeout,
				len(image.URLs),
				idempotencyKey,
			)
			return image, nil
		}
		lastErr = err
		retryable := isRetryableImageGenerationError(err)
		terminal := isImageGenerationTerminalError(err)
		logx.Errorf(
			"gpt image generation failed model=%s size=%s count=%d prompt_runes=%d attempt=%d elapsed=%s timeout=%s retryable=%t terminal=%t idempotency_key=%s err=%v",
			model,
			size,
			count,
			promptRunes,
			attempt+1,
			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):
		}
	}
	return GeneratedImage{}, lastErr
}

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

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

func jsonImageEditBody(model string, prompt string, size string, count int, images []string) ([]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
	}
	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) ([]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)
	}
	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) (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 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 {
		var detail struct {
			Error any `json:"error"`
		}
		_ = json.NewDecoder(resp.Body).Decode(&detail)
		if detail.Error != nil {
			return GeneratedImage{}, fmt.Errorf("image generation failed: %s: %v", resp.Status, detail.Error)
		}
		return GeneratedImage{}, fmt.Errorf("image generation failed: %s", resp.Status)
	}

	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
}

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":
		return true
	default:
		return false
	}
}

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
	}
	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")
}