package conprof

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/storage"
)

type mockStorage struct {
	mock.Mock
}

func (m *mockStorage) Objects(ctx context.Context, bucket, prefix string) ([]storage.ObjectMeta, error) {
	args := m.Called(ctx, bucket, prefix)
	if args.Get(0) == nil {
		return nil, args.Error(1)
	}
	return args.Get(0).([]storage.ObjectMeta), args.Error(1)
}

func (m *mockStorage) ObjectsCallback(ctx context.Context, bucket, prefix string, fn func(meta storage.ObjectMeta) bool) error {
	args := m.Called(ctx, bucket, prefix, fn)
	return args.Error(0)
}

func (m *mockStorage) ObjectsCallbackWithStartKey(ctx context.Context, bucket, prefix, startKey string, fn func(meta storage.ObjectMeta) bool) error {
	args := m.Called(ctx, bucket, prefix, startKey, fn)
	return args.Error(0)
}

func (m *mockStorage) Object(ctx context.Context, bucket, key string) (*storage.Object, error) {
	args := m.Called(ctx, bucket, key)
	if args.Get(0) == nil {
		return nil, args.Error(1)
	}
	return args.Get(0).(*storage.Object), args.Error(1)
}

func (m *mockStorage) ObjectExists(ctx context.Context, bucket, key string) (bool, error) {
	args := m.Called(ctx, bucket, key)
	return args.Bool(0), args.Error(1)
}

func (m *mockStorage) ObjectSize(ctx context.Context, bucket, key string) (int64, error) {
	args := m.Called(ctx, bucket, key)
	return args.Get(0).(int64), args.Error(1)
}

func (m *mockStorage) ObjectsDelimiter(ctx context.Context, bucket, prefix string) ([]string, error) {
	args := m.Called(ctx, bucket, prefix)
	if args.Get(0) == nil {
		return nil, args.Error(1)
	}
	return args.Get(0).([]string), args.Error(1)
}

func (m *mockStorage) SaveObject(ctx context.Context, bucket, key string, body io.Reader) error {
	args := m.Called(ctx, bucket, key, body)
	return args.Error(0)
}

func (m *mockStorage) DeleteObject(ctx context.Context, bucket, key string) error {
	args := m.Called(ctx, bucket, key)
	return args.Error(0)
}

func (m *mockStorage) DeleteObjects(ctx context.Context, bucket string, keys []string) error {
	args := m.Called(ctx, bucket, keys)
	return args.Error(0)
}

