sub2api/backend/internal/service/openai_reasoning_effort_policy.go

package service

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"github.com/Wei-Shaw/sub2api/internal/domain"
	"github.com/tidwall/gjson"
	"github.com/tidwall/sjson"
)

const (
	maxReasoningEffortMappings = 64
	maxReasoningEffortValueLen = 64
	maxReasoningEffortModelLen = 200

	// ReasoningEffortOverLimitDowngrade rewrites values above the ceiling.
	ReasoningEffortOverLimitDowngrade = "downgrade"
	// ReasoningEffortOverLimitDeny rejects the request when the ceiling is exceeded.
	ReasoningEffortOverLimitDeny = "deny"
	// ReasoningEffortMappingDeny is a mapping target that rejects the request
	// when the matching source reasoning effort is present.
	ReasoningEffortMappingDeny = "deny"
)

var openAIReasoningEffortValues = []string{"minimal", "low", "medium", "high", "xhigh", "max"}
var anthropicReasoningEffortValues = []string{"low", "medium", "high", "xhigh", "max"}

func normalizeReasoningEffortMappingSource(raw string) string {
	if strings.EqualFold(strings.TrimSpace(raw), "none") {
		return "none"
	}
	return NormalizeMaxReasoningEffort(raw)
}

func normalizeReasoningEffortMappingTarget(raw string) string {
	if isReasoningEffortMappingDeny(raw) {
		return ReasoningEffortMappingDeny
	}
	return NormalizeMaxReasoningEffort(raw)
}

func isReasoningEffortMappingDeny(raw string) bool {
	return strings.EqualFold(strings.TrimSpace(raw), ReasoningEffortMappingDeny)
}

type openAIReasoningEffortPolicyContextKey struct{}
type requestedReasoningEffortContextKey struct{}

type openAIReasoningEffortPolicy struct {
	maxEffort string
	overLimit string
	mappings  []ReasoningEffortMapping
}

// ReasoningEffortOverLimitError is returned when a group policy is set to deny
// requests whose explicit reasoning effort exceeds the ceiling.
type ReasoningEffortOverLimitError struct {
	Requested string
	Max       string
}

func (e *ReasoningEffortOverLimitError) Error() string {
	if e == nil {
		return "reasoning effort exceeds this group's limit"
	}
	requested := strings.TrimSpace(e.Requested)
	max := strings.TrimSpace(e.Max)
	if requested == "" && max == "" {
		return "reasoning effort exceeds this group's limit"
	}
	if requested == "" {
		return fmt.Sprintf("reasoning effort exceeds this group's limit of %q", max)
	}
	if max == "" {
		return fmt.Sprintf("reasoning effort %q exceeds this group's limit", requested)
	}
	return fmt.Sprintf("reasoning effort %q exceeds this group's limit of %q", requested, max)
}

// ReasoningEffortMappingDeniedError is returned when a group mapping target is
// set to deny the explicit reasoning effort on the request.
type ReasoningEffortMappingDeniedError struct {
	Requested string
}

func (e *ReasoningEffortMappingDeniedError) Error() string {
	if e == nil {
		return "reasoning effort is denied by this group's mapping policy"
	}
	requested := strings.TrimSpace(e.Requested)
	if requested == "" {
		return "reasoning effort is denied by this group's mapping policy"
	}
	return fmt.Sprintf("reasoning effort %q is denied by this group's mapping policy", requested)
}

// IsReasoningEffortPolicyDenied reports whether err is a local reasoning-effort
// policy rejection (ceiling deny or mapping deny).
func IsReasoningEffortPolicyDenied(err error) bool {
	if err == nil {
		return false
	}
	var overLimit *ReasoningEffortOverLimitError
	if errors.As(err, &overLimit) {
		return true
	}
	var mappingDenied *ReasoningEffortMappingDeniedError
	return errors.As(err, &mappingDenied)
}

