package runner

import (
	"context"
	"strings"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func (s *Service) rateLimitReservations(ctx context.Context, user *auth.User, candidate store.RuntimeModelCandidate, body map[string]any) []store.RateLimitReservation {
	out := make([]store.RateLimitReservation, 0)
	out = append(out, reservationsFromPolicy("platform_model", candidate.PlatformModelID, effectiveRateLimitPolicy(candidate), body)...)
	if group, err := s.store.ResolveUserGroupPolicy(ctx, user); err == nil && group.ID != "" {
		out = append(out, reservationsFromPolicy("user_group", group.ID, group.RateLimitPolicy, body)...)
	}
	return out
}

func effectiveRateLimitPolicy(candidate store.RuntimeModelCandidate) map[string]any {
	if hasRules(candidate.ModelRateLimitPolicy) {
		return candidate.ModelRateLimitPolicy
	}
	if hasRules(candidate.PlatformRateLimitPolicy) {
		return candidate.PlatformRateLimitPolicy
	}
	return nil
}

func reservationsFromPolicy(scopeType string, scopeKey string, policy map[string]any, body map[string]any) []store.RateLimitReservation {
	if scopeKey == "" || !hasRules(policy) {
		return nil
	}
	rules, _ := policy["rules"].([]any)
	out := make([]store.RateLimitReservation, 0, len(rules))
	estimatedTokens := estimateRequestTokens(body)
	for _, rawRule := range rules {
		rule, _ := rawRule.(map[string]any)
		metric := strings.TrimSpace(stringFromMap(rule, "metric"))
		limit := floatFromAny(rule["limit"])
		amount := 1.0
		if strings.HasPrefix(metric, "tpm") {
			amount = float64(estimatedTokens)
		}
		out = append(out, store.RateLimitReservation{
			ScopeType:       scopeType,
			ScopeKey:        scopeKey,
			Metric:          metric,
			Limit:           limit,
			Amount:          amount,
			WindowSeconds:   int(floatFromAny(rule["windowSeconds"])),
			LeaseTTLSeconds: int(floatFromAny(rule["leaseTtlSeconds"])),
		})
	}
	return out
}

func hasRules(policy map[string]any) bool {
	rules, _ := policy["rules"].([]any)
	return len(rules) > 0
}

func estimateRequestTokens(body map[string]any) int {
	text := ""
	if prompt := stringFromMap(body, "prompt"); prompt != "" {
		text += prompt
	}
	if input := stringFromMap(body, "input"); input != "" {
		text += input
	}
	if messages, ok := body["messages"].([]any); ok {
		for _, raw := range messages {
			message, _ := raw.(map[string]any)
			switch content := message["content"].(type) {
			case string:
				text += content
			case []any:
				for _, rawPart := range content {
					part, _ := rawPart.(map[string]any)
					text += stringFromMap(part, "text")
				}
			}
		}
	}
	if text == "" {
		return 1
	}
	return len([]rune(text))/4 + 1
}