easyai-ai-gateway/apps/api/internal/clients/vectorizer.go

package clients

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"mime/multipart"
	"net"
	"net/http"
	"net/textproto"
	"net/url"
	"strings"
	"time"
)

const vectorizerMaxResponseBytes = 128 << 20

// VectorizerClient implements the Vectorizer.AI binary vectorization API.
// Gateway async semantics are supplied by the outer River-backed task runner.
type VectorizerClient struct {
	HTTPClient *http.Client
	LookupIP   func(context.Context, string) ([]net.IPAddr, error)
}

func (c VectorizerClient) Run(ctx context.Context, request Request) (Response, error) {
	startedAt := time.Now()
	format := strings.ToLower(strings.TrimSpace(firstNonEmptyString(request.Body["format"], request.Body["output_format"])))
	if format == "" {
		format = "svg"
	}
	if !vectorizerFormatAllowed(format) {
		return Response{}, &ClientError{Code: "invalid_parameter", Message: "vectorizer format must be svg, eps, pdf, dxf, or png", Param: "format", StatusCode: http.StatusBadRequest}
	}

	imageToken := strings.TrimSpace(firstNonEmptyString(request.Body["_vectorizer_image_token"]))
	receipt := strings.TrimSpace(firstNonEmptyString(request.Body["_vectorizer_receipt"]))
	endpoint := "vectorize"
	fields := map[string]string{"output.file_format": format}
	if imageToken != "" {
		endpoint = "download"
		fields["image.token"] = imageToken
		if receipt != "" {
			fields["receipt"] = receipt
		}
	} else {
		imageURL := vectorizerImageURL(request.Body)
		if imageURL == "" {
			return Response{}, &ClientError{Code: "invalid_parameter", Message: "vectorizer source image URL is required", Param: "source.url", StatusCode: http.StatusBadRequest}
		}
		if err := c.validateSourceURL(ctx, imageURL, boolishDefault(request.Candidate.PlatformConfig["allowPrivateSourceDownloads"], false)); err != nil {
			return Response{}, err
		}
		fields["image.url"] = imageURL
		fields["mode"] = firstNonEmptyString(request.Candidate.PlatformConfig["mode"], "production")
		fields["policy.retention_days"] = firstNonEmptyString(request.Candidate.PlatformConfig["retentionDays"], request.Candidate.PlatformConfig["retention_days"], "7")
		fields["processing.shapes.min_area_px"] = vectorizerCleanupMinArea(request.Body)
		if maxColors := vectorizerMaxColors(request.Body); maxColors > 0 {
			fields["processing.max_colors"] = fmt.Sprint(maxColors)
		}
	}
	appendVectorizerOutputFields(fields, format)

	payload, contentType, err := vectorizerMultipartBody(fields)
	if err != nil {
		return Response{}, err
	}
	url := providerURL(request.Candidate.BaseURL, endpoint)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, payload)
	if err != nil {
		return Response{}, err
	}
	req.Header.Set("Content-Type", contentType)
	applyVectorizerAuth(req, request.Candidate.Credentials)
	client := httpClient(request.HTTPClient, c.HTTPClient)
	resp, err := client.Do(req)
	if err != nil {
		return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
	}
	defer resp.Body.Close()
	requestID := requestIDFromHTTPResponse(resp)
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		result, decodeErr := decodeHTTPResponse(resp)
		if decodeErr != nil {
			return Response{}, annotateResponseError(decodeErr, requestID, startedAt, time.Now())
		}
		return Response{}, &ClientError{Code: "provider_failed", Message: firstNonEmptyString(result["message"], result["error"], resp.Status), RequestID: requestID, StatusCode: resp.StatusCode, Retryable: resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500}
	}
	raw, err := io.ReadAll(io.LimitReader(resp.Body, vectorizerMaxResponseBytes+1))
	if err != nil {
		return Response{}, &ClientError{Code: "invalid_response", Message: err.Error(), RequestID: requestID, Retryable: true}
	}
	if len(raw) == 0 || len(raw) > vectorizerMaxResponseBytes {
		return Response{}, &ClientError{Code: "invalid_response", Message: "vectorizer response is empty or too large", RequestID: requestID, Retryable: false}
	}
	returnedToken := strings.TrimSpace(resp.Header.Get("X-Image-Token"))
	returnedReceipt := strings.TrimSpace(resp.Header.Get("X-Receipt"))
	if request.OnRemoteTaskSubmitted != nil && returnedToken != "" {
		digest := sha256.Sum256([]byte(returnedToken))
		if err := request.OnRemoteTaskSubmitted("vectorizer-"+hex.EncodeToString(digest[:8]), map[string]any{
			"imageToken": returnedToken,
			"receipt":    returnedReceipt,
		}); err != nil {
			return Response{}, err
		}
	}
	finishedAt := time.Now()
	mimeType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
	if mimeType == "" || mimeType == "application/octet-stream" {
		mimeType = vectorizerContentType(format)
	}
	return Response{
		Result: map[string]any{
			"status": "success",
			"model":  request.Model,
			"data": []any{map[string]any{
				"type":      vectorizerOutputKind(format),
				"b64_json":  base64.StdEncoding.EncodeToString(raw),
				"mime_type": mimeType,
				"format":    format,
			}},
			"vectorizer": map[string]any{
				"format":            format,
				"creditsCharged":    numericHeader(resp.Header.Get("X-Credits-Charged")),
				"creditsCalculated": numericHeader(resp.Header.Get("X-Credits-Calculated")),
			},
		},
		RequestID:          requestID,
		Progress:           append(providerProgress(request), Progress{Phase: "uploading", Progress: 0.9, Message: "vectorizer result received"}),
		ResponseStartedAt:  startedAt,
		ResponseFinishedAt: finishedAt,
		ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
	}, nil
}