// NormalizeMaxReasoningEffort validates and canonicalizes a group policy value.
// Empty means that the group does not impose a ceiling.
func NormalizeMaxReasoningEffort(raw string) string {
	value := strings.ToLower(strings.TrimSpace(raw))
	value = strings.NewReplacer("-", "", "_", "", " ", "").Replace(value)
	switch value {
	case "":
		return ""
	case "minimal":
		return "minimal"
	case "low":
		return "low"
	case "medium":
		return "medium"
	case "high":
		return "high"
	case "xhigh", "extrahigh":
		return "xhigh"
	case "max":
		return "max"
	default:
		return ""
	}
}

func reasoningEffortValuesForPlatform(platform string) []string {
	switch platform {
	case PlatformOpenAI, PlatformComposite:
		return openAIReasoningEffortValues
	case PlatformAnthropic:
		return anthropicReasoningEffortValues
	default:
		return nil
	}
}

func normalizeMaxReasoningEffortForPlatform(platform, raw string) (string, error) {
	if strings.TrimSpace(raw) == "" {
		return "", nil
	}

	allowedValues := reasoningEffortValuesForPlatform(platform)
	if len(allowedValues) == 0 {
		return "", fmt.Errorf(
			"reasoning effort policy is only supported for platforms %q, %q, and %q",
			PlatformAnthropic,
			PlatformOpenAI,
			PlatformComposite,
		)
	}

	value := NormalizeMaxReasoningEffort(raw)
	for _, allowed := range allowedValues {
		if value == allowed {
			return value, nil
		}
	}
	return "", fmt.Errorf(
		"reasoning effort %q is not supported for platform %q; allowed values: %s",
		raw,
		platform,
		strings.Join(allowedValues, ", "),
	)
}

// NormalizeMaxReasoningEffortOverLimit canonicalizes the over-limit action.
// Empty means the historical default: automatically downgrade to the ceiling.
func NormalizeMaxReasoningEffortOverLimit(raw string) string {
	switch strings.ToLower(strings.TrimSpace(raw)) {
	case "", ReasoningEffortOverLimitDowngrade:
		return ReasoningEffortOverLimitDowngrade
	case ReasoningEffortOverLimitDeny:
		return ReasoningEffortOverLimitDeny
	default:
		return ""
	}
}

func normalizeMaxReasoningEffortOverLimitForPlatform(platform, raw string) (string, error) {
	value := NormalizeMaxReasoningEffortOverLimit(raw)
	if value == "" {
		return "", fmt.Errorf(
			"reasoning effort over-limit action %q is not supported; allowed values: %s, %s",
			raw,
			ReasoningEffortOverLimitDowngrade,
			ReasoningEffortOverLimitDeny,
		)
	}
	if value == ReasoningEffortOverLimitDowngrade {
		return value, nil
	}
	if len(reasoningEffortValuesForPlatform(platform)) == 0 {
		return "", fmt.Errorf(
			"reasoning effort policy is only supported for platforms %q, %q, and %q",
			PlatformAnthropic,
			PlatformOpenAI,
			PlatformComposite,
		)
	}
	return value, nil
}

func reasoningEffortRank(raw string) (int, bool) {
	switch NormalizeMaxReasoningEffort(raw) {
	case "minimal":
		return 1, true
	case "low":
		return 2, true
	case "medium":
		return 3, true
	case "high":
		return 4, true
	case "xhigh":
		return 5, true
	case "max":
		return 6, true
	default:
		return 0, false
	}
}

func normalizeReasoningEffortMatchType(matchType, model string) (string, error) {
	model = strings.TrimSpace(model)
	matchType = strings.ToLower(strings.TrimSpace(matchType))
	if model == "" {
		return "", nil
	}
	switch matchType {
	case "", domain.ReasoningEffortMatchExact:
		return domain.ReasoningEffortMatchExact, nil
	case domain.ReasoningEffortMatchPrefix:
		return domain.ReasoningEffortMatchPrefix, nil
	case domain.ReasoningEffortMatchSuffix:
		return domain.ReasoningEffortMatchSuffix, nil
	default:
		return "", fmt.Errorf("invalid match_type %q", matchType)
	}
}

func reasoningEffortMappingDuplicateKey(from, matchType, model string) string {
	return from + "\x00" + matchType + "\x00" + strings.ToLower(strings.TrimSpace(model))
}

