package topsqlv2

import (
	"context"
	"errors"
	"sync"

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

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

var _ duckDBAdapter = &rateLimitedQuery{}

type rateLimitedQuery struct {
	inner           duckDBAdapter
	limiters        sync.Map
	perClusterLimit rate.Limit
	perClusterB     int
}

func newRateLimitedQuery(inner duckDBAdapter, perClusterQPS float64, perClusterB int) *rateLimitedQuery {
	return &rateLimitedQuery{
		inner:           inner,
		limiters:        sync.Map{},
		perClusterLimit: rate.Limit(perClusterQPS),
		perClusterB:     perClusterB,
	}
}

func (s *rateLimitedQuery) 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
}

func (s *rateLimitedQuery) SelectSlowlogInDLWithoutFilterPushdown(ctx context.Context, dest interface{}, tableAlias, table, query string) error {
	// Extract clusterID from tableAlias if possible, otherwise skip rate limiting
	// For now, we'll skip rate limiting for this method as it's used internally
	return s.inner.SelectSlowlogInDLWithoutFilterPushdown(ctx, dest, tableAlias, table, query)
}