package slowquery

import (
	"context"
	"errors"
	"sync"

	"golang.org/x/time/rate"
)

var ErrRateLimited = errors.New("rate limited")

var _ SlowQuery = &rateLimitedSlowQuery{}

type rateLimitedSlowQuery struct {
	inner           SlowQuery
	limiters        sync.Map
	perClusterLimit rate.Limit
	perClusterB     int
}

func newRateLimitedSlowQuery(inner SlowQuery, perClusterQPS float64, perClusterB int) *rateLimitedSlowQuery {
	return &rateLimitedSlowQuery{
		inner:           inner,
		limiters:        sync.Map{},
		perClusterLimit: rate.Limit(perClusterQPS),
		perClusterB:     perClusterB,
	}
}

func (s *rateLimitedSlowQuery) List(ctx context.Context, params *ListParams) ([]*SlowLog, error) {
	if params.SharedTier {
		return s.inner.List(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return nil, err
	}
	return s.inner.List(ctx, params)
}

func (s *rateLimitedSlowQuery) Detail(ctx context.Context, params *DetailParams) (*SlowLog, error) {
	if params.SharedTier {
		return s.inner.Detail(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return nil, err
	}
	return s.inner.Detail(ctx, params)
}

func (s *rateLimitedSlowQuery) Analyze(ctx context.Context, params *AnalyzeParams) error {
	if params.SharedTier {
		return s.inner.Analyze(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return err
	}
	return s.inner.Analyze(ctx, params)
}

func (s *rateLimitedSlowQuery) AvailableFields(ctx context.Context) ([]string, error) {
	return s.inner.AvailableFields(ctx)
}

func (s *rateLimitedSlowQuery) Download(ctx context.Context, params *DownloadParams) (*DownloadData, error) {
	return s.inner.Download(ctx, params)
}

func (s *rateLimitedSlowQuery) checkRateLimit(clusterID string) error {
	limiterAny, ok := s.limiters.Load(clusterID)
	if !ok {
		limiterAny = rate.NewLimiter(s.perClusterLimit, s.perClusterB)
		s.limiters.Store(clusterID, limiterAny)
	}
	limiter := limiterAny.(*rate.Limiter)
	if !limiter.Allow() {
		return ErrRateLimited
	}
	return nil
}