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 (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 {
	policy := candidate.PlatformRateLimitPolicy
	if hasRules(candidate.RuntimeRateLimitPolicy) {
		policy = mergeMap(policy, candidate.RuntimeRateLimitPolicy)
	}
	if nested, ok := candidate.RuntimePolicyOverride["rateLimitPolicy"].(map[string]any); ok && len(nested) > 0 {
		policy = mergeMap(policy, nested)
	}
	if hasRules(candidate.ModelRateLimitPolicy) {
		policy = mergeMap(policy, candidate.ModelRateLimitPolicy)
	}
	if hasRules(policy) {
		return policy
	}
	return nil
}

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, 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
}

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
}