package sqls

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) ListTables(ctx context.Context, params *ListTableParams) ([]string, error) {
	if params.SharedTier {
		return s.inner.ListTables(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return nil, err
	}
	return s.inner.ListTables(ctx, params)
}

// getSQLsBaseURL delegates the call to the inner adapter
func (s *rateLimitedQuery) getSQLsBaseURL(tenantID, o11yClusterID, env, provider, region string, sharedTier, nextGen bool) string {
	return s.inner.getSQLsBaseURL(tenantID, o11yClusterID, env, provider, region, sharedTier, nextGen)
}

func (s *rateLimitedQuery) ListTableFields(ctx context.Context, params *ListTableFieldsParams) ([]FieldInfo, error) {
	if params.SharedTier {
		return s.inner.ListTableFields(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return nil, err
	}
	return s.inner.ListTableFields(ctx, params)
}

func (s *rateLimitedQuery) ListTableData(ctx context.Context, params *ListTableDataParams) (*ListTableDataResponse, error) {
	if params.SharedTier {
		return s.inner.ListTableData(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return nil, err
	}
	return s.inner.ListTableData(ctx, params)
}

func (s *rateLimitedQuery) GetTableDataByID(ctx context.Context, params *GetTableDataByIDParams) (map[string]interface{}, error) {
	if params.SharedTier {
		return s.inner.GetTableDataByID(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return nil, err
	}
	return s.inner.GetTableDataByID(ctx, params)
}

func (s *rateLimitedQuery) QueryStatementsSummary(ctx context.Context, params *StatementsSummaryParams) ([]StatementsSummaryItem, error) {
	if params.SharedTier {
		return s.inner.QueryStatementsSummary(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return nil, err
	}
	return s.inner.QueryStatementsSummary(ctx, params)
}

func (s *rateLimitedQuery) QueryStatementsTopStats(ctx context.Context, params *StatementsTopStatsParams) ([]StatementsTopStatsItem, error) {
	if params.SharedTier {
		return s.inner.QueryStatementsTopStats(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return nil, err
	}
	return s.inner.QueryStatementsTopStats(ctx, params)
}

func (s *rateLimitedQuery) QueryStatementsFieldValues(ctx context.Context, params *StatementsFieldValuesParams) ([]string, error) {
	if params.SharedTier {
		return s.inner.QueryStatementsFieldValues(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return nil, err
	}
	return s.inner.QueryStatementsFieldValues(ctx, params)
}

func (s *rateLimitedQuery) ProbeStatementsDataSources(ctx context.Context, params *StatementsDataSourceProbeParams) (*StatementsDataSourcesResponse, error) {
	if params.SharedTier {
		return s.inner.ProbeStatementsDataSources(ctx, params)
	}
	if err := s.checkRateLimit(params.ClusterID); err != nil {
		return nil, err
	}
	return s.inner.ProbeStatementsDataSources(ctx, params)
}