new-api/relay/helper/dimension_resolver.go

package helper

import (
	"fmt"
	"io"
	"mime/multipart"
	"strconv"
	"strings"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/types"
	"github.com/gin-gonic/gin"
)

type derivedPricingFunc func(req, resp map[string]any) (any, error)

var derivedPricingFuncs = map[string]derivedPricingFunc{}

type ResolvedDimensionValues struct {
	Values   map[string]any
	Missing  map[string]bool
	RawInput map[string]any
}

type DimensionResolver struct {
	c        *gin.Context
	request  map[string]any
	response map[string]any
}

func NewDimensionResolver(c *gin.Context) (*DimensionResolver, error) {
	contentType := c.Request.Header.Get("Content-Type")
	req := map[string]any{}

	switch {
	case strings.HasPrefix(contentType, "application/json"):
		storage, err := common.GetBodyStorage(c)
		if err != nil {
			return nil, err
		}
		body, err := storage.Bytes()
		if err != nil {
			return nil, err
		}
		if len(strings.TrimSpace(string(body))) > 0 {
			if err := common.Unmarshal(body, &req); err != nil {
				return nil, err
			}
		}
		if err := resetRequestBody(c, storage); err != nil {
			return nil, err
		}
	case strings.Contains(contentType, gin.MIMEMultipartPOSTForm):
		form, err := common.ParseMultipartFormReusable(c)
		if err != nil {
			return nil, err
		}
		req = requestMapFromMultipart(form)
	default:
		return nil, fmt.Errorf("unsupported content-type for pricing dimensions: %s", contentType)
	}

	return &DimensionResolver{
		c:        c,
		request:  req,
		response: map[string]any{},
	}, nil
}

func (r *DimensionResolver) ResolveDimensions(dims []types.PricingDimension) (*ResolvedDimensionValues, error) {
	resolved := &ResolvedDimensionValues{
		Values:   make(map[string]any, len(dims)),
		Missing:  make(map[string]bool),
		RawInput: types.CloneMapAny(r.request),
	}

	for _, dim := range dims {
		value, exists, err := r.resolveSource(dim.Source)
		if err != nil {
			return nil, err
		}
		if !exists || !hasNonEmptyValue(value) {
			if dim.Optional {
				resolved.Values[dim.Key] = dim.Default
				continue
			}
			resolved.Values[dim.Key] = nil
			resolved.Missing[dim.Key] = true
			continue
		}
		converted, err := convertPricingDimensionValue(value, dim.Type)
		if err != nil {
			return nil, fmt.Errorf("dimension %q: %w", dim.Key, err)
		}
		resolved.Values[dim.Key] = converted
	}

	return resolved, nil
}

func IsAllowedPricingSource(source string) bool {
	if strings.HasPrefix(source, "response.") {
		return false
	}
	if strings.HasPrefix(source, "request.") {
		return len(strings.TrimPrefix(source, "request.")) > 0
	}
	if _, ok := derivedPricingFuncs[strings.TrimPrefix(source, "derived.")]; ok && strings.HasPrefix(source, "derived.") {
		return true
	}
	return false
}

func registerDerivedPricingFunc(name string, fn derivedPricingFunc) {
	derivedPricingFuncs[name] = fn
}

func (r *DimensionResolver) resolveSource(source string) (any, bool, error) {
	if strings.HasPrefix(source, "response.") {
		return nil, false, fmt.Errorf("pricing source %q is unsupported in schema v1", source)
	}
	if strings.HasPrefix(source, "request.") {
		return getNestedValue(r.request, strings.TrimPrefix(source, "request."))
	}
	if strings.HasPrefix(source, "derived.") {
		name := strings.TrimPrefix(source, "derived.")
		fn, ok := derivedPricingFuncs[name]
		if !ok {
			return nil, false, fmt.Errorf("derived pricing source %q is not registered", source)
		}
		value, err := fn(r.request, r.response)
		if err != nil {
			return nil, false, err
		}
		return value, hasNonEmptyValue(value), nil
	}
	return nil, false, fmt.Errorf("pricing source %q is not allowed", source)
}

func getNestedValue(src map[string]any, path string) (any, bool, error) {
	if path == "" {
		return nil, false, nil
	}
	var cur any = src
	for _, part := range strings.Split(path, ".") {
		m, ok := cur.(map[string]any)
		if !ok {
			return nil, false, nil
		}
		cur, ok = m[part]
		if !ok {
			return nil, false, nil
		}
	}
	return cur, true, nil
}

func convertPricingDimensionValue(value any, typ string) (any, error) {
	switch typ {
	case "string":
		if s, ok := value.(string); ok {
			return s, nil
		}
		return fmt.Sprintf("%v", value), nil
	case "number":
		switch v := value.(type) {
		case float64:
			return v, nil
		case float32:
			return float64(v), nil
		case int:
			return float64(v), nil
		case int64:
			return float64(v), nil
		case int32:
			return float64(v), nil
		case string:
			f, err := strconv.ParseFloat(strings.TrimSpace(v), 64)
			if err != nil {
				return nil, fmt.Errorf("value %q is not a number", v)
			}
			return f, nil
		default:
			return nil, fmt.Errorf("value is not a number")
		}
	case "boolean":
		switch v := value.(type) {
		case bool:
			return v, nil
		case string:
			b, err := strconv.ParseBool(strings.TrimSpace(v))
			if err != nil {
				return nil, fmt.Errorf("value %q is not a boolean", v)
			}
			return b, nil
		default:
			return nil, fmt.Errorf("value is not a boolean")
		}
	default:
		return nil, fmt.Errorf("unknown dimension type %q", typ)
	}
}

func requestMapFromMultipart(form *multipart.Form) map[string]any {
	req := map[string]any{}
	for key, vals := range form.Value {
		if len(vals) == 1 {
			assignNestedValue(req, key, vals[0])
			continue
		}
		copied := append([]string(nil), vals...)
		assignNestedValue(req, key, copied)
	}
	for key, files := range form.File {
		if len(files) == 1 {
			assignNestedValue(req, key, files[0].Filename)
			continue
		}
		names := make([]string, 0, len(files))
		for _, file := range files {
			names = append(names, file.Filename)
		}
		assignNestedValue(req, key, names)
	}
	return req
}

func assignNestedValue(dst map[string]any, key string, value any) {
	parts := strings.Split(key, ".")
	cur := dst
	for i, part := range parts {
		if i == len(parts)-1 {
			cur[part] = value
			return
		}
		next, ok := cur[part].(map[string]any)
		if !ok {
			next = map[string]any{}
			cur[part] = next
		}
		cur = next
	}
}

func resetRequestBody(c *gin.Context, storage common.BodyStorage) error {
	if _, err := storage.Seek(0, io.SeekStart); err != nil {
		return err
	}
	c.Request.Body = io.NopCloser(storage)
	return nil
}