package slowquery

import (
	"context"
	"errors"
	"sync"

	"golang.org/x/time/rate"

	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/slowquery/slowstats"
)

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

// StatsMetric returns time-bucketed slow query metrics.
// It forwards the call to the inner SlowQuery implementation with rate limiting.
func (s *rateLimitedSlowQuery) StatsMetric(
	ctx context.Context,
	cluster slowstats.ClusterInfo,
	beginHourTs int64,
	hours int,
	name string,
) ([][]int64, error) {
	if !cluster.SharedTier {
		if err := s.checkRateLimit(cluster.ClusterID); err != nil {
			return nil, err
		}
	}
	return s.inner.StatsMetric(ctx, cluster, beginHourTs, hours, name)
}

// Stats returns aggregated statistics about slow queries.
// It forwards the call to the inner SlowQuery implementation with rate limiting.
func (s *rateLimitedSlowQuery) Stats(
	ctx context.Context,
	cluster slowstats.ClusterInfo,
	beginHourTs int64,
	hours int,
	statementTypes []string,
	statementTypesNotIn bool,
	databases []string,
	internal bool,
	order string,
	limit int,
) ([]slowstats.SlowStats, error) {
	if !cluster.SharedTier {
		if err := s.checkRateLimit(cluster.ClusterID); err != nil {
			return nil, err
		}
	}
	return s.inner.Stats(ctx, cluster, beginHourTs, hours, statementTypes, statementTypesNotIn, databases, internal, order, limit)
}

// Databases returns a list of distinct database names from slow logs.
// It forwards the call to the inner SlowQuery implementation with rate limiting.
func (s *rateLimitedSlowQuery) Databases(ctx context.Context, cluster slowstats.ClusterInfo, beginHourTs int64, hours int, internal bool) ([]string, error) {
	if !cluster.SharedTier {
		if err := s.checkRateLimit(cluster.ClusterID); err != nil {
			return nil, err
		}
	}
	return s.inner.Databases(ctx, cluster, beginHourTs, hours, internal)
}

func (s *rateLimitedSlowQuery) AvailableFields(ctx context.Context, params *AvailableFieldsParams) ([]string, error) {
	return s.inner.AvailableFields(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
}