package topsqlv2

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"golang.org/x/time/rate"
)

// mockDuckDBAdapterForRateLimit is a mock implementation of duckDBAdapter for rate limit testing
type mockDuckDBAdapterForRateLimit struct {
	selectFunc func(ctx context.Context, dest interface{}, tableAlias, table, query string) error
}

func (m *mockDuckDBAdapterForRateLimit) SelectSlowlogInDLWithoutFilterPushdown(ctx context.Context, dest interface{}, tableAlias, table, query string) error {
	if m.selectFunc != nil {
		return m.selectFunc(ctx, dest, tableAlias, table, query)
	}
	return nil
}

func TestNewRateLimitedQuery(t *testing.T) {
	mockAdapter := &mockDuckDBAdapterForRateLimit{}
	rlq := newRateLimitedQuery(mockAdapter, 10.0, 10)

	require.NotNil(t, rlq)
	require.Equal(t, mockAdapter, rlq.inner)
	require.Equal(t, rate.Limit(10.0), rlq.perClusterLimit)
	require.Equal(t, 10, rlq.perClusterB)
}

func TestRateLimitedQuery_SelectSlowlogInDLWithoutFilterPushdown(t *testing.T) {
	mockAdapter := &mockDuckDBAdapterForRateLimit{
		selectFunc: func(ctx context.Context, dest interface{}, tableAlias, table, query string) error {
			return nil
		},
	}
	rlq := newRateLimitedQuery(mockAdapter, 10.0, 10)

	var dest []interface{}
	err := rlq.SelectSlowlogInDLWithoutFilterPushdown(context.Background(), &dest, "alias", "table", "SELECT * FROM table")
	require.NoError(t, err)
}

func TestRateLimitedQuery_SelectSlowlogInDLWithoutFilterPushdown_Error(t *testing.T) {
	testErr := errors.New("test error")
	mockAdapter := &mockDuckDBAdapterForRateLimit{
		selectFunc: func(ctx context.Context, dest interface{}, tableAlias, table, query string) error {
			return testErr
		},
	}
	rlq := newRateLimitedQuery(mockAdapter, 10.0, 10)

	var dest []interface{}
	err := rlq.SelectSlowlogInDLWithoutFilterPushdown(context.Background(), &dest, "alias", "table", "SELECT * FROM table")
	require.Error(t, err)
	require.Equal(t, testErr, err)
}

func TestCheckRateLimit(t *testing.T) {
	mockAdapter := &mockDuckDBAdapterForRateLimit{}
	rlq := newRateLimitedQuery(mockAdapter, 100.0, 100)

	// First request should succeed
	err := rlq.checkRateLimit("cluster-1")
	require.NoError(t, err)

	// Multiple requests should succeed with high limit
	// Use a reasonable number that won't trigger rate limiting
	for i := 0; i < 50; i++ {
		err := rlq.checkRateLimit("cluster-1")
		require.NoError(t, err)
	}
}

func TestCheckRateLimit_RateLimited(t *testing.T) {
	mockAdapter := &mockDuckDBAdapterForRateLimit{}
	// Set a very low rate limit to trigger rate limiting
	rlq := newRateLimitedQuery(mockAdapter, 0.1, 1)

	// First request should succeed
	err := rlq.checkRateLimit("cluster-1")
	require.NoError(t, err)

	// Second request immediately after should be rate limited
	err = rlq.checkRateLimit("cluster-1")
	require.Error(t, err)
	require.Equal(t, ErrRateLimited, err)
}

func TestCheckRateLimit_MultipleClusters(t *testing.T) {
	mockAdapter := &mockDuckDBAdapterForRateLimit{}
	rlq := newRateLimitedQuery(mockAdapter, 0.1, 1)

	// Different clusters should have separate rate limiters
	err1 := rlq.checkRateLimit("cluster-1")
	require.NoError(t, err1)

	err2 := rlq.checkRateLimit("cluster-2")
	require.NoError(t, err2)

	// Both clusters should be able to make requests
	err3 := rlq.checkRateLimit("cluster-1")
	require.Error(t, err3) // Should be rate limited for cluster-1

	err4 := rlq.checkRateLimit("cluster-2")
	require.Error(t, err4) // Should be rate limited for cluster-2
}

func TestCheckRateLimit_RateLimitRecovery(t *testing.T) {
	mockAdapter := &mockDuckDBAdapterForRateLimit{}
	rlq := newRateLimitedQuery(mockAdapter, 1.0, 1) // 1 QPS, bucket size 1

	// First request should succeed
	err := rlq.checkRateLimit("cluster-1")
	require.NoError(t, err)

	// Second request immediately after should be rate limited
	err = rlq.checkRateLimit("cluster-1")
	require.Error(t, err)
	require.Equal(t, ErrRateLimited, err)

	// Wait for rate limit to recover (1 second for 1 QPS)
	time.Sleep(1100 * time.Millisecond)

	// Third request should succeed after waiting
	err = rlq.checkRateLimit("cluster-1")
	require.NoError(t, err)
}

func TestCheckRateLimit_Concurrent(t *testing.T) {
	mockAdapter := &mockDuckDBAdapterForRateLimit{}
	rlq := newRateLimitedQuery(mockAdapter, 100.0, 100)

	// Test concurrent requests
	done := make(chan error, 10)
	for i := 0; i < 10; i++ {
		go func() {
			err := rlq.checkRateLimit("cluster-1")
			done <- err
		}()
	}

	// Wait for all requests to complete
	successCount := 0
	for i := 0; i < 10; i++ {
		err := <-done
		if err == nil {
			successCount++
		}
	}
	// With high limit, most should succeed
	require.Greater(t, successCount, 0)
}