new-api/relay/helper/derived_funcs.go

package helper

import (
	"fmt"
	"math"
	"strconv"
	"strings"
)

func init() {
	registerDerivedPricingFunc("megapixels", derivedMegapixels)
	registerDerivedPricingFunc("aspect_ratio", derivedAspectRatio)
	registerDerivedPricingFunc("input_mode", derivedInputMode)
}

func derivedMegapixels(req, resp map[string]any) (any, error) {
	resolution, ok, _ := getNestedValue(req, "resolution")
	if !ok || !hasNonEmptyValue(resolution) {
		resolution, ok, _ = getNestedValue(req, "options.size")
	}
	if !ok || !hasNonEmptyValue(resolution) {
		return nil, nil
	}
	width, height, err := parseResolution(fmt.Sprintf("%v", resolution))
	if err != nil {
		return nil, err
	}
	return float64(width*height) / 1_000_000, nil
}

func derivedAspectRatio(req, resp map[string]any) (any, error) {
	resolution, ok, _ := getNestedValue(req, "resolution")
	if !ok || !hasNonEmptyValue(resolution) {
		resolution, ok, _ = getNestedValue(req, "options.size")
	}
	if !ok || !hasNonEmptyValue(resolution) {
		return nil, nil
	}
	width, height, err := parseResolution(fmt.Sprintf("%v", resolution))
	if err != nil {
		return nil, err
	}
	if height == 0 {
		return nil, fmt.Errorf("resolution height must be greater than zero")
	}
	gcd := intGCD(width, height)
	return fmt.Sprintf("%d:%d", width/gcd, height/gcd), nil
}

func derivedInputMode(req, resp map[string]any) (any, error) {
	if hasAnyNonEmpty(req, "video", "video_url", "reference_video", "first_frame_video") {
		return "v2v", nil
	}
	if hasAnyNonEmpty(req, "image", "image_url", "image[]", "images", "first_frame_image", "last_frame_image") {
		return "i2v", nil
	}
	return "t2v", nil
}

func hasNonEmptyValue(v any) bool {
	switch value := v.(type) {
	case nil:
		return false
	case string:
		return strings.TrimSpace(value) != ""
	case []any:
		return len(value) > 0
	case []string:
		return len(value) > 0
	case map[string]any:
		return len(value) > 0
	default:
		return true
	}
}

func hasAnyNonEmpty(req map[string]any, keys ...string) bool {
	for _, key := range keys {
		value, ok, _ := getNestedValue(req, key)
		if ok && hasNonEmptyValue(value) {
			return true
		}
	}
	return false
}

func parseResolution(raw string) (int, int, error) {
	normalized := strings.ToLower(strings.TrimSpace(raw))
	if strings.HasSuffix(normalized, "p") {
		height, err := strconv.Atoi(strings.TrimSuffix(normalized, "p"))
		if err != nil || height <= 0 {
			return 0, 0, fmt.Errorf("invalid resolution %q", raw)
		}
		width := int(math.Round(float64(height) * 16 / 9))
		return width, height, nil
	}

	parts := strings.FieldsFunc(normalized, func(r rune) bool {
		return r == 'x' || r == '*' || r == '×'
	})
	if len(parts) != 2 {
		return 0, 0, fmt.Errorf("invalid resolution %q", raw)
	}
	width, err := strconv.Atoi(strings.TrimSpace(parts[0]))
	if err != nil || width <= 0 {
		return 0, 0, fmt.Errorf("invalid resolution width %q", raw)
	}
	height, err := strconv.Atoi(strings.TrimSpace(parts[1]))
	if err != nil || height <= 0 {
		return 0, 0, fmt.Errorf("invalid resolution height %q", raw)
	}
	return width, height, nil
}

func intGCD(a, b int) int {
	for b != 0 {
		a, b = b, a%b
	}
	if a < 0 {
		return -a
	}
	if a == 0 {
		return 1
	}
	return a
}