gen2d/backend/internal/service/gptimage.go

package service

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"

	"gen2d/internal/config"
)

// gptImage2Cfg 保存 GPT Image 2 配置。
var gptImage2Cfg config.GptImage2Config

// InitGptImage2Config 注入 GPT Image 2 配置。
func InitGptImage2Config(cfg config.GptImage2Config) {
	gptImage2Cfg = cfg
}

// ======================== GPT Image 2 API 类型 ========================

// gpt2SubmitRequest 提交生图任务请求体。
type gpt2SubmitRequest struct {
	Prompt          string   `json:"prompt"`
	AspectRatio     string   `json:"aspect_ratio,omitempty"`
	ReferenceImages []string `json:"reference_images,omitempty"`
	XChannel        string   `json:"x_channel,omitempty"`
}

// gpt2SubmitResponse 提交生图任务响应体。
type gpt2SubmitResponse struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Data    struct {
		TaskID string `json:"task_id"`
		Status string `json:"status"`
	} `json:"data"`
}

// gpt2StatusRequest 查询任务状态请求体。
type gpt2StatusRequest struct {
	TaskID string `json:"task_id"`
}

// gpt2StatusResponse 查询任务状态响应体。
type gpt2StatusResponse struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Data    struct {
		TaskID         string `json:"task_id"`
		Status         string `json:"status"`
		Progress       int    `json:"progress"`
		ResultImageURL string `json:"result_image_url"`
		ErrorMessage   string `json:"error_message"`
	} `json:"data"`
}

// ======================== 公开接口 ========================

// GenerateImagesGpt2 通过 GPT Image 2 异步 API 生成图片。
// count 张图通过并行提交+轮询实现。
func GenerateImagesGpt2(ctx context.Context, prompt string, count int) ([]GeneratedImage, error) {
	if gptImage2Cfg.APIKey == "" {
		return nil, fmt.Errorf("gpt_image2 api_key not configured")
	}

	// 并行提交任务
	type submitResult struct {
		index  int
		taskID string
		err    error
	}
	results := make(chan submitResult, count)
	for i := 0; i < count; i++ {
		go func(idx int) {
			taskID, err := submitGpt2Task(ctx, prompt)
			results <- submitResult{index: idx, taskID: taskID, err: err}
		}(i)
	}

	// 收集 taskID
	taskIDs := make([]string, count)
	for i := 0; i < count; i++ {
		r := <-results
		if r.err != nil {
			return nil, fmt.Errorf("submit task %d: %w", r.index, r.err)
		}
		taskIDs[r.index] = r.taskID
	}

	// 并行轮询+下载
	type imageResult struct {
		index int
		image GeneratedImage
		err   error
	}
	imgResults := make(chan imageResult, count)
	for i, tid := range taskIDs {
		go func(idx int, taskID string) {
			img, err := pollAndDownloadGpt2(ctx, taskID, prompt)
			imgResults <- imageResult{index: idx, image: img, err: err}
		}(i, tid)
	}

	images := make([]GeneratedImage, count)
	for i := 0; i < count; i++ {
		r := <-imgResults
		if r.err != nil {
			return nil, fmt.Errorf("task %d: %w", r.index, r.err)
		}
		images[r.index] = r.image
	}

	return images, nil
}

// ======================== 内部实现 ========================

// submitGpt2Task 提交生图任务，返回 taskID。
func submitGpt2Task(ctx context.Context, prompt string) (string, error) {
	reqBody := gpt2SubmitRequest{
		Prompt:      prompt,
		AspectRatio: gptImage2Cfg.AspectRatio,
		XChannel:    gptImage2Cfg.XChannel,
	}

	body, err := json.Marshal(reqBody)
	if err != nil {
		return "", fmt.Errorf("marshal: %w", err)
	}

	url := strings.TrimRight(gptImage2Cfg.BaseURL, "/") + "/gpt-image2/generate"
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return "", fmt.Errorf("create request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+gptImage2Cfg.APIKey)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("send: %w", err)
	}
	defer resp.Body.Close()

	var submitResp gpt2SubmitResponse
	if err := json.NewDecoder(resp.Body).Decode(&submitResp); err != nil {
		return "", fmt.Errorf("decode: %w", err)
	}

	if submitResp.Code != 200 {
		return "", fmt.Errorf("submit failed: %s", submitResp.Message)
	}

	log.Printf("[gpt_image2] task submitted: %s", submitResp.Data.TaskID)
	return submitResp.Data.TaskID, nil
}

// pollAndDownloadGpt2 轮询任务状态直到完成，然后下载图片。
func pollAndDownloadGpt2(ctx context.Context, taskID, prompt string) (GeneratedImage, error) {
	pollInterval := gptImage2Cfg.PollInterval
	if pollInterval <= 0 {
		pollInterval = 3
	}
	maxWait := gptImage2Cfg.PollMaxWait
	if maxWait <= 0 {
		maxWait = 120
	}

	deadline := time.Now().Add(time.Duration(maxWait) * time.Second)
	ticker := time.NewTicker(time.Duration(pollInterval) * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return GeneratedImage{}, ctx.Err()
		case <-ticker.C:
			status, err := queryGpt2Status(ctx, taskID)
			if err != nil {
				return GeneratedImage{}, fmt.Errorf("query status: %w", err)
			}

			switch status.Data.Status {
			case "completed":
				log.Printf("[gpt_image2] task %s completed, downloading from %s", taskID, status.Data.ResultImageURL)
				data, err := downloadGpt2Image(ctx, status.Data.ResultImageURL)
				if err != nil {
					return GeneratedImage{}, fmt.Errorf("download: %w", err)
				}
				return GeneratedImage{Data: data, Format: "png"}, nil

			case "failed":
				errMsg := status.Data.ErrorMessage
				if errMsg == "" {
					errMsg = "unknown error"
				}
				return GeneratedImage{}, fmt.Errorf("generation failed: %s", errMsg)

			default:
				log.Printf("[gpt_image2] task %s status=%s progress=%d", taskID, status.Data.Status, status.Data.Progress)
			}

			if time.Now().After(deadline) {
				return GeneratedImage{}, fmt.Errorf("poll timeout after %ds", maxWait)
			}
		}
	}
}

// queryGpt2Status 查询任务状态。
func queryGpt2Status(ctx context.Context, taskID string) (*gpt2StatusResponse, error) {
	reqBody := gpt2StatusRequest{TaskID: taskID}
	body, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("marshal: %w", err)
	}

	url := strings.TrimRight(gptImage2Cfg.BaseURL, "/") + "/gpt-image2/status"
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+gptImage2Cfg.APIKey)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("send: %w", err)
	}
	defer resp.Body.Close()

	var statusResp gpt2StatusResponse
	if err := json.NewDecoder(resp.Body).Decode(&statusResp); err != nil {
		return nil, fmt.Errorf("decode: %w", err)
	}

	if statusResp.Code != 200 {
		return nil, fmt.Errorf("status query failed: %s", statusResp.Message)
	}

	return &statusResp, nil
}

// downloadGpt2Image 下载生成的图片。
func downloadGpt2Image(ctx context.Context, imageURL string) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil)
	if err != nil {
		return nil, fmt.Errorf("create download request: %w", err)
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("download: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("download status %d", resp.StatusCode)
	}

	return io.ReadAll(resp.Body)
}