package atlassian

import (
	"context"
	"errors"
	"math/rand/v2"
	"net/http"
	"sync"
	"time"

	"golang.org/x/time/rate"

	"github.com/tidbcloud/aisupport/packages/sdk-go/infra/logger"
	"github.com/tidbcloud/aisupport/services/jsm-syncer/config"
)

// RateLimitError preserves Atlassian's rate-limit metadata after the client
// has exhausted or declined retries.
type RateLimitError struct {
	Err        error
	RetryAfter time.Duration
	Reason     string
}

func (e *RateLimitError) Error() string {
	return e.Err.Error()
}

func (e *RateLimitError) Unwrap() error {
	return e.Err
}

type requestRateLimitPolicy struct {
	limiter          *rate.Limiter
	cooldown         sharedCooldown
	cooldownDuration time.Duration
	maxRetries       int
	retryDelay       func(int, time.Duration) time.Duration
}

func newDefaultRequestRateLimitPolicy(
	conf config.AtlassianRequestPolicyConfig,
) *requestRateLimitPolicy {
	return newRequestRateLimitPolicy(
		conf,
		defaultRequestsPerSecond,
		defaultRequestBurst,
		defaultRequestMaxRetries,
		defaultRateLimitCooldown,
	)
}

func newRequestRateLimitPolicy(
	conf config.AtlassianRequestPolicyConfig,
	defaultRequestsPerSecond float64,
	defaultBurst int,
	defaultMaxRetries int,
	defaultCooldown time.Duration,
) *requestRateLimitPolicy {
	requestsPerSecond := positiveFloatOrDefault(conf.RequestsPerSecond, defaultRequestsPerSecond)
	burst := positiveIntOrDefault(conf.Burst, defaultBurst)
	maxRetries := positiveIntOrDefault(conf.MaxRetries, defaultMaxRetries)
	return &requestRateLimitPolicy{
		limiter:          rate.NewLimiter(rate.Limit(requestsPerSecond), burst),
		cooldownDuration: positiveDurationOrDefault(conf.Cooldown, defaultCooldown),
		maxRetries:       maxRetries,
		retryDelay:       rateLimitRetryDelay,
	}
}

func (c *client) doWithRateLimitPolicy(
	ctx context.Context,
	policy *requestRateLimitPolicy,
	action string,
	do func() (*http.Response, error),
) (*http.Response, error) {
	for attempt := 0; ; attempt++ {
		if err := c.waitForRateLimitPolicy(ctx, policy); err != nil {
			return nil, err
		}

		resp, err := do()
		if err != nil {
			return nil, err
		}
		if resp.StatusCode != http.StatusTooManyRequests {
			return resp, nil
		}

		err = c.statusError(action, resp)
		_ = resp.Body.Close()
		var limitErr *RateLimitError
		if !errors.As(err, &limitErr) {
			return nil, err
		}
		if isGlobalRateLimit(limitErr.Reason) || attempt >= policy.maxRetries {
			return nil, err
		}

		delay := max(
			policy.cooldownDuration,
			policy.retryDelay(attempt, limitErr.RetryAfter),
		)
		retryAt := policy.cooldown.Extend(delay)
		logger.FromContext(ctx).Info(
			"Atlassian request was rate limited; cooling down endpoint before retry",
			"action", action,
			"attempt", attempt+1,
			"max_retries", policy.maxRetries,
			"rate_limit_reason", limitErr.Reason,
			"retry_after", limitErr.RetryAfter.String(),
			"retry_at", retryAt,
		)
	}
}

func (c *client) waitForRateLimitPolicy(ctx context.Context, policy *requestRateLimitPolicy) error {
	for {
		if err := c.globalRateLimitCooldown.Wait(ctx); err != nil {
			return err
		}
		if err := policy.cooldown.Wait(ctx); err != nil {
			return err
		}
		if err := policy.limiter.Wait(ctx); err != nil {
			return err
		}
		// A different request may have extended a cooldown while this request
		// was waiting for a rate-limit token. Re-enter the limiter after the
		// cooldown so all waiters do not wake up as a burst.
		if !c.globalRateLimitCooldown.Active() && !policy.cooldown.Active() {
			return nil
		}
	}
}

func (c *client) statusError(action string, resp *http.Response) error {
	err := statusError(action, resp)
	var limitErr *RateLimitError
	if errors.As(err, &limitErr) && isGlobalRateLimit(limitErr.Reason) {
		c.globalRateLimitCooldown.Extend(max(
			c.globalRateLimitCooldownDuration,
			limitErr.RetryAfter,
		))
	}
	return err
}

func isGlobalRateLimit(reason string) bool {
	return reason == rateLimitReasonGlobalQuota || reason == rateLimitReasonTenantQuota
}

func rateLimitRetryDelay(attempt int, retryAfter time.Duration) time.Duration {
	backoff := rateLimitBackoffBase * time.Duration(1<<min(attempt, 5))
	if backoff > rateLimitBackoffMax {
		backoff = rateLimitBackoffMax
	}
	jittered := time.Duration(float64(backoff) * (0.7 + rand.Float64()*0.6))
	if jittered > rateLimitBackoffMax {
		jittered = rateLimitBackoffMax
	}
	if retryAfter > jittered {
		return retryAfter
	}
	return jittered
}

type sharedCooldown struct {
	mu    sync.Mutex
	until time.Time
}

func (c *sharedCooldown) Extend(delay time.Duration) time.Time {
	retryAt := time.Now().Add(delay)
	c.mu.Lock()
	defer c.mu.Unlock()
	if retryAt.After(c.until) {
		c.until = retryAt
	}
	return c.until
}

func (c *sharedCooldown) Wait(ctx context.Context) error {
	for {
		c.mu.Lock()
		delay := time.Until(c.until)
		c.mu.Unlock()
		if delay <= 0 {
			return nil
		}

		timer := time.NewTimer(delay)
		select {
		case <-ctx.Done():
			timer.Stop()
			return ctx.Err()
		case <-timer.C:
		}
	}
}

func (c *sharedCooldown) Active() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return time.Now().Before(c.until)
}

func positiveIntOrDefault(value int, fallback int) int {
	if value > 0 {
		return value
	}
	return fallback
}

func positiveFloatOrDefault(value float64, fallback float64) float64 {
	if value > 0 {
		return value
	}
	return fallback
}

func positiveDurationOrDefault(value time.Duration, fallback time.Duration) time.Duration {
	if value > 0 {
		return value
	}
	return fallback
}