image2api/backend/internal/provider/imagine/image.go

package imagine

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"strconv"
	"strings"
	"time"

	http "github.com/bogdanfinn/fhttp"
)

// GenerateImage runs the full Imagine.art pipeline: submit the txt2img job, poll
// the org objects feed until the batch finishes, then download the produced
// image. styleID picks the model (41001 = 1.5 / 2K, 41004 = 1.5pro / 4K).
// HTTP 402 → ErrQuotaExhausted. These models are pure text2img (no refs).
func (c *Client) GenerateImage(ctx context.Context, cred string, styleID int, resolution, aspectRatio, prompt string) ([]byte, map[string]any, error) {
	cr, ok := parseCred(cred)
	if !ok {
		return nil, nil, ErrAuth
	}
	userID := userIDFromToken(cr.Token)

	metadata, _ := json.Marshal(map[string]any{
		"placeholderUuid":           uuid4(),
		"promptWithoutManipulation": prompt,
		"modeId":                    0,
	})
	// parent_id MUST be a canvas node this account owns — the server rejects a
	// foreign id ("user does not have access to parent asset") and silently
	// orphans a parent-less generation (charged but never produced). It's supplied
	// at import (credential.parentId).
	fields := map[string]string{
		"style_id":      strconv.Itoa(styleID),
		"aspect_ratio":  aspectRatio,
		"resolution":    resolution,
		"variation":     "txt2img",
		"prompt":        prompt,
		"is_enhance":    "0",
		"count":         "1",
		"clientVersion": "1",
		"org_id":        userID,
		"use_plugin":    "false",
		"metadata":      string(metadata),
	}
	if pid := strings.TrimSpace(cr.ParentID); pid != "" {
		fields["parent_id"] = pid
	}

	var buf bytes.Buffer
	w := multipart.NewWriter(&buf)
	for k, v := range fields {
		_ = w.WriteField(k, v)
	}
	_ = w.Close()

	body, status, err := c.apiPost(ctx, cr.Token, apiBase+"/v1/image/generations/upload", w.FormDataContentType(), buf.Bytes())
	if err != nil {
		return nil, nil, fmt.Errorf("%w: submit: %s", ErrTemporaryUpstream, err.Error())
	}
	if status == 401 || status == 403 {
		return nil, nil, ErrAuth
	}
	if status == 402 {
		return nil, nil, ErrQuotaExhausted
	}
	if status != 200 && status != 201 {
		return nil, nil, fmt.Errorf("%w: submit http %d: %s", ErrTemporaryUpstream, status, clip(body, 200))
	}
	var jobs []struct {
		BatchID string `json:"batchId"`
		ID      string `json:"id"`
		Status  string `json:"status"`
	}
	if err := json.Unmarshal(body, &jobs); err != nil || len(jobs) == 0 || jobs[0].BatchID == "" {
		return nil, nil, fmt.Errorf("%w: no batch id: %s", ErrTemporaryUpstream, clip(body, 200))
	}
	batchID := jobs[0].BatchID

	imageURL, err := c.pollImage(ctx, cr.Token, userID, batchID)
	if err != nil {
		return nil, nil, err
	}
	data, err := c.download(ctx, imageURL)
	if err != nil {
		return nil, nil, err
	}
	return data, map[string]any{"batch_id": batchID, "image_url": imageURL, "org_id": userID}, nil
}

// pollImage polls the org objects feed until the entry for our batch finishes,
// then extracts its asset URL (image_url is a JSON-encoded array string).
func (c *Client) pollImage(ctx context.Context, token, userID, batchID string) (string, error) {
	ticker := time.NewTicker(3 * time.Second)
	defer ticker.Stop()
	// Poll for the full generation budget (caller's genCtx), leaving headroom for
	// the download, instead of a shorter hardcoded cap that killed slow jobs early.
	deadline := time.Now().Add(4 * time.Minute)
	if dl, ok := ctx.Deadline(); ok {
		deadline = dl.Add(-60 * time.Second)
	}

	url := teamsBase + "/v1/org/" + userID + "/objects?batch=true&limit=50&service=image,chat-image"
	for {
		body, status, err := c.apiGetP(ctx, token, url, false)
		if err == nil && status == 200 {
			var resp struct {
				Data []struct {
					BatchID string `json:"batch_id"`
					Status  string `json:"status"`
					Code    int    `json:"code"`
					// Current shape: the produced asset lives at url.generation[0].
					URL struct {
						Generation []string `json:"generation"`
					} `json:"url"`
					ImageURL string `json:"image_url"` // legacy fallback
				} `json:"data"`
			}
			if json.Unmarshal(body, &resp) == nil {
				for _, o := range resp.Data {
					if o.BatchID != batchID {
						continue
					}
					st := strings.ToLower(strings.TrimSpace(o.Status))
					switch {
					case st == "finished" || o.Code == 2:
						if u := firstNonEmpty(o.URL.Generation); u != "" {
							return u, nil
						}
						if u := firstImageURL(o.ImageURL); u != "" {
							return u, nil
						}
					case st == "failed" || st == "error":
						return "", fmt.Errorf("%w: job %s", ErrTemporaryUpstream, st)
					}
				}
			}
		} else if status == 401 || status == 403 {
			return "", ErrAuth
		}
		if time.Now().After(deadline) {
			return "", fmt.Errorf("%w: generation timed out", ErrTemporaryUpstream)
		}
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-ticker.C:
		}
	}
}

// firstNonEmpty returns the first non-blank string in a slice.
func firstNonEmpty(ss []string) string {
	for _, s := range ss {
		if strings.TrimSpace(s) != "" {
			return strings.TrimSpace(s)
		}
	}
	return ""
}

// firstImageURL parses the image_url field — a JSON-encoded array of URLs — and
// returns the first one. Tolerates a bare string too.
func firstImageURL(raw string) string {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return ""
	}
	var urls []string
	if json.Unmarshal([]byte(raw), &urls) == nil {
		for _, u := range urls {
			if strings.TrimSpace(u) != "" {
				return strings.TrimSpace(u)
			}
		}
		return ""
	}
	if strings.HasPrefix(raw, "http") {
		return raw
	}
	return ""
}

func (c *Client) download(ctx context.Context, url string) ([]byte, error) {
	client, err := c.newDirectTLSClient()
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	req = req.WithContext(ctx)
	req.Header = http.Header{
		"accept":     {"image/avif,image/webp,image/png,image/*,*/*;q=0.8"},
		"user-agent": {userAgent},
		"referer":    {webOrigin + "/"},
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	b, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("%w: image download http %d", ErrTemporaryUpstream, resp.StatusCode)
	}
	return b, nil
}

// apiPost issues a POST with a raw body + content-type, carrying the bearer token.
func (c *Client) apiPost(ctx context.Context, token, url, contentType string, body []byte) ([]byte, int, error) {
	client, err := c.newTLSClient()
	if err != nil {
		return nil, 0, err
	}
	req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return nil, 0, err
	}
	req = req.WithContext(ctx)
	req.Header = http.Header{
		"accept":        {"application/json, text/plain, */*"},
		"authorization": {"Bearer " + token},
		"content-type":  {contentType},
		"origin":        {webOrigin},
		"referer":       {webOrigin + "/"},
		"user-agent":    {userAgent},
		http.HeaderOrderKey: {
			"accept", "authorization", "content-type", "origin", "referer", "user-agent",
		},
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, 0, err
	}
	defer resp.Body.Close()
	b, err := io.ReadAll(resp.Body)
	return b, resp.StatusCode, err
}