func requestModelMatchesReasoningEffortMapping(requestModel, matchType, mappingModel string) bool {
	scope := strings.ToLower(strings.TrimSpace(mappingModel))
	req := strings.ToLower(strings.TrimSpace(requestModel))
	switch matchType {
	case "":
		return true
	case domain.ReasoningEffortMatchExact:
		return scope != "" && scope == req
	case domain.ReasoningEffortMatchPrefix:
		return scope != "" && req != "" && strings.HasPrefix(req, scope)
	case domain.ReasoningEffortMatchSuffix:
		return scope != "" && req != "" && strings.HasSuffix(req, scope)
	default:
		return false
	}
}

func selectReasoningEffortMapping(mappings []ReasoningEffortMapping, from, requestModel string) (ReasoningEffortMapping, bool) {
	type candidate struct {
		mapping       ReasoningEffortMapping
		matchStrength int
		patternLen    int
		index         int
	}
	candidates := make([]candidate, 0, len(mappings))
	for i, mapping := range mappings {
		if normalizeReasoningEffortMappingSource(mapping.From) != from {
			continue
		}
		model := strings.TrimSpace(mapping.Model)
		matchType, err := normalizeReasoningEffortMatchType(mapping.MatchType, model)
		if err != nil {
			continue
		}
		if !requestModelMatchesReasoningEffortMapping(requestModel, matchType, model) {
			continue
		}
		strength := 1
		patternLen := 0
		switch matchType {
		case domain.ReasoningEffortMatchExact:
			strength = 3
		case domain.ReasoningEffortMatchPrefix, domain.ReasoningEffortMatchSuffix:
			strength = 2
			patternLen = len(strings.ToLower(model))
		}
		candidates = append(candidates, candidate{
			mapping:       mapping,
			matchStrength: strength,
			patternLen:    patternLen,
			index:         i,
		})
	}
	if len(candidates) == 0 {
		return ReasoningEffortMapping{}, false
	}
	best := candidates[0]
	for _, item := range candidates[1:] {
		if item.matchStrength != best.matchStrength {
			if item.matchStrength > best.matchStrength {
				best = item
			}
			continue
		}
		if item.patternLen != best.patternLen {
			if item.patternLen > best.patternLen {
				best = item
			}
			continue
		}
		if item.index < best.index {
			best = item
		}
	}
	return best.mapping, true
}

// NormalizeReasoningEffortMappings validates group mapping rules against the
// fixed effort values supported by the selected platform. Optional model scopes are
// canonicalized to exact/prefix/suffix match; empty type and model keep the
// global all-models rule.
func NormalizeReasoningEffortMappings(platform string, raw []ReasoningEffortMapping) ([]ReasoningEffortMapping, error) {
	if len(raw) > maxReasoningEffortMappings {
		return nil, fmt.Errorf("reasoning effort mappings cannot exceed %d entries", maxReasoningEffortMappings)
	}

	normalized := make([]ReasoningEffortMapping, 0, len(raw))
	seen := make(map[string]struct{}, len(raw))
	for i, mapping := range raw {
		from := normalizeReasoningEffortMappingSource(mapping.From)
		to := normalizeReasoningEffortMappingTarget(mapping.To)
		if from == "" || to == "" {
			return nil, fmt.Errorf("reasoning effort mapping %d contains an empty or unknown value", i+1)
		}
		if len(from) > maxReasoningEffortValueLen || len(to) > maxReasoningEffortValueLen {
			return nil, fmt.Errorf("reasoning effort mapping %d values cannot exceed %d characters", i+1, maxReasoningEffortValueLen)
		}
		if from != "none" {
			if _, err := normalizeMaxReasoningEffortForPlatform(platform, from); err != nil {
				return nil, fmt.Errorf("reasoning effort mapping %d source: %w", i+1, err)
			}
		} else if len(reasoningEffortValuesForPlatform(platform)) == 0 {
			return nil, fmt.Errorf(
				"reasoning effort mapping %d source: reasoning effort policy is only supported for platforms %q and %q",
				i+1, PlatformOpenAI, PlatformComposite,
			)
		}
		if to != ReasoningEffortMappingDeny {
			if _, err := normalizeMaxReasoningEffortForPlatform(platform, to); err != nil {
				return nil, fmt.Errorf("reasoning effort mapping %d target: %w", i+1, err)
			}
		} else if len(reasoningEffortValuesForPlatform(platform)) == 0 {
			return nil, fmt.Errorf(
				"reasoning effort mapping %d target: reasoning effort policy is only supported for platforms %q, %q, and %q",
				i+1,
				PlatformAnthropic,
				PlatformOpenAI,
				PlatformComposite,
			)
		}
		model := strings.TrimSpace(mapping.Model)
		if len(model) > maxReasoningEffortModelLen {
			return nil, fmt.Errorf("reasoning effort mapping %d model cannot exceed %d characters", i+1, maxReasoningEffortModelLen)
		}
		matchType, err := normalizeReasoningEffortMatchType(mapping.MatchType, model)
		if err != nil {
			return nil, fmt.Errorf("reasoning effort mapping %d: %w", i+1, err)
		}
		key := reasoningEffortMappingDuplicateKey(from, matchType, model)
		if _, exists := seen[key]; exists {
			if model == "" {
				return nil, fmt.Errorf("duplicate reasoning effort mapping source %q", from)
			}
			return nil, fmt.Errorf("duplicate reasoning effort mapping source %q for %s model %q", from, matchType, strings.ToLower(model))
		}
		seen[key] = struct{}{}
		normalized = append(normalized, ReasoningEffortMapping{
			From:      from,
			To:        to,
			MatchType: matchType,
			Model:     model,
		})
	}
	return normalized, nil
}