func TestConProf_List_Validation(t *testing.T) {
	tests := []struct {
		name      string
		params    *ListParams
		expError  string
		setupMock func(*mockStorage)
	}{
		{
			name:     "nil params",
			params:   nil,
			expError: "invalid parameters",
		},
		{
			name: "invalid begin time",
			params: &ListParams{
				ClusterInfo: ClusterInfo{
					Region: "us-west-2",
				},
				BeginTime: 0,
				EndTime:   1000,
			},
			expError: "invalid begin time",
			setupMock: func(m *mockStorage) {
				// No mock setup needed as the error should be returned before any storage calls
			},
		},
		{
			name: "invalid end time",
			params: &ListParams{
				ClusterInfo: ClusterInfo{
					Region: "us-west-2",
				},
				BeginTime: 1000,
				EndTime:   0,
			},
			expError: "invalid end time",
			setupMock: func(m *mockStorage) {
				// No mock setup needed as the error should be returned before any storage calls
			},
		},
		{
			name: "end before begin",
			params: &ListParams{
				ClusterInfo: ClusterInfo{
					Region: "us-west-2",
				},
				BeginTime: 2000,
				EndTime:   1000,
			},
			expError: "invalid time range",
			setupMock: func(m *mockStorage) {
				// No mock setup needed since we expect validation to fail before any storage calls
			},
		},
		{
			name: "time range too large",
			params: &ListParams{
				ClusterInfo: ClusterInfo{
					Region: "us-west-2",
				},
				BeginTime: 960,
				EndTime:   960 + maxTimeRangeSecs + 61,
			},
			expError: "time range too large",
			setupMock: func(m *mockStorage) {
				// Setup mock for the Objects call in tryFindFirstKey
				m.On("Objects", mock.Anything, storage.BucketName("test", "us-west-2", false, false), mock.Anything).
					Return([]storage.ObjectMeta{}, nil).Maybe()

				// Setup mock for ObjectsCallback that might be called
				m.On("ObjectsCallback", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
					Return(nil).Maybe()
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			store := &mockStorage{}
			if tt.setupMock != nil {
				tt.setupMock(store)
			}
			cp := &conprof{env: "test", store: store}

			_, err := cp.List(context.Background(), tt.params)
			require.Error(t, err)
			assert.Contains(t, err.Error(), tt.expError)
			store.AssertExpectations(t)
		})
	}
}

func TestConProf_List_Success(t *testing.T) {
	now := time.Now().Unix()
	begin := now - 300 // 5 minutes ago
	end := now

	tests := []struct {
		name      string
		setupMock func(*mockStorage, *ListParams)
		provider  string
		expected  int
		beginTime int64
		endTime   int64
	}{
		{
			name: "single timestamp",
			setupMock: func(m *mockStorage, params *ListParams) {
				// The actual prefix will include the timestamp rounded down to the nearest minute
				roundedBegin := params.BeginTime - (params.BeginTime % 60)
				prefix := fmt.Sprintf("0/t1/p1/c1/profiles/%d", roundedBegin)
				m.On("Objects", mock.Anything, storage.BucketName("test", params.Region, false, false), prefix).
					Return([]storage.ObjectMeta{
						{Key: prefix + "/1698657360-tidb-cpu-" + base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080"))},
						{Key: prefix + "/1698657360-tikv-cpu-" + base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10081"))},
					}, nil)
			},
			expected:  2,
			beginTime: 1698657360,
			endTime:   1698657360,
		},
		{
			name:     "with first key",
			provider: "aws",
			setupMock: func(m *mockStorage, params *ListParams) {
				// Setup tryFindFirstKey - it will try multiple times with different timestamps
				basePrefix := commonKeyPrefix(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID)
				m.On("Objects", mock.Anything, storage.BucketName("test", params.Region, false, false), mock.MatchedBy(func(prefix string) bool {
					return strings.HasPrefix(prefix, basePrefix)
				})).
					Return([]storage.ObjectMeta{
						{Key: basePrefix + "1698657360-tidb-cpu-" + base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080"))},
					}, nil).Maybe()

				// Setup listObjectMetasFromFirstKey
				dir := commonKeyPrefix(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID)
				m.On("ObjectsCallbackWithStartKey", mock.Anything, storage.BucketName("test", params.Region, false, false), dir,
					mock.Anything, mock.Anything).
					Run(func(args mock.Arguments) {
						fn := args.Get(4).(func(meta storage.ObjectMeta) bool)
						fn(storage.ObjectMeta{Key: dir + "1698657360-tidb-cpu-" + base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080"))})
						fn(storage.ObjectMeta{Key: dir + "1698657360-tikv-cpu-" + base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10081"))})
					}).Return(nil)
			},
			expected:  2,
			beginTime: begin,
			endTime:   end,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			store := &mockStorage{}
			cp := &conprof{
				env:   "test",
				store: store,
			}

			beginTime := tt.beginTime
			endTime := tt.endTime
			if beginTime == 0 {
				beginTime = begin
			}
			if endTime == 0 {
				endTime = end
			}
			params := &ListParams{
				ClusterInfo: ClusterInfo{
					Provider:   tt.provider,
					Region:     "us-west-2",
					TenantID:   "t1",
					ProjectID:  "p1",
					ClusterID:  "c1",
					DeployType: "dedicated",
				},
				BeginTime: beginTime,
				EndTime:   endTime,
			}

			tt.setupMock(store, params)

			result, err := cp.List(context.Background(), params)
			require.NoError(t, err)
			assert.Len(t, result, tt.expected)
		})
	}
}

func TestConProf_Download(t *testing.T) {
	tests := []struct {
		name      string
		provider  string
		setupMock func(*mockStorage, *DownloadParams)
		expError  bool
	}{
		{
			name:     "successful download",
			provider: "aws",
			setupMock: func(m *mockStorage, params *DownloadParams) {
				key := fmt.Sprintf("%s%d-%s-%s-%s.log.gz",
					commonKeyPrefix(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID),
					params.Ts, params.Component, params.ProfileType,
					base64.RawURLEncoding.EncodeToString([]byte(params.Address)))

				// Create a base64-encoded reader
				encodedData := base64.StdEncoding.EncodeToString([]byte("test profile data"))
				m.On("Object", mock.Anything, storage.BucketName("test", params.Region, false, false), key).
					Return(&storage.Object{Reader: io.NopCloser(strings.NewReader(encodedData))}, nil)
			},
			expError: false,
		},
		{
			name:     "not found",
			provider: "gcp",
			setupMock: func(m *mockStorage, params *DownloadParams) {
				key := fmt.Sprintf("%s%d-%s-%s-%s.log",
					commonKeyPrefix(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID),
					params.Ts, params.Component, params.ProfileType,
					base64.RawURLEncoding.EncodeToString([]byte(params.Address)))

				m.On("Object", mock.Anything, storage.BucketName("test", params.Region, false, false), key).
					Return(nil, errors.New("not found"))
			},
			expError: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			store := &mockStorage{}
			cp := &conprof{
				env:   "test",
				store: store,
			}

			params := &DownloadParams{
				ClusterInfo: ClusterInfo{
					Provider:   tt.provider,
					Region:     "us-west-2",
					TenantID:   "t1",
					ProjectID:  "p1",
					ClusterID:  "c1",
					DeployType: "dedicated",
				},
				ProfileMeta: ProfileMeta{
					Ts:          1234567890,
					Address:     "127.0.0.1:10080",
					Component:   "tidb",
					ProfileType: "cpu",
				},
			}

			tt.setupMock(store, params)

			_, err := cp.Download(context.Background(), params)
			if tt.expError {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

func TestHelperFunctions(t *testing.T) {
	t.Run("bucketName", func(t *testing.T) {
		tests := []struct {
			env    string
			region string
			exp    string
		}{
			{"prod", "us-west-2", "o11y-prod-shared-us-west-2"},
			{"staging", "us-east-1", "o11y-test-shared-us-east-1"},
			{"dev", "ap-northeast-1", "o11y-test-shared-ap-northeast-1"},
		}

		for _, tt := range tests {
			t.Run(tt.env+"-"+tt.region, func(t *testing.T) {
				assert.Equal(t, tt.exp, storage.BucketName(tt.env, tt.region, false, false))
			})
		}
	})

	t.Run("commonKeyPrefix", func(t *testing.T) {
		prefix := commonKeyPrefix("dedicated", "t1", "p1", "c1")
		assert.Equal(t, "0/t1/p1/c1/profiles/", prefix)
	})

	t.Run("commonKeyPrefixBetweenBeginEnd", func(t *testing.T) {
		prefix := commonKeyPrefixBetweenBeginEnd("dedicated", "t1", "p1", "c1", 1234567800, 1234567860)
		assert.Contains(t, prefix, "0/t1/p1/c1/profiles/")
	})

	t.Run("extractProfileMetasFromObjectMetas", func(t *testing.T) {
		metas := extractProfileMetasFromObjectMetas([]storage.ObjectMeta{
			{Key: "0/t1/p1/c1/profiles/1698657360-pd-cpu-address1"},
			{Key: "0/t1/p1/c1/profiles/1698657360-tikv-cpu-address2"},
		})
		assert.Len(t, metas, 2)
		assert.Equal(t, "pd", metas[0].Component)
		assert.Equal(t, "tikv", metas[1].Component)
	})
}

func TestExtractProfileMetaFromFilename(t *testing.T) {
	tests := []struct {
		name     string
		filename string
		expected ProfileMeta
	}{
		{
			name:     "valid filename",
			filename: "1698657360-pd-cpu-ZGItcGQtMC5kYi1wZC1wZWVyLnRpZGIxMzczOTMzMDc2NjYyOTIwNTczLnN2YzoyMzc5",
			expected: ProfileMeta{
				Ts:          1698657360,
				Component:   "pd",
				ProfileType: "cpu",
				Address:     "db-pd-0.db-pd-peer.tidb1373933076662920573.svc:2379",
			},
		},
		{
			name:     "invalid format",
			filename: "invalid-format",
			expected: ProfileMeta{},
		},
		{
			name:     "invalid base64 address",
			filename: "1698657360-pd-cpu-invalid-base64!!!",
			expected: ProfileMeta{}, // Invalid format doesn't match regex
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := extractProfileMetaFromFilename(tt.filename)
			assert.Equal(t, tt.expected, result)
		})
	}
}

func TestNewConProf(t *testing.T) {
	store := &mockStorage{}
	cp := NewConProf("test", store)
	assert.NotNil(t, cp)
	_, ok := cp.(*conprof)
	assert.True(t, ok)
}

func TestConProf_List_AzureProvider(t *testing.T) {
	now := time.Now().Unix()
	begin := now - 300
	end := now

	store := &mockStorage{}
	cp := &conprof{
		env:   "test",
		store: store,
	}

	params := &ListParams{
		ClusterInfo: ClusterInfo{
			Provider:   "azure",
			Region:     "us-west-2",
			TenantID:   "t1",
			ProjectID:  "p1",
			ClusterID:  "c1",
			DeployType: "dedicated",
		},
		BeginTime: begin,
		EndTime:   end,
	}

	// Setup tryFindFirstKey - return empty (no first key found)
	basePrefix := commonKeyPrefix(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID)
	store.On("Objects", mock.Anything, storage.BucketName("test", params.Region, false, false), mock.MatchedBy(func(prefix string) bool {
		return strings.HasPrefix(prefix, basePrefix)
	})).
		Return([]storage.ObjectMeta{}, nil).Times(10) // tryFindFirstKey tries up to 10 times

	// Setup listObjectMetasBetweenBeginEndCommonPrefix
	prefix := commonKeyPrefixBetweenBeginEnd(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID, begin, end)
	store.On("ObjectsCallback", mock.Anything, storage.BucketName("test", params.Region, false, false), prefix, mock.Anything).
		Run(func(args mock.Arguments) {
			fn := args.Get(3).(func(meta storage.ObjectMeta) bool)
			ts1 := begin + 60
			ts2 := begin + 120
			baseDir := commonKeyPrefix(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID)
			fn(storage.ObjectMeta{Key: baseDir + fmt.Sprintf("%d-tidb-cpu-%s", ts1, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080")))})
			fn(storage.ObjectMeta{Key: baseDir + fmt.Sprintf("%d-tikv-cpu-%s", ts2, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10081")))})
		}).Return(nil)

	result, err := cp.List(context.Background(), params)
	require.NoError(t, err)
	assert.Len(t, result, 2)
	store.AssertExpectations(t)
}

func TestConProf_List_NoFirstKey(t *testing.T) {
	now := time.Now().Unix()
	begin := now - 300
	end := now

	store := &mockStorage{}
	cp := &conprof{
		env:   "test",
		store: store,
	}

	params := &ListParams{
		ClusterInfo: ClusterInfo{
			Provider:   "gcp",
			Region:     "us-west-2",
			TenantID:   "t1",
			ProjectID:  "p1",
			ClusterID:  "c1",
			DeployType: "dedicated",
		},
		BeginTime: begin,
		EndTime:   end,
	}

	// Setup tryFindFirstKey - return empty (no first key found)
	basePrefix := commonKeyPrefix(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID)
	store.On("Objects", mock.Anything, storage.BucketName("test", params.Region, false, false), mock.MatchedBy(func(prefix string) bool {
		return strings.HasPrefix(prefix, basePrefix)
	})).
		Return([]storage.ObjectMeta{}, nil).Times(10) // tryFindFirstKey tries up to 10 times

	// Setup listObjectMetasBetweenBeginEndCommonPrefix
	prefix := commonKeyPrefixBetweenBeginEnd(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID, begin, end)
	store.On("ObjectsCallback", mock.Anything, storage.BucketName("test", params.Region, false, false), prefix, mock.Anything).
		Run(func(args mock.Arguments) {
			fn := args.Get(3).(func(meta storage.ObjectMeta) bool)
			ts1 := begin + 60
			ts2 := begin + 120
			baseDir := commonKeyPrefix(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID)
			fn(storage.ObjectMeta{Key: baseDir + fmt.Sprintf("%d-tidb-cpu-%s", ts1, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080")))})
			fn(storage.ObjectMeta{Key: baseDir + fmt.Sprintf("%d-tikv-cpu-%s", ts2, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10081")))})
		}).Return(nil)

	result, err := cp.List(context.Background(), params)
	require.NoError(t, err)
	assert.Len(t, result, 2)
	store.AssertExpectations(t)
}

func TestConProf_List_StorageError(t *testing.T) {
	now := time.Now().Unix()
	begin := now - 300
	end := now

	store := &mockStorage{}
	cp := &conprof{
		env:   "test",
		store: store,
	}

	params := &ListParams{
		ClusterInfo: ClusterInfo{
			Provider:   "gcp",
			Region:     "us-west-2",
			TenantID:   "t1",
			ProjectID:  "p1",
			ClusterID:  "c1",
			DeployType: "dedicated",
		},
		BeginTime: begin,
		EndTime:   end,
	}

	// Setup tryFindFirstKey to return error
	basePrefix := commonKeyPrefix(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID)
	store.On("Objects", mock.Anything, storage.BucketName("test", params.Region, false, false), mock.MatchedBy(func(prefix string) bool {
		return strings.HasPrefix(prefix, basePrefix)
	})).
		Return(nil, errors.New("storage error"))

	_, err := cp.List(context.Background(), params)
	require.Error(t, err)
	store.AssertExpectations(t)
}

func TestConProf_List_ObjectsCallbackError(t *testing.T) {
	now := time.Now().Unix()
	begin := now - 300
	end := now

	store := &mockStorage{}
	cp := &conprof{
		env:   "test",
		store: store,
	}

	params := &ListParams{
		ClusterInfo: ClusterInfo{
			Provider:   "azure",
			Region:     "us-west-2",
			TenantID:   "t1",
			ProjectID:  "p1",
			ClusterID:  "c1",
			DeployType: "dedicated",
		},
		BeginTime: begin,
		EndTime:   end,
	}

	// Setup tryFindFirstKey - return empty
	basePrefix := commonKeyPrefix(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID)
	store.On("Objects", mock.Anything, storage.BucketName("test", params.Region, false, false), mock.MatchedBy(func(prefix string) bool {
		return strings.HasPrefix(prefix, basePrefix)
	})).
		Return([]storage.ObjectMeta{}, nil).Times(10)

	// Setup listObjectMetasBetweenBeginEndCommonPrefix to return error
	prefix := commonKeyPrefixBetweenBeginEnd(params.DeployType, params.TenantID, params.ProjectID, params.ClusterID, begin, end)
	store.On("ObjectsCallback", mock.Anything, storage.BucketName("test", params.Region, false, false), prefix, mock.Anything).
		Return(errors.New("callback error"))

	_, err := cp.List(context.Background(), params)
	require.Error(t, err)
	store.AssertExpectations(t)
}

func TestListObjectMetasFromFirstKey_EdgeCases(t *testing.T) {
	dir := "0/t1/p1/c1/profiles/"
	startKey := dir + "1698657360-tidb-cpu-" + base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080"))
	end := int64(1698657500)

	t.Run("invalid filename", func(t *testing.T) {
		store := &mockStorage{}
		store.On("ObjectsCallbackWithStartKey", mock.Anything, "bucket", dir, startKey, mock.Anything).
			Run(func(args mock.Arguments) {
				fn := args.Get(4).(func(meta storage.ObjectMeta) bool)
				// Invalid key (no filename)
				fn(storage.ObjectMeta{Key: dir})
			}).Return(nil)

		result, err := listObjectMetasFromFirstKey(context.Background(), store, "bucket", dir, startKey, end)
		require.NoError(t, err)
		assert.Len(t, result, 0)
	})

	t.Run("invalid timestamp", func(t *testing.T) {
		store := &mockStorage{}
		store.On("ObjectsCallbackWithStartKey", mock.Anything, "bucket", dir, startKey, mock.Anything).
			Run(func(args mock.Arguments) {
				fn := args.Get(4).(func(meta storage.ObjectMeta) bool)
				// Invalid timestamp
				fn(storage.ObjectMeta{Key: dir + "invalid-tidb-cpu-address"})
			}).Return(nil)

		result, err := listObjectMetasFromFirstKey(context.Background(), store, "bucket", dir, startKey, end)
		require.NoError(t, err)
		assert.Len(t, result, 0)
	})

	t.Run("timestamp exceeds end", func(t *testing.T) {
		store := &mockStorage{}
		store.On("ObjectsCallbackWithStartKey", mock.Anything, "bucket", dir, startKey, mock.Anything).
			Run(func(args mock.Arguments) {
				fn := args.Get(4).(func(meta storage.ObjectMeta) bool)
				// Timestamp exceeds end
				fn(storage.ObjectMeta{Key: dir + fmt.Sprintf("%d-tidb-cpu-%s", end+100, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080")))})
			}).Return(nil)

		result, err := listObjectMetasFromFirstKey(context.Background(), store, "bucket", dir, startKey, end)
		require.NoError(t, err)
		assert.Len(t, result, 0)
	})

	t.Run("callback error", func(t *testing.T) {
		store := &mockStorage{}
		store.On("ObjectsCallbackWithStartKey", mock.Anything, "bucket", dir, startKey, mock.Anything).
			Return(errors.New("callback error"))

		_, err := listObjectMetasFromFirstKey(context.Background(), store, "bucket", dir, startKey, end)
		require.Error(t, err)
	})
}

func TestListObjectMetasBetweenBeginEndCommonPrefix(t *testing.T) {
	prefix := "0/t1/p1/c1/profiles/1698"
	begin := int64(1698657360)
	end := int64(1698657500)

	t.Run("success", func(t *testing.T) {
		store := &mockStorage{}
		store.On("ObjectsCallback", mock.Anything, "bucket", prefix, mock.Anything).
			Run(func(args mock.Arguments) {
				fn := args.Get(3).(func(meta storage.ObjectMeta) bool)
				ts1 := begin + 60
				ts2 := begin + 120
				baseDir := "0/t1/p1/c1/profiles/"
				fn(storage.ObjectMeta{Key: baseDir + fmt.Sprintf("%d-tidb-cpu-%s", ts1, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080")))})
				fn(storage.ObjectMeta{Key: baseDir + fmt.Sprintf("%d-tikv-cpu-%s", ts2, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10081")))})
			}).Return(nil)

		result, err := listObjectMetasBetweenBeginEndCommonPrefix(context.Background(), store, "bucket", prefix, begin, end)
		require.NoError(t, err)
		assert.Len(t, result, 2)
	})

	t.Run("timestamp before begin", func(t *testing.T) {
		store := &mockStorage{}
		baseDir := "0/t1/p1/c1/profiles/"
		store.On("ObjectsCallback", mock.Anything, "bucket", prefix, mock.Anything).
			Run(func(args mock.Arguments) {
				fn := args.Get(3).(func(meta storage.ObjectMeta) bool)
				// Timestamp before begin (should be skipped)
				fn(storage.ObjectMeta{Key: baseDir + fmt.Sprintf("%d-tidb-cpu-%s", begin-100, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080")))})
				// Valid timestamp
				fn(storage.ObjectMeta{Key: baseDir + fmt.Sprintf("%d-tikv-cpu-%s", begin+60, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10081")))})
			}).Return(nil)

		result, err := listObjectMetasBetweenBeginEndCommonPrefix(context.Background(), store, "bucket", prefix, begin, end)
		require.NoError(t, err)
		assert.Len(t, result, 1)
	})

	t.Run("timestamp exceeds end", func(t *testing.T) {
		store := &mockStorage{}
		baseDir := "0/t1/p1/c1/profiles/"
		store.On("ObjectsCallback", mock.Anything, "bucket", prefix, mock.Anything).
			Run(func(args mock.Arguments) {
				fn := args.Get(3).(func(meta storage.ObjectMeta) bool)
				// Valid timestamp
				fn(storage.ObjectMeta{Key: baseDir + fmt.Sprintf("%d-tidb-cpu-%s", begin+60, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080")))})
				// Timestamp exceeds end (should stop)
				fn(storage.ObjectMeta{Key: baseDir + fmt.Sprintf("%d-tikv-cpu-%s", end+100, base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10081")))})
			}).Return(nil)

		result, err := listObjectMetasBetweenBeginEndCommonPrefix(context.Background(), store, "bucket", prefix, begin, end)
		require.NoError(t, err)
		assert.Len(t, result, 1)
	})

	t.Run("invalid filename", func(t *testing.T) {
		store := &mockStorage{}
		store.On("ObjectsCallback", mock.Anything, "bucket", prefix, mock.Anything).
			Run(func(args mock.Arguments) {
				fn := args.Get(3).(func(meta storage.ObjectMeta) bool)
				fn(storage.ObjectMeta{Key: prefix})
			}).Return(nil)

		result, err := listObjectMetasBetweenBeginEndCommonPrefix(context.Background(), store, "bucket", prefix, begin, end)
		require.NoError(t, err)
		assert.Len(t, result, 0)
	})

	t.Run("callback error", func(t *testing.T) {
		store := &mockStorage{}
		store.On("ObjectsCallback", mock.Anything, "bucket", prefix, mock.Anything).
			Return(errors.New("callback error"))

		_, err := listObjectMetasBetweenBeginEndCommonPrefix(context.Background(), store, "bucket", prefix, begin, end)
		require.Error(t, err)
	})
}

func TestTryFindFirstKey(t *testing.T) {
	bucket := "bucket"
	deployType := "dedicated"
	tenant := "t1"
	project := "p1"
	cluster := "c1"
	begin := int64(1698657360)

	t.Run("found on first try", func(t *testing.T) {
		store := &mockStorage{}
		prefix := commonKeyPrefixBetweenBeginEnd(deployType, tenant, project, cluster, begin, begin)
		store.On("Objects", mock.Anything, bucket, prefix).
			Return([]storage.ObjectMeta{
				{Key: prefix + "1698657360-tidb-cpu-address"},
			}, nil)

		result, err := tryFindFirstKey(context.Background(), store, bucket, deployType, tenant, project, cluster, begin)
		require.NoError(t, err)
		assert.Equal(t, prefix+"1698657360-tidb-cpu-address", result)
	})

	t.Run("found on second try", func(t *testing.T) {
		store := &mockStorage{}
		prefix1 := commonKeyPrefixBetweenBeginEnd(deployType, tenant, project, cluster, begin, begin)
		prefix2 := commonKeyPrefixBetweenBeginEnd(deployType, tenant, project, cluster, begin+60, begin+60)
		store.On("Objects", mock.Anything, bucket, prefix1).
			Return([]storage.ObjectMeta{}, nil)
		store.On("Objects", mock.Anything, bucket, prefix2).
			Return([]storage.ObjectMeta{
				{Key: prefix2 + "1698657420-tidb-cpu-address"},
			}, nil)

		result, err := tryFindFirstKey(context.Background(), store, bucket, deployType, tenant, project, cluster, begin)
		require.NoError(t, err)
		assert.Equal(t, prefix2+"1698657420-tidb-cpu-address", result)
	})

	t.Run("not found", func(t *testing.T) {
		store := &mockStorage{}
		store.On("Objects", mock.Anything, bucket, mock.MatchedBy(func(p string) bool {
			return strings.HasPrefix(p, "0/t1/p1/c1/profiles/")
		})).
			Return([]storage.ObjectMeta{}, nil).Times(10)

		result, err := tryFindFirstKey(context.Background(), store, bucket, deployType, tenant, project, cluster, begin)
		require.NoError(t, err)
		assert.Equal(t, "", result)
	})

	t.Run("error", func(t *testing.T) {
		store := &mockStorage{}
		prefix := commonKeyPrefixBetweenBeginEnd(deployType, tenant, project, cluster, begin, begin)
		store.On("Objects", mock.Anything, bucket, prefix).
			Return(nil, errors.New("storage error"))

		_, err := tryFindFirstKey(context.Background(), store, bucket, deployType, tenant, project, cluster, begin)
		require.Error(t, err)
	})
}

func TestExtractProfileMetasFromObjectMetas_EdgeCases(t *testing.T) {
	t.Run("empty metas", func(t *testing.T) {
		result := extractProfileMetasFromObjectMetas([]storage.ObjectMeta{})
		assert.Len(t, result, 0)
	})

	t.Run("invalid filename", func(t *testing.T) {
		result := extractProfileMetasFromObjectMetas([]storage.ObjectMeta{
			{Key: "0/t1/p1/c1/profiles/"},
		})
		assert.Len(t, result, 0)
	})

	t.Run("invalid profile meta", func(t *testing.T) {
		result := extractProfileMetasFromObjectMetas([]storage.ObjectMeta{
			{Key: "0/t1/p1/c1/profiles/invalid-format"},
		})
		assert.Len(t, result, 0)
	})

	t.Run("missing required fields", func(t *testing.T) {
		// Create a filename that matches regex but has missing fields
		filename := "1698657360-pd-cpu-"
		result := extractProfileMetasFromObjectMetas([]storage.ObjectMeta{
			{Key: "0/t1/p1/c1/profiles/" + filename},
		})
		assert.Len(t, result, 0)
	})
}

func TestBucketName_WithEnvVar(t *testing.T) {
	originalValue := os.Getenv("BUCKET_NAME")
	defer func() {
		if originalValue != "" {
			os.Setenv("BUCKET_NAME", originalValue)
		} else {
			os.Unsetenv("BUCKET_NAME")
		}
	}()

	os.Setenv("BUCKET_NAME", "custom-bucket")
	result := storage.BucketName("test", "us-west-2", false, false)
	assert.Equal(t, "custom-bucket", result)
}

func TestCommonKeyPrefix_Premium(t *testing.T) {
	prefix := commonKeyPrefix("premium", "t1", "p1", "c1")
	assert.Equal(t, "0/shared_pool/c1/profiles/", prefix)
}

func TestLongestCommonPrefix_EdgeCases(t *testing.T) {
	t.Run("empty slice", func(t *testing.T) {
		result := longestCommonPrefix([]string{})
		assert.Equal(t, "", result)
	})

	t.Run("single string", func(t *testing.T) {
		result := longestCommonPrefix([]string{"abc"})
		assert.Equal(t, "abc", result)
	})

	t.Run("no common prefix", func(t *testing.T) {
		result := longestCommonPrefix([]string{"abc", "def", "ghi"})
		assert.Equal(t, "", result)
	})

	t.Run("one string shorter", func(t *testing.T) {
		result := longestCommonPrefix([]string{"abc", "ab"})
		assert.Equal(t, "ab", result)
	})

	t.Run("all same", func(t *testing.T) {
		result := longestCommonPrefix([]string{"abc", "abc", "abc"})
		assert.Equal(t, "abc", result)
	})
}

func TestConProf_Download_MoreProviders(t *testing.T) {
	tests := []struct {
		name     string
		provider string
		expKey   string
	}{
		{
			name:     "azure provider",
			provider: "azure",
			expKey: fmt.Sprintf("%s%d-%s-%s-%s.log.gz",
				commonKeyPrefix("dedicated", "t1", "p1", "c1"),
				1234567890, "tidb", "cpu",
				base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080"))),
		},
		{
			name:     "gcp provider",
			provider: "gcp",
			expKey: fmt.Sprintf("%s%d-%s-%s-%s.log",
				commonKeyPrefix("dedicated", "t1", "p1", "c1"),
				1234567890, "tidb", "cpu",
				base64.RawURLEncoding.EncodeToString([]byte("127.0.0.1:10080"))),
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			store := &mockStorage{}
			cp := &conprof{
				env:   "test",
				store: store,
			}

			params := &DownloadParams{
				ClusterInfo: ClusterInfo{
					Provider:   tt.provider,
					Region:     "us-west-2",
					TenantID:   "t1",
					ProjectID:  "p1",
					ClusterID:  "c1",
					DeployType: "dedicated",
				},
				ProfileMeta: ProfileMeta{
					Ts:          1234567890,
					Address:     "127.0.0.1:10080",
					Component:   "tidb",
					ProfileType: "cpu",
				},
			}

			encodedData := base64.StdEncoding.EncodeToString([]byte("test profile data"))
			store.On("Object", mock.Anything, storage.BucketName("test", params.Region, false, false), tt.expKey).
				Return(&storage.Object{Reader: io.NopCloser(strings.NewReader(encodedData))}, nil)

			result, err := cp.Download(context.Background(), params)
			require.NoError(t, err)
			assert.NotEmpty(t, result)
			store.AssertExpectations(t)
		})
	}
}