package slowquery

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

	"github.com/stretchr/testify/require"

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

// mockSlowQuery is a mock implementation of SlowQuery for testing
type mockSlowQuery struct {
	listFunc            func(ctx context.Context, params *ListParams) ([]*SlowLog, error)
	detailFunc          func(ctx context.Context, params *DetailParams) (*SlowLog, error)
	statsMetricFunc     func(ctx context.Context, cluster slowstats.ClusterInfo, beginHourTs int64, hours int, name string) ([][]int64, error)
	statsFunc           func(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)
	databasesFunc       func(ctx context.Context, cluster slowstats.ClusterInfo, beginHourTs int64, hours int, internal bool) ([]string, error)
	availableFieldsFunc func(ctx context.Context, params *AvailableFieldsParams) ([]string, error)
}

func (m *mockSlowQuery) List(ctx context.Context, params *ListParams) ([]*SlowLog, error) {
	if m.listFunc != nil {
		return m.listFunc(ctx, params)
	}
	return []*SlowLog{{Digest: "test-digest"}}, nil
}

func (m *mockSlowQuery) Detail(ctx context.Context, params *DetailParams) (*SlowLog, error) {
	if m.detailFunc != nil {
		return m.detailFunc(ctx, params)
	}
	return &SlowLog{Digest: "test-digest"}, nil
}

func (m *mockSlowQuery) StatsMetric(ctx context.Context, cluster slowstats.ClusterInfo, beginHourTs int64, hours int, name string) ([][]int64, error) {
	if m.statsMetricFunc != nil {
		return m.statsMetricFunc(ctx, cluster, beginHourTs, hours, name)
	}
	return [][]int64{{1, 2, 3}}, nil
}

func (m *mockSlowQuery) 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 m.statsFunc != nil {
		return m.statsFunc(ctx, cluster, beginHourTs, hours, statementTypes, statementTypesNotIn, databases, internal, order, limit)
	}
	return []slowstats.SlowStats{}, nil
}

func (m *mockSlowQuery) Databases(ctx context.Context, cluster slowstats.ClusterInfo, beginHourTs int64, hours int, internal bool) ([]string, error) {
	if m.databasesFunc != nil {
		return m.databasesFunc(ctx, cluster, beginHourTs, hours, internal)
	}
	return []string{"db1", "db2"}, nil
}

func (m *mockSlowQuery) AvailableFields(ctx context.Context, params *AvailableFieldsParams) ([]string, error) {
	if m.availableFieldsFunc != nil {
		return m.availableFieldsFunc(ctx, params)
	}
	return []string{"field1", "field2"}, nil
}

func TestNewRateLimitedSlowQuery(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 10.0, 10)

	require.NotNil(t, rlq)
	require.Equal(t, mockQuery, rlq.inner)
	require.Equal(t, 10.0, float64(rlq.perClusterLimit))
	require.Equal(t, 10, rlq.perClusterB)
}

func TestRateLimitedSlowQuery_List_SharedTier(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 10.0, 10)

	params := &ListParams{
		ClusterID:  "cluster-1",
		SharedTier: true,
	}

	logs, err := rlq.List(context.Background(), params)
	require.NoError(t, err)
	require.NotNil(t, logs)
	require.Len(t, logs, 1)
}

func TestRateLimitedSlowQuery_List_RateLimited(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	// Set a very low rate limit to trigger rate limiting
	rlq := newRateLimitedSlowQuery(mockQuery, 0.1, 1)

	params := &ListParams{
		ClusterID:  "cluster-1",
		SharedTier: false,
	}

	// First request should succeed
	logs, err := rlq.List(context.Background(), params)
	require.NoError(t, err)
	require.NotNil(t, logs)

	// Second request immediately after should be rate limited
	_, err = rlq.List(context.Background(), params)
	require.Error(t, err)
	require.Equal(t, ErrRateLimited, err)
}

func TestRateLimitedSlowQuery_Detail_SharedTier(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 10.0, 10)

	params := &DetailParams{
		ClusterID:  "cluster-1",
		SharedTier: true,
		ConnID:     123,
		Digest:     "test-digest",
		Timestamp:  1234567890.0,
	}

	log, err := rlq.Detail(context.Background(), params)
	require.NoError(t, err)
	require.NotNil(t, log)
}

func TestRateLimitedSlowQuery_Detail_RateLimited(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 0.1, 1)

	params := &DetailParams{
		ClusterID:  "cluster-1",
		SharedTier: false,
		ConnID:     123,
		Digest:     "test-digest",
		Timestamp:  1234567890.0,
	}

	// First request should succeed
	_, err := rlq.Detail(context.Background(), params)
	require.NoError(t, err)

	// Second request should be rate limited
	_, err = rlq.Detail(context.Background(), params)
	require.Error(t, err)
	require.Equal(t, ErrRateLimited, err)
}

func TestRateLimitedSlowQuery_StatsMetric_SharedTier(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 10.0, 10)

	cluster := slowstats.ClusterInfo{
		ClusterID:  "cluster-1",
		SharedTier: true,
	}

	result, err := rlq.StatsMetric(context.Background(), cluster, 1234567890, 24, "count_per_minute")
	require.NoError(t, err)
	require.NotNil(t, result)
}