// WithRequestedReasoningEffort stores the client-requested effort captured from
// the inbound body before group policy or model-family remapping.
func WithRequestedReasoningEffort(ctx context.Context, effort string) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	effort = strings.TrimSpace(effort)
	if effort == "" {
		return ctx
	}
	return context.WithValue(ctx, requestedReasoningEffortContextKey{}, effort)
}

// RequestedReasoningEffortFromContext returns the inbound requested effort bound
// to ctx, or nil when none was captured.
func RequestedReasoningEffortFromContext(ctx context.Context) *string {
	if ctx == nil {
		return nil
	}
	value, ok := ctx.Value(requestedReasoningEffortContextKey{}).(string)
	if !ok {
		return nil
	}
	value = strings.TrimSpace(value)
	if value == "" {
		return nil
	}
	return &value
}

// WithOpenAIReasoningEffortPolicy binds a group policy to a request after its
// concrete target platform has been resolved to OpenAI. The policy is copied so
// retries and asynchronous forwarding cannot observe later slice mutations.
func WithOpenAIReasoningEffortPolicy(ctx context.Context, maxEffort string, mappings []ReasoningEffortMapping, overLimit string) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	policy := openAIReasoningEffortPolicy{
		maxEffort: maxEffort,
		overLimit: overLimit,
		mappings:  append([]ReasoningEffortMapping(nil), mappings...),
	}
	return context.WithValue(ctx, openAIReasoningEffortPolicyContextKey{}, policy)
}

// ApplyOpenAIReasoningEffortPolicyFromContext applies a policy previously bound
// to the request. An unbound request is returned byte-for-byte unchanged.
func ApplyOpenAIReasoningEffortPolicyFromContext(ctx context.Context, body []byte) ([]byte, bool, error) {
	if ctx == nil {
		return body, false, nil
	}
	policy, ok := ctx.Value(openAIReasoningEffortPolicyContextKey{}).(openAIReasoningEffortPolicy)
	if !ok {
		return body, false, nil
	}
	return ApplyOpenAIReasoningEffortPolicy(body, policy.maxEffort, policy.mappings, policy.overLimit)
}

func mapReasoningEffort(raw string, mappings []ReasoningEffortMapping, requestModel string) (string, bool) {
	value := strings.TrimSpace(raw)
	canonical := normalizeReasoningEffortMappingSource(value)
	if canonical == "" {
		return value, false
	}
	mapping, ok := selectReasoningEffortMapping(mappings, canonical, requestModel)
	if !ok {
		return value, false
	}
	return strings.TrimSpace(mapping.To), true
}

