package runner

import (
	"context"
	"errors"
	"strings"
	"time"

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

type localRateLimitError struct {
	clientErr  *clients.ClientError
	cause      error
	retryAfter time.Duration
}

func (e *localRateLimitError) Error() string {
	if e == nil || e.clientErr == nil {
		return store.ErrRateLimited.Error()
	}
	return e.clientErr.Error()
}

func (e *localRateLimitError) Unwrap() []error {
	if e == nil || e.clientErr == nil {
		if e != nil && e.cause != nil {
			return []error{e.cause}
		}
		return []error{store.ErrRateLimited}
	}
	if e.cause != nil {
		return []error{e.clientErr, e.cause}
	}
	return []error{e.clientErr, store.ErrRateLimited}
}

func localRateLimitRetryAfter(err error) time.Duration {
	var limitErr *localRateLimitError
	if errors.As(err, &limitErr) && limitErr.retryAfter > 0 {
		return limitErr.retryAfter
	}
	return store.RateLimitRetryAfter(err)
}

func isLocalRateLimitError(err error) bool {
	var limitErr *localRateLimitError
	return errors.As(err, &limitErr)
}

func platformModelRateLimitError(err error) (*store.RateLimitExceededError, bool) {
	var limitErr *store.RateLimitExceededError
	if !errors.As(err, &limitErr) || limitErr.ScopeType != "platform_model" {
		return nil, false
	}
	return limitErr, true
}

func (s *Service) rateLimitReservations(ctx context.Context, user *auth.User, candidate store.RuntimeModelCandidate, body map[string]any) []store.RateLimitReservation {
	out := make([]store.RateLimitReservation, 0)
	// Cross-platform scopes must be evaluated first so a user-group limit can
	// never be bypassed by rotating to another platform model.
	if group, err := s.store.ResolveUserGroupPolicy(ctx, user); err == nil && group.ID != "" {
		out = append(out, reservationsFromPolicy(
			"user_group",
			group.ID,
			firstNonEmptyString(group.Name, group.GroupKey),
			map[string]any{
				"groupKey": group.GroupKey,
				"name":     group.Name,
			},
			store.NormalizeRateLimitPolicy(group.RateLimitPolicy),
			body,
		)...)
	}
	out = append(out, reservationsFromPolicy(
		"platform_model",
		candidate.PlatformModelID,
		firstNonEmptyString(candidate.DisplayName, candidate.ModelAlias, candidate.ModelName),
		map[string]any{
			"platformId":   candidate.PlatformID,
			"platformName": candidate.PlatformName,
			"modelAlias":   candidate.ModelAlias,
			"modelName":    candidate.ModelName,
		},
		effectiveRateLimitPolicy(candidate),
		body,
	)...)
	return out
}

func (s *Service) admissionScopes(ctx context.Context, user *auth.User, candidate store.RuntimeModelCandidate) ([]store.AdmissionScope, string) {
	scopes := make([]store.AdmissionScope, 0, 2)
	platformPolicy := effectiveRateLimitPolicy(candidate)
	if scope, ok := admissionScopeFromPolicy(
		"platform_model",
		candidate.PlatformModelID,
		firstNonEmptyString(candidate.DisplayName, candidate.ModelAlias, candidate.ModelName),
		map[string]any{
			"platformId":   candidate.PlatformID,
			"platformName": candidate.PlatformName,
			"modelAlias":   candidate.ModelAlias,
			"modelName":    candidate.ModelName,
		},
		platformPolicy,
	); ok {
		scopes = append(scopes, scope)
	}
	groupID := ""
	if group, err := s.store.ResolveUserGroupPolicy(ctx, user); err == nil && group.ID != "" {
		groupPolicy := store.NormalizeRateLimitPolicy(group.RateLimitPolicy)
		if scope, ok := admissionScopeFromPolicy(
			"user_group",
			group.ID,
			firstNonEmptyString(group.Name, group.GroupKey),
			map[string]any{"groupKey": group.GroupKey, "name": group.Name},
			groupPolicy,
		); ok {
			groupID = group.ID
			scopes = append(scopes, scope)
		}
	}
	return scopes, groupID
}

func admissionScopeFromPolicy(scopeType string, scopeKey string, scopeName string, metadata map[string]any, policy map[string]any) (store.AdmissionScope, bool) {
	concurrentLimit, hasConcurrent := store.RateLimitPolicyMetric(policy, "concurrent")
	queueRule, hasQueue := store.QueueRuleFromPolicy(policy)
	if !hasConcurrent && !hasQueue {
		return store.AdmissionScope{}, false
	}
	scope := store.AdmissionScope{
		ScopeType:       scopeType,
		ScopeKey:        scopeKey,
		ScopeName:       scopeName,
		ScopeMetadata:   metadata,
		ConcurrentLimit: concurrentLimit,
		Amount:          1,
		LeaseTTLSeconds: 120,
		Policy:          policy,
	}
	if hasQueue {
		scope.QueueLimit = queueRule.Limit
		scope.MaxWaitSeconds = queueRule.MaxWaitSeconds
	}
	rules, _ := store.NormalizeRateLimitPolicy(policy)["rules"].([]any)
	for _, rawRule := range rules {
		rule, _ := rawRule.(map[string]any)
		if strings.TrimSpace(stringFromMap(rule, "metric")) != "concurrent" {
			continue
		}
		if ttl := int(floatFromAny(rule["leaseTtlSeconds"])); ttl > 0 {
			scope.LeaseTTLSeconds = ttl
		}
		break
	}
	return scope, true
}

func effectiveRateLimitPolicy(candidate store.RuntimeModelCandidate) map[string]any {
	return store.EffectiveRateLimitPolicy(store.EffectiveRateLimitPolicyInput{
		BasePolicy:            candidate.BaseRateLimitPolicy,
		PlatformPolicy:        candidate.PlatformRateLimitPolicy,
		RuntimePolicy:         candidate.RuntimeRateLimitPolicy,
		RuntimePolicyExplicit: candidate.RuntimePolicyExplicit,
		RuntimePolicyOverride: candidate.RateLimitRuntimeOverride,
		ModelPolicy:           candidate.ModelRateLimitPolicy,
		ModelPolicyMode:       candidate.ModelRateLimitPolicyMode,
	})
}

func effectiveRetryPolicy(candidate store.RuntimeModelCandidate) map[string]any {
	policy := candidate.PlatformRetryPolicy
	if len(candidate.RuntimeRetryPolicy) > 0 {
		policy = mergeMap(policy, candidate.RuntimeRetryPolicy)
	}
	if nested, ok := candidate.RuntimePolicyOverride["retryPolicy"].(map[string]any); ok && len(nested) > 0 {
		policy = mergeMap(policy, nested)
	}
	if len(candidate.ModelRetryPolicy) > 0 {
		policy = mergeMap(policy, candidate.ModelRetryPolicy)
	}
	return policy
}

func reservationsFromPolicy(scopeType string, scopeKey string, scopeName string, scopeMetadata map[string]any, 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"])
		if metric == "" || metric == "queue_size" || limit <= 0 {
			continue
		}
		amount := 1.0
		if strings.HasPrefix(metric, "tpm") {
			amount = float64(estimatedTokens)
		}
		out = append(out, store.RateLimitReservation{
			ScopeType:       scopeType,
			ScopeKey:        scopeKey,
			ScopeName:       scopeName,
			ScopeMetadata:   scopeMetadata,
			Metric:          metric,
			Limit:           limit,
			Amount:          amount,
			WindowSeconds:   int(floatFromAny(rule["windowSeconds"])),
			LeaseTTLSeconds: int(floatFromAny(rule["leaseTtlSeconds"])),
			Policy:          policy,
		})
	}
	return out
}

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

