package sqls

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

	"github.com/stretchr/testify/require"
)

// mockDuckDBAdapter is a mock implementation of duckDBAdapter for testing
type mockDuckDBAdapter struct {
	listTablesFunc                 func(ctx context.Context, params *ListTableParams) ([]string, error)
	listTableFieldsFunc            func(ctx context.Context, params *ListTableFieldsParams) ([]FieldInfo, error)
	listTableDataFunc              func(ctx context.Context, params *ListTableDataParams) (*ListTableDataResponse, error)
	getTableDataByIDFunc           func(ctx context.Context, params *GetTableDataByIDParams) (map[string]interface{}, error)
	queryStatementsSummaryFunc     func(ctx context.Context, params *StatementsSummaryParams) ([]StatementsSummaryItem, error)
	queryStatementsTopStatsFunc    func(ctx context.Context, params *StatementsTopStatsParams) ([]StatementsTopStatsItem, error)
	queryStatementsFieldValuesFunc func(ctx context.Context, params *StatementsFieldValuesParams) ([]string, error)
	probeStatementsDataSourcesFunc func(ctx context.Context, params *StatementsDataSourceProbeParams) (*StatementsDataSourcesResponse, error)
	getSQLsBaseURLFunc             func(tenantID, o11yClusterID, env, provider, region string, sharedTier, nextGen bool) string
}

func (m *mockDuckDBAdapter) ListTables(ctx context.Context, params *ListTableParams) ([]string, error) {
	if m.listTablesFunc != nil {
		return m.listTablesFunc(ctx, params)
	}
	return []string{"table1", "table2"}, nil
}

func (m *mockDuckDBAdapter) ListTableFields(ctx context.Context, params *ListTableFieldsParams) ([]FieldInfo, error) {
	if m.listTableFieldsFunc != nil {
		return m.listTableFieldsFunc(ctx, params)
	}
	return []FieldInfo{{Name: "field1", Type: "string"}}, nil
}

func (m *mockDuckDBAdapter) ListTableData(ctx context.Context, params *ListTableDataParams) (*ListTableDataResponse, error) {
	if m.listTableDataFunc != nil {
		return m.listTableDataFunc(ctx, params)
	}
	return &ListTableDataResponse{}, nil
}

func (m *mockDuckDBAdapter) GetTableDataByID(ctx context.Context, params *GetTableDataByIDParams) (map[string]interface{}, error) {
	if m.getTableDataByIDFunc != nil {
		return m.getTableDataByIDFunc(ctx, params)
	}
	return map[string]interface{}{"id": params.ID}, nil
}

func (m *mockDuckDBAdapter) getSQLsBaseURL(tenantID, o11yClusterID, env, provider, region string, sharedTier, nextGen bool) string {
	if m.getSQLsBaseURLFunc != nil {
		return m.getSQLsBaseURLFunc(tenantID, o11yClusterID, env, provider, region, sharedTier, nextGen)
	}
	return "s3://bucket/path"
}

func (m *mockDuckDBAdapter) QueryStatementsSummary(ctx context.Context, params *StatementsSummaryParams) ([]StatementsSummaryItem, error) {
	if m.queryStatementsSummaryFunc != nil {
		return m.queryStatementsSummaryFunc(ctx, params)
	}
	return []StatementsSummaryItem{}, nil
}

func (m *mockDuckDBAdapter) QueryStatementsTopStats(ctx context.Context, params *StatementsTopStatsParams) ([]StatementsTopStatsItem, error) {
	if m.queryStatementsTopStatsFunc != nil {
		return m.queryStatementsTopStatsFunc(ctx, params)
	}
	return []StatementsTopStatsItem{}, nil
}

func (m *mockDuckDBAdapter) QueryStatementsFieldValues(ctx context.Context, params *StatementsFieldValuesParams) ([]string, error) {
	if m.queryStatementsFieldValuesFunc != nil {
		return m.queryStatementsFieldValuesFunc(ctx, params)
	}
	return []string{}, nil
}

func (m *mockDuckDBAdapter) ProbeStatementsDataSources(ctx context.Context, params *StatementsDataSourceProbeParams) (*StatementsDataSourcesResponse, error) {
	if m.probeStatementsDataSourcesFunc != nil {
		return m.probeStatementsDataSourcesFunc(ctx, params)
	}
	return &StatementsDataSourcesResponse{V2Available: false}, nil
}

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

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

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

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

	tables, err := rlq.ListTables(context.Background(), params)
	require.NoError(t, err)
	require.Equal(t, []string{"table1", "table2"}, tables)
}

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

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

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

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

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

	params := &ListTableFieldsParams{
		ClusterID:  "cluster-1",
		SharedTier: true,
		TableName:  "table1",
	}

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

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

	params := &ListTableFieldsParams{
		ClusterID:  "cluster-1",
		SharedTier: false,
		TableName:  "table1",
	}

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

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

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

	params := &ListTableDataParams{
		ClusterID:  "cluster-1",
		SharedTier: true,
		TableName:  "table1",
	}

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

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

	params := &ListTableDataParams{
		ClusterID:  "cluster-1",
		SharedTier: false,
		TableName:  "table1",
	}

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

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

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

	params := &GetTableDataByIDParams{
		ClusterID:  "cluster-1",
		SharedTier: true,
		TableName:  "table1",
		ID:         123,
	}

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

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

	params := &GetTableDataByIDParams{
		ClusterID:  "cluster-1",
		SharedTier: false,
		TableName:  "table1",
		ID:         123,
	}

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

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

func TestRateLimitedQuery_GetSQLsBaseURL(t *testing.T) {
	mockAdapter := &mockDuckDBAdapter{
		getSQLsBaseURLFunc: func(tenantID, o11yClusterID, env, provider, region string, sharedTier, nextGen bool) string {
			return "s3://test-bucket/path"
		},
	}
	rlq := newRateLimitedQuery(mockAdapter, 10.0, 10)

	url := rlq.getSQLsBaseURL("tenant-1", "cluster-1", "prod", "aws", "us-west-2", false, false)
	require.Equal(t, "s3://test-bucket/path", url)
}

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

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

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

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

func TestRateLimitedQuery_ErrorPropagation(t *testing.T) {
	testErr := errors.New("test error")
	mockAdapter := &mockDuckDBAdapter{
		listTablesFunc: func(ctx context.Context, params *ListTableParams) ([]string, error) {
			return nil, testErr
		},
	}
	rlq := newRateLimitedQuery(mockAdapter, 10.0, 10)

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

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

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

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

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

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