func (c VectorizerClient) validateSourceURL(ctx context.Context, rawURL string, allowPrivate bool) error {
	parsed, err := url.Parse(rawURL)
	if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || strings.TrimSpace(parsed.Hostname()) == "" || parsed.User != nil {
		return &ClientError{Code: "invalid_parameter", Message: "vectorizer source URL must be a public http(s) URL without userinfo", Param: "source.url", StatusCode: http.StatusBadRequest}
	}
	if allowPrivate {
		return nil
	}
	lookup := c.LookupIP
	if lookup == nil {
		lookup = net.DefaultResolver.LookupIPAddr
	}
	addresses, err := lookup(ctx, parsed.Hostname())
	if err != nil || len(addresses) == 0 {
		return &ClientError{Code: "invalid_parameter", Message: "vectorizer source DNS resolution failed", Param: "source.url", StatusCode: http.StatusBadRequest}
	}
	for _, address := range addresses {
		if topazBlockedAddress(address.IP) {
			return &ClientError{Code: "invalid_parameter", Message: "vectorizer source resolved to a blocked network", Param: "source.url", StatusCode: http.StatusBadRequest}
		}
	}
	return nil
}

func vectorizerImageURL(body map[string]any) string {
	if source, ok := body["source"].(map[string]any); ok {
		return firstNonEmptyString(source["url"], source["image_url"], source["imageUrl"])
	}
	return firstNonEmptyString(body["image_url"], body["imageUrl"])
}

func vectorizerMultipartBody(fields map[string]string) (*bytes.Buffer, string, error) {
	var payload bytes.Buffer
	writer := multipart.NewWriter(&payload)
	for key, value := range fields {
		header := make(textproto.MIMEHeader)
		header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"`, strings.ReplaceAll(key, `"`, `\"`)))
		part, err := writer.CreatePart(header)
		if err != nil {
			return nil, "", err
		}
		if _, err := io.WriteString(part, value); err != nil {
			return nil, "", err
		}
	}
	if err := writer.Close(); err != nil {
		return nil, "", err
	}
	return &payload, writer.FormDataContentType(), nil
}

func applyVectorizerAuth(req *http.Request, credentials map[string]any) {
	username := credential(credentials, "username", "accessKey", "access_key", "apiId", "api_id", "id")
	password := credential(credentials, "password", "secretKey", "secret_key", "apiSecret", "api_secret", "secret")
	if username != "" || password != "" {
		req.SetBasicAuth(username, password)
		return
	}
	if apiKey := credential(credentials, "apiKey", "api_key", "token"); apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+apiKey)
	}
}

func vectorizerFormatAllowed(format string) bool {
	switch format {
	case "svg", "eps", "pdf", "dxf", "png":
		return true
	default:
		return false
	}
}

func vectorizerCleanupMinArea(body map[string]any) string {
	cleanup := strings.ToLower(strings.TrimSpace(firstNonEmptyString(body["cleanupLevel"], body["cleanup_level"])))
	switch cleanup {
	case "low":
		return "0"
	case "strong":
		return "1"
	default:
		return "0.125"
	}
}

func vectorizerMaxColors(body map[string]any) int {
	for _, key := range []string{"maxColors", "max_colors"} {
		if value := int(numericValue(body[key], 0)); value == 0 || value == 2 || value == 4 || value == 8 || value == 16 || value == 32 {
			return value
		}
	}
	return 0
}

func appendVectorizerOutputFields(fields map[string]string, format string) {
	if format == "svg" {
		fields["output.svg.version"] = "svg_1_1"
		fields["output.svg.fixed_size"] = "false"
		fields["output.svg.adobe_compatibility_mode"] = "true"
	}
	if format == "dxf" {
		fields["output.dxf.compatibility_level"] = "lines_and_arcs"
	}
}

func vectorizerContentType(format string) string {
	switch format {
	case "svg":
		return "image/svg+xml"
	case "eps":
		return "application/postscript"
	case "pdf":
		return "application/pdf"
	case "dxf":
		return "application/dxf"
	default:
		return "image/png"
	}
}

func vectorizerOutputKind(format string) string {
	if format == "svg" || format == "png" {
		return "image"
	}
	return "file"
}

func numericHeader(value string) any {
	value = strings.TrimSpace(value)
	if value == "" {
		return nil
	}
	return numericValue(value, 0)
}