new-api/relay/helper/derived_funcs.go

package helper

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

	"github.com/QuantumNous/new-api/common"
)

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

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 derivedVideoInput(req, resp map[string]any) (any, error) {
	if hasAnyNonEmpty(req, "video", "video_url", "reference_video", "first_frame_video") {
		return true, nil
	}
	metadata, ok, _ := getNestedValue(req, "metadata")
	if !ok || !hasNonEmptyValue(metadata) {
		return false, nil
	}
	metadataMap, err := normalizeMetadataMap(metadata)
	if err != nil {
		return nil, err
	}
	return metadataContainsVideoInput(metadataMap), nil
}

func normalizeMetadataMap(raw any) (map[string]any, error) {
	switch value := raw.(type) {
	case map[string]any:
		return value, nil
	case string:
		if strings.TrimSpace(value) == "" {
			return nil, nil
		}
		metadata := map[string]any{}
		if err := common.Unmarshal([]byte(value), &metadata); err != nil {
			return nil, err
		}
		return metadata, nil
	default:
		return nil, nil
	}
}

func metadataContainsVideoInput(metadata map[string]any) bool {
	if metadata == nil {
		return false
	}
	content, ok := metadata["content"].([]any)
	if !ok {
		return false
	}
	for _, item := range content {
		m, ok := item.(map[string]any)
		if !ok {
			continue
		}
		if typ, _ := m["type"].(string); typ == "video_url" {
			return true
		}
		if value, ok := m["video_url"]; ok && hasNonEmptyValue(value) {
			return true
		}
	}
	return false
}

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
}