func sanitizeGroupReasoningEffortPolicy(group *Group) {
	if group == nil {
		return
	}
	maxEffort, maxErr := normalizeMaxReasoningEffortForPlatform(group.Platform, group.MaxReasoningEffort)
	mappings, mappingsErr := NormalizeReasoningEffortMappings(group.Platform, group.ReasoningEffortMappings)
	overLimit := NormalizeMaxReasoningEffortOverLimit(group.MaxReasoningEffortOverLimit)
	if maxErr != nil {
		maxEffort = ""
	}
	if mappingsErr != nil {
		mappings = []ReasoningEffortMapping{}
	}
	if overLimit == "" || (group.Platform != PlatformAnthropic && group.Platform != PlatformOpenAI && group.Platform != PlatformComposite) {
		overLimit = ReasoningEffortOverLimitDowngrade
	}
	group.MaxReasoningEffort = maxEffort
	group.MaxReasoningEffortOverLimit = overLimit
	group.ReasoningEffortMappings = mappings
}

// ApplyReasoningEffortPolicy applies one mapping (optionally scoped to
// the request model by exact name, prefix, or suffix) and then either caps
// known effort levels or rejects the request when the group is configured to
// deny values above the ceiling. A mapping whose target is deny rejects the
// request when that source value is present. Omitted values remain untouched
// so upstream defaults stay in control. It understands both OpenAI and
// Anthropic request field shapes.
func ApplyReasoningEffortPolicy(body []byte, maxEffort string, mappings []ReasoningEffortMapping, overLimit string) ([]byte, bool, error) {
	maxRank, hasMax := reasoningEffortRank(maxEffort)
	if len(body) == 0 || (!hasMax && len(mappings) == 0) {
		return body, false, nil
	}
	deny := hasMax && NormalizeMaxReasoningEffortOverLimit(overLimit) == ReasoningEffortOverLimitDeny
	canonicalMax := NormalizeMaxReasoningEffort(maxEffort)

	requestModel := strings.TrimSpace(gjson.GetBytes(body, "model").String())
	result := body
	changed := false
	for _, path := range []string{"reasoning.effort", "reasoning_effort", "output_config.effort"} {
		field := gjson.GetBytes(result, path)
		if !field.Exists() || field.Type != gjson.String {
			continue
		}
		original := strings.TrimSpace(field.String())
		if original == "" {
			continue
		}

		effective, mapped := mapReasoningEffort(original, mappings, requestModel)
		if mapped && isReasoningEffortMappingDeny(effective) {
			requested := normalizeReasoningEffortMappingSource(original)
			if requested == "" {
				requested = original
			}
			return body, false, &ReasoningEffortMappingDeniedError{Requested: requested}
		}
		if currentRank, recognized := reasoningEffortRank(effective); recognized {
			effective = NormalizeMaxReasoningEffort(effective)
			if hasMax && currentRank > maxRank {
				if deny {
					return body, false, &ReasoningEffortOverLimitError{Requested: effective, Max: canonicalMax}
				}
				effective = canonicalMax
			}
		}
		if effective == original {
			continue
		}

		updated, err := sjson.SetBytes(result, path, effective)
		if err != nil {
			continue
		}
		result = updated
		changed = true
	}
	return result, changed, nil
}

func applyOpenAIWSReasoningEffortPolicy(payload []byte, hooks *OpenAIWSIngressHooks) ([]byte, error) {
	if hooks == nil || (hooks.MaxReasoningEffort == "" && len(hooks.ReasoningEffortMappings) == 0) {
		return payload, nil
	}
	capped, changed, err := ApplyOpenAIReasoningEffortPolicy(payload, hooks.MaxReasoningEffort, hooks.ReasoningEffortMappings, hooks.MaxReasoningEffortOverLimit)
	if err != nil {
		return payload, err
	}
	if changed {
		return capped, nil
	}
	return payload, nil
}

// ApplyOpenAIReasoningEffortPolicy is retained for OpenAI forwarding callers.
func ApplyOpenAIReasoningEffortPolicy(body []byte, maxEffort string, mappings []ReasoningEffortMapping, overLimit string) ([]byte, bool, error) {
	return ApplyReasoningEffortPolicy(body, maxEffort, mappings, overLimit)
}