func TestRateLimitedSlowQuery_StatsMetric_RateLimited(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 0.1, 1)

	cluster := slowstats.ClusterInfo{
		ClusterID:  "cluster-1",
		SharedTier: false,
	}

	// First request should succeed
	_, err := rlq.StatsMetric(context.Background(), cluster, 1234567890, 24, "count_per_minute")
	require.NoError(t, err)

	// Second request should be rate limited
	_, err = rlq.StatsMetric(context.Background(), cluster, 1234567890, 24, "count_per_minute")
	require.Error(t, err)
	require.Equal(t, ErrRateLimited, err)
}

func TestRateLimitedSlowQuery_Stats_SharedTier(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 10.0, 10)

	cluster := slowstats.ClusterInfo{
		ClusterID:  "cluster-1",
		SharedTier: true,
	}

	result, err := rlq.Stats(context.Background(), cluster, 1234567890, 24, []string{"Select"}, false, []string{"db1"}, false, "count", 10)
	require.NoError(t, err)
	require.NotNil(t, result)
}

func TestRateLimitedSlowQuery_Stats_RateLimited(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 0.1, 1)

	cluster := slowstats.ClusterInfo{
		ClusterID:  "cluster-1",
		SharedTier: false,
	}

	// First request should succeed
	_, err := rlq.Stats(context.Background(), cluster, 1234567890, 24, []string{"Select"}, false, []string{"db1"}, false, "count", 10)
	require.NoError(t, err)

	// Second request should be rate limited
	_, err = rlq.Stats(context.Background(), cluster, 1234567890, 24, []string{"Select"}, false, []string{"db1"}, false, "count", 10)
	require.Error(t, err)
	require.Equal(t, ErrRateLimited, err)
}

func TestRateLimitedSlowQuery_Databases_SharedTier(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 10.0, 10)

	cluster := slowstats.ClusterInfo{
		ClusterID:  "cluster-1",
		SharedTier: true,
	}

	dbs, err := rlq.Databases(context.Background(), cluster, 1234567890, 24, false)
	require.NoError(t, err)
	require.NotNil(t, dbs)
	require.Len(t, dbs, 2)
}

func TestRateLimitedSlowQuery_Databases_RateLimited(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 0.1, 1)

	cluster := slowstats.ClusterInfo{
		ClusterID:  "cluster-1",
		SharedTier: false,
	}

	// First request should succeed
	_, err := rlq.Databases(context.Background(), cluster, 1234567890, 24, false)
	require.NoError(t, err)

	// Second request should be rate limited
	_, err = rlq.Databases(context.Background(), cluster, 1234567890, 24, false)
	require.Error(t, err)
	require.Equal(t, ErrRateLimited, err)
}

func TestRateLimitedSlowQuery_AvailableFields(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 10.0, 10)

	// AvailableFields should not be rate limited
	fields, err := rlq.AvailableFields(context.Background(), &AvailableFieldsParams{ClusterID: "c1"})
	require.NoError(t, err)
	require.NotNil(t, fields)
	require.Len(t, fields, 2)
}

func TestRateLimitedSlowQuery_MultipleClusters(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 10.0, 10)

	// Different clusters should have separate rate limiters
	params1 := &ListParams{
		ClusterID:  "cluster-1",
		SharedTier: false,
	}
	params2 := &ListParams{
		ClusterID:  "cluster-2",
		SharedTier: false,
	}

	// Both should succeed
	logs1, err1 := rlq.List(context.Background(), params1)
	require.NoError(t, err1)
	require.NotNil(t, logs1)

	logs2, err2 := rlq.List(context.Background(), params2)
	require.NoError(t, err2)
	require.NotNil(t, logs2)
}

func TestRateLimitedSlowQuery_ErrorPropagation(t *testing.T) {
	testErr := errors.New("test error")
	mockQuery := &mockSlowQuery{
		listFunc: func(ctx context.Context, params *ListParams) ([]*SlowLog, error) {
			return nil, testErr
		},
	}
	rlq := newRateLimitedSlowQuery(mockQuery, 10.0, 10)

	params := &ListParams{
		ClusterID:  "cluster-1",
		SharedTier: false,
	}

	_, err := rlq.List(context.Background(), params)
	require.Error(t, err)
	require.Equal(t, testErr, err)
}

func TestRateLimitedSlowQuery_RateLimitRecovery(t *testing.T) {
	mockQuery := &mockSlowQuery{}
	rlq := newRateLimitedSlowQuery(mockQuery, 1.0, 1) // 1 QPS, bucket size 1

	params := &ListParams{
		ClusterID:  "cluster-1",
		SharedTier: false,
	}

	// First request should succeed
	_, err := rlq.List(context.Background(), params)
	require.NoError(t, err)

	// Second request immediately after should be rate limited
	_, err = rlq.List(context.Background(), params)
	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.List(context.Background(), params)
	require.NoError(t, err)
}