func estimateRequestTokens(body map[string]any) int {
	var text strings.Builder
	appendTokenEstimateText(&text, body["prompt"])
	appendTokenEstimateText(&text, body["input"])
	appendTokenEstimateText(&text, body["query"])
	appendTokenEstimateText(&text, body["documents"])
	for _, item := range contentItems(body["content"]) {
		if stringFromAny(item["type"]) == "text" {
			appendTokenEstimateText(&text, item["text"])
		}
	}
	if messages, ok := body["messages"].([]any); ok {
		for _, raw := range messages {
			message, _ := raw.(map[string]any)
			switch content := message["content"].(type) {
			case string:
				appendTokenEstimateText(&text, content)
			case []any:
				for _, rawPart := range content {
					part, _ := rawPart.(map[string]any)
					appendTokenEstimateText(&text, part["text"])
				}
			}
		}
	}
	estimatedText := text.String()
	if estimatedText == "" {
		return 1
	}
	return len([]rune(estimatedText))/4 + 1
}

func appendTokenEstimateText(out *strings.Builder, value any) {
	switch typed := value.(type) {
	case string:
		out.WriteString(typed)
	case []string:
		for _, item := range typed {
			out.WriteString(item)
		}
	case []any:
		for _, item := range typed {
			appendTokenEstimateText(out, item)
		}
	case map[string]any:
		for _, key := range []string{"text", "content", "query", "document"} {
			if text := stringFromAny(typed[key]); text != "" {
				out.WriteString(text)
			}
		}
	}
}

func tokenUsageAmounts(usage clients.Usage) map[string]float64 {
	out := map[string]float64{}
	if usage.InputTokens > 0 {
		out["tpm_input"] = float64(usage.InputTokens)
	}
	if usage.OutputTokens > 0 {
		out["tpm_output"] = float64(usage.OutputTokens)
	}
	total := usage.TotalTokens
	if total <= 0 {
		total = usage.InputTokens + usage.OutputTokens
	}
	if total > 0 {
		out["tpm_total"] = float64(total)
	}
	return out
}
