package process

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"testing"
	"time"

	_ "github.com/mattn/go-sqlite3"
	"github.com/parquet-go/parquet-go"
	"github.com/stretchr/testify/assert"
	"github.com/tidbcloud/observability/cmd/diagnosis-compute/pkg/config"
	parquetpkg "github.com/tidbcloud/observability/cmd/diagnosis-compute/pkg/parquet"
	"github.com/tidbcloud/observability/cmd/diagnosis-compute/pkg/types"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/storage"
	"golang.org/x/sync/semaphore"
)

// Mock Storage to avoid S3 dependency
type MockStorage2 struct {
	dir string
}

func NewMockStorage2(dir string) *MockStorage2 {
	return &MockStorage2{
		dir: dir,
	}
}

// Implement storage.Storage interface
func (m *MockStorage2) SaveObject(ctx context.Context, bucket, key string, body io.Reader) error {
	// write to local filesystem under m.dir/<bucket>/<key>
	base := filepath.Join(m.dir, bucket)
	fullPath := filepath.Join(base, filepath.FromSlash(key))
	if err := os.MkdirAll(filepath.Dir(fullPath), 0o755); err != nil {
		return err
	}
	f, err := os.Create(fullPath)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(f, body)
	return err
}

// Implement storage.Storage interface
func (m *MockStorage2) Object(ctx context.Context, bucket, key string) (*storage.Object, error) {
	base := filepath.Join(m.dir, bucket)
	fullPath := filepath.Join(base, filepath.FromSlash(key))
	fi, err := os.Stat(fullPath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, storage.ErrNotFound
		}
		return nil, err
	}
	f, err := os.Open(fullPath)
	if err != nil {
		return nil, err
	}
	size := fi.Size()
	return &storage.Object{
		ContentLength: &size,
		Reader:        f,
	}, nil
}

// Implement storage.Storage interface
func (m *MockStorage2) ObjectExists(ctx context.Context, bucket, key string) (bool, error) {
	base := filepath.Join(m.dir, bucket)
	fullPath := filepath.Join(base, filepath.FromSlash(key))
	_, err := os.Stat(fullPath)
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

// Implement storage.Storage interface
func (m *MockStorage2) ObjectSize(ctx context.Context, bucket, key string) (int64, error) {
	base := filepath.Join(m.dir, bucket)
	fullPath := filepath.Join(base, filepath.FromSlash(key))
	fi, err := os.Stat(fullPath)
	if err != nil {
		if os.IsNotExist(err) {
			return 0, nil
		}
		return 0, err
	}
	return fi.Size(), nil
}

// Implement storage.Storage interface
func (m *MockStorage2) Objects(ctx context.Context, bucket, prefix string) ([]storage.ObjectMeta, error) {
	base := filepath.Join(m.dir, bucket)
	var metas []storage.ObjectMeta
	// walk the bucket directory
	err := filepath.Walk(base, func(pathName string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		rel, err := filepath.Rel(base, pathName)
		if err != nil {
			return err
		}
		key := filepath.ToSlash(rel)
		if prefix == "" || strings.HasPrefix(key, prefix) {
			metas = append(metas, storage.ObjectMeta{Key: key, Size: info.Size()})
		}
		return nil
	})
	if err != nil {
		// if bucket dir does not exist, return empty list
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	return metas, nil
}

// Implement storage.Storage interface
func (m *MockStorage2) ObjectsCallback(ctx context.Context, bucket, prefix string, callback func(storage.ObjectMeta) bool) error {
	metas, err := m.Objects(ctx, bucket, prefix)
	if err != nil {
		return err
	}
	// deterministic order
	keys := make([]string, 0, len(metas))
	metaMap := make(map[string]storage.ObjectMeta, len(metas))
	for _, mm := range metas {
		keys = append(keys, mm.Key)
		metaMap[mm.Key] = mm
	}
	sort.Strings(keys)
	for _, k := range keys {
		if !callback(metaMap[k]) {
			return nil
		}
	}
	return nil
}

// Implement storage.Storage interface
func (m *MockStorage2) ObjectsCallbackWithStartKey(ctx context.Context, bucket, prefix, startKey string, callback func(storage.ObjectMeta) bool) error {
	metas, err := m.Objects(ctx, bucket, prefix)
	if err != nil {
		return err
	}
	// deterministic order
	keys := make([]string, 0, len(metas))
	metaMap := make(map[string]storage.ObjectMeta, len(metas))
	for _, mm := range metas {
		keys = append(keys, mm.Key)
		metaMap[mm.Key] = mm
	}
	sort.Strings(keys)
	skip := startKey != ""
	for _, k := range keys {
		if skip {
			if k == startKey {
				skip = false
			}
			continue
		}
		if !callback(metaMap[k]) {
			return nil
		}
	}
	return nil
}

// Implement storage.Storage interface
func (m *MockStorage2) ObjectsDelimiter(ctx context.Context, bucket, prefix string) ([]string, error) {
	metas, err := m.Objects(ctx, bucket, "")
	if err != nil {
		return nil, err
	}
	prefixes := make(map[string]struct{})
	for _, mm := range metas {
		key := mm.Key
		if prefix != "" && !strings.HasPrefix(key, prefix) {
			continue
		}
		start := len(prefix)
		if prefix == "" {
			start = 0
		}
		nextSlash := -1
		for i := start; i < len(key); i++ {
			if key[i] == '/' {
				nextSlash = i
				break
			}
		}
		if nextSlash != -1 {
			prefixes[key[:nextSlash+1]] = struct{}{}
		}
	}
	result := make([]string, 0, len(prefixes))
	for p := range prefixes {
		result = append(result, p)
	}
	return result, nil
}

func (m *MockStorage2) DeleteObject(ctx context.Context, bucket, key string) error {
	base := filepath.Join(m.dir, bucket)
	fullPath := filepath.Join(base, filepath.FromSlash(key))
	err := os.Remove(fullPath)
	if err != nil && os.IsNotExist(err) {
		return nil
	}
	return err
}

func (m *MockStorage2) DeleteObjects(ctx context.Context, bucket string, keys []string) error {
	for _, key := range keys {
		if err := m.DeleteObject(ctx, bucket, key); err != nil {
			return err
		}
	}
	return nil
}

func prepareObjectsForTest(ctx context.Context, store *MockStorage2, keyPrefix string) {
	// Save TiDB logs
	store.SaveObject(ctx, config.BucketName(), keyPrefix+"tidb/1700464131-a.log", strings.NewReader(`2023-11-20T07:05:56.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tidb-0;tidb [2023/11/20 07:05:56.858 +00:00] [INFO] [cpuprofile.go:113] ["test"]
	2023-11-20T07:05:59.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tidb-0;tidb [2023/11/20 07:05:59.858 +00:00] [INFO] [cpuprofile.go:113] ["test1"]
	`))

	store.SaveObject(ctx, config.BucketName(), keyPrefix+"tidb/1700464191-b.log", strings.NewReader(`2023-11-20T07:06:56.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tidb-0;tidb [2023/11/20 07:06:56.858 +00:00] [INFO] [cpuprofile.go:113] ["test2"]
	2023-11-20T07:06:59.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tidb-0;tidb [2023/11/20 07:06:59.858 +00:00] [INFO] [cpuprofile.go:113] ["test3"]
	`))

	store.SaveObject(ctx, config.BucketName(), keyPrefix+"tidb/1700464141-a.log", strings.NewReader(`2023-11-20T07:06:06.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tidb-1;tidb [2023/11/20 07:06:06.858 +00:00] [INFO] [cpuprofile.go:113] ["test"]
	2023-11-20T07:06:09.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tidb-1;tidb [2023/11/20 07:06:09.858 +00:00] [INFO] [cpuprofile.go:113] ["test1"]
	`))

	// Save TiKV logs
	store.SaveObject(ctx, config.BucketName(), keyPrefix+"tikv/1700464131-a.log", strings.NewReader(`2023-11-20T07:05:56.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tikv-0;tikv [2023/11/20 07:05:56.858 +00:00] [INFO] [lib.rs:89] ["test"]
	2023-11-20T07:05:59.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tikv-0;tikv [2023/11/20 07:05:59.858 +00:00] [INFO] [lib.rs:89] ["test1"]
	`))

	store.SaveObject(ctx, config.BucketName(), keyPrefix+"tikv/1700464191-b.log", strings.NewReader(`2023-11-20T07:06:56.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tikv-0;tikv [2023/11/20 07:06:56.858 +00:00] [INFO] [lib.rs:89] ["test2"]
	2023-11-20T07:06:59.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tikv-0;tikv [2023/11/20 07:06:59.858 +00:00] [INFO] [lib.rs:89] ["test3"]
	`))

	store.SaveObject(ctx, config.BucketName(), keyPrefix+"tikv/1700464141-a.log", strings.NewReader(`2023-11-20T07:06:06.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tikv-1;tikv [2023/11/20 07:06:06.858 +00:00] [INFO] [lib.rs:89] ["test"]
	2023-11-20T07:06:09.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tikv-1;tikv [2023/11/20 07:06:09.858 +00:00] [INFO] [lib.rs:89] ["test1"]
	`))

	// Save PD logs
	store.SaveObject(ctx, config.BucketName(), keyPrefix+"pd/1700464131-a.log", strings.NewReader(`2023-11-20T07:05:56.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-pd-0;pd [2023/11/20 07:05:56.858 +00:00] [INFO] [lib.rs:89] ["test"]
	2023-11-20T07:05:59.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-pd-0;pd [2023/11/20 07:05:59.858 +00:00] [INFO] [lib.rs:89] ["test1"]
	`))

	store.SaveObject(ctx, config.BucketName(), keyPrefix+"pd/1700464191-b.log", strings.NewReader(`2023-11-20T07:06:56.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-pd-0;pd [2023/11/20 07:06:56.858 +00:00] [INFO] [lib.rs:89] ["test2"]
	2023-11-20T07:06:59.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-pd-0;pd [2023/11/20 07:06:59.858 +00:00] [INFO] [lib.rs:89] ["test3"]
	`))

	store.SaveObject(ctx, config.BucketName(), keyPrefix+"pd/1700464141-a.log", strings.NewReader(`2023-11-20T07:06:06.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-pd-1;pd [2023/11/20 07:06:06.858 +00:00] [INFO] [lib.rs:89] ["test"]
	2023-11-20T07:06:09.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-pd-1;pd [2023/11/20 07:06:09.858 +00:00] [INFO] [lib.rs:89] ["test1"]
	`))
}

func TestScanScanComponents(t *testing.T) {
	ctx := context.Background()
	storeDir := t.TempDir()
	store := NewMockStorage2(storeDir)

	clusterInfo := types.ClusterInfo{
		Type:      0,
		TenantID:  "0",
		ProjectID: "1",
		ClusterID: "2",
		StorageID: "2",
	}
	keyPrefix := fmt.Sprintf("%d/%s/%s/%s/logs/", clusterInfo.Type, clusterInfo.TenantID, clusterInfo.ProjectID, clusterInfo.ClusterID)
	prepareObjectsForTest(ctx, store, keyPrefix)

	components, err := scanComponents(ctx, store, clusterInfo)
	assert.NoError(t, err)
	assert.ElementsMatch(t, []string{"pd", "tidb", "tikv"}, components)

	// Add an extra component not in whitelist
	err = store.SaveObject(ctx, config.BucketName(), keyPrefix+"unknown/1700464131-a.log", strings.NewReader(`
	2023-11-20T07:05:56.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-unknown-0;unknown [2023/11/20 07:05:56.858 +00:00] [INFO] [lib.rs:89] ["test"]
	`))
	assert.NoError(t, err)
	assert.ElementsMatch(t, []string{"pd", "tidb", "tikv"}, components)

	// Test with another clusterID
	components, err = scanComponents(ctx, store, types.ClusterInfo{
		Type:      0,
		TenantID:  "0",
		ProjectID: "1",
		ClusterID: "3",
		StorageID: "3",
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []string{}, components)
}

func TestRunClusterHourlyTask(t *testing.T) {
	ctx := context.Background()
	storeDir := t.TempDir()
	store := NewMockStorage2(storeDir)

	config.S3sem = semaphore.NewWeighted(5)
	config.ClusterTaskSem = semaphore.NewWeighted(5)
	config.LocalMergedFileDir = t.TempDir()

	clusterInfo := types.ClusterInfo{
		Type:        0,
		Env:         "prod",
		TenantID:    "0",
		ProjectID:   "1",
		ClusterID:   "2",
		StorageID:   "2",
		ClusterType: types.ClusterTypeDedicated,
	}
	keyPrefix := fmt.Sprintf("%d/%s/%s/%s/logs/", clusterInfo.Type, clusterInfo.TenantID, clusterInfo.ProjectID, clusterInfo.ClusterID)
	prepareObjectsForTest(ctx, store, keyPrefix)

	hourBegin := int64(1700463600) // 2023-11-20 07:00:00 UTC
	assert.True(t, runClusterHourlyTask(ctx, store, clusterInfo, hourBegin))

	parquetKey := fmt.Sprintf(
		"deltalake/global/logs/env=prod/deployType=dedicated/date=2023-11-20/component=tidb/07_%s_01.parquet",
		clusterInfo.ClusterID,
	)
	exists, err := store.ObjectExists(ctx, config.BucketName(), parquetKey)
	assert.NoError(t, err)
	assert.True(t, exists)

	obj, err := store.Object(ctx, config.BucketName(), parquetKey)
	assert.NoError(t, err)
	assert.NotNil(t, obj)
	parquetBytes, err := io.ReadAll(obj.Reader)
	assert.NoError(t, err)
	assert.NoError(t, obj.Reader.Close())

	rows, err := parquet.Read[parquetpkg.LogRow](bytes.NewReader(parquetBytes), int64(len(parquetBytes)))
	assert.NoError(t, err)
	assert.NotEmpty(t, rows)
	assert.Equal(t, clusterInfo.TenantID, rows[0].TenantID)
	assert.Equal(t, clusterInfo.ClusterID, rows[0].Cluster)

	mergedKey := fmt.Sprintf("diagnosis/data/%s/merged-logs/", clusterInfo.ClusterID)
	metas, err := store.Objects(ctx, config.BucketName(), mergedKey)
	assert.NoError(t, err)
	assert.Empty(t, metas)
}

func TestAnalyzeClusterTask(t *testing.T) {
	ctx := context.Background()
	storeDir := t.TempDir()
	store := NewMockStorage2(storeDir)

	config.S3sem = semaphore.NewWeighted(5)
	config.ClusterTaskSem = semaphore.NewWeighted(5)
	config.LocalMergedFileDir = t.TempDir()

	clusterInfo := types.ClusterInfo{
		Type:        0,
		Env:         "prod",
		TenantID:    "0",
		ProjectID:   "1",
		ClusterID:   "2",
		StorageID:   "2",
		ClusterType: types.ClusterTypeDedicated,
	}

	keyPrefix := fmt.Sprintf("%d/%s/%s/%s/logs/", clusterInfo.Type, clusterInfo.TenantID, clusterInfo.ProjectID, clusterInfo.ClusterID)
	prepareObjectsForTest(ctx, store, keyPrefix)

	store.SaveObject(ctx, config.BucketName(), keyPrefix+"tidb/1700467143-a.log", strings.NewReader("2023-11-20T07:45:36.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tidb-0;tidb [2023/11/20 07:45:56.858 +00:00] [INFO] [lib.rs:89] [\"test\"]\n2023-11-20T07:45:36.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tidb-0;tidb [2023/11/20 07:45:56.858 +00:00] [INFO] [lib.rs:89] [\"test\"]"))
	store.SaveObject(ctx, config.BucketName(), keyPrefix+"tidb/1700467144-a.log", strings.NewReader("2023-11-20T07:45:36.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tidb-0;tidb [2023/11/20 07:45:56.858 +00:00] [INFO] [lib.rs:89] [\"test\"]\n2023-11-20T07:45:36.858176156Z;ip-172-30-8-211.ec2.internal;tidb1379661944646413139;db-tidb-0;tidb [2023/11/20 07:45:56.858 +00:00] [INFO] [lib.rs:89] [\"test\"]\n"+`2025-11-17T02:32:25.521558569Z;ip-172-21-93-194.ec2.internal;tidb10913508581581239630;db-tidb-0;tidb [2025/12/11 07:59:28.182 +00:00] [INFO] [telemetry.go:66] [category=telemetry] ["telemetry data"="{\"multiSchemaChange\":{\"multi_schema_change_used\":0}}"]`))

	hourBegin := int64(1700463600)
	assert.True(t, runClusterHourlyTask(ctx, store, clusterInfo, hourBegin))

	s3key := fmt.Sprintf("diagnosis/data/%s/log-stats/1h/%s.db", clusterInfo.StorageID, "2023112007")
	obj, err := store.Object(ctx, config.BucketName(), s3key)
	assert.NoError(t, err)
	assert.NotNil(t, obj)

	statsBytes, err := io.ReadAll(obj.Reader)
	assert.NoError(t, err)
	assert.Greater(t, len(statsBytes), 0)

	// Check the stats value
	sqlDBPath := filepath.Join(storeDir, "test.db")
	f, err := os.Create(sqlDBPath)
	assert.NoError(t, err)
	_, err = io.Copy(f, strings.NewReader(string(statsBytes)))
	assert.NoError(t, err)
	assert.NoError(t, f.Close())

	db, err := sql.Open("sqlite3", sqlDBPath)
	assert.NoError(t, err)

	// check logstats results
	res, err := db.Query("select component, level, count from log_stats where component='tidb' order by count desc limit 1")
	assert.NoError(t, err)
	assert.True(t, res.Next())
	var component, level string
	var count int64
	err = res.Scan(&component, &level, &count)
	assert.NoError(t, err)
	assert.Equal(t, "tidb", component)
	assert.Equal(t, "INFO", level)
	assert.Greater(t, count, int64(0))
	assert.NoError(t, res.Close())

	// check telemetry results
	res, err = db.Query("select ts,cluster_id,instance,telemetry_data from telemetry_stats")
	assert.NoError(t, err)
	assert.True(t, res.Next())
	var ts int64
	var clusterID, instance, telemetryData string
	err = res.Scan(&ts, &clusterID, &instance, &telemetryData)
	assert.NoError(t, err)
	assert.Equal(t, clusterInfo.ClusterID, clusterID)
	assert.Equal(t, "db-tidb-0", instance)
	assert.True(t, json.Valid([]byte(telemetryData)))
	assert.Contains(t, telemetryData, "\"multi_schema_change_used\":0")
	assert.NoError(t, res.Close())

	assert.NoError(t, db.Close())
}

func TestRunClusterHourlyTaskAggregatesUnknownLogfmtByPattern(t *testing.T) {
	ctx := context.Background()
	storeDir := t.TempDir()
	store := NewMockStorage2(storeDir)

	config.S3sem = semaphore.NewWeighted(5)
	config.ClusterTaskSem = semaphore.NewWeighted(5)
	config.LocalMergedFileDir = t.TempDir()

	clusterInfo := types.ClusterInfo{
		Type:        0,
		Env:         "prod",
		TenantID:    "0",
		ProjectID:   "1",
		ClusterID:   "2",
		StorageID:   "2",
		ClusterType: types.ClusterTypeDedicated,
	}

	keyPrefix := fmt.Sprintf("%d/%s/%s/%s/logs/", clusterInfo.Type, clusterInfo.TenantID, clusterInfo.ProjectID, clusterInfo.ClusterID)
	prepareObjectsForTest(ctx, store, keyPrefix)

	hourBegin := time.Date(2022, 10, 16, 21, 0, 0, 0, time.UTC).Unix()
	logTs := hourBegin + 17*60 + 44
	logLines := strings.Join([]string{
		"2022-10-16T21:17:44.085464273Z;gke-prod-seed-asia-n-gardener-node-po-31add738-37yl;monitoring;prometheus-operator-b4bf9fb96-s6hd4;prometheus-operator level=info ts=2022-10-16T21:17:44.085055757Z caller=operator.go:200 component=alertmanageroperator msg=\"connection established\" cluster-version=v1.21.14-gke.2700",
		"2022-10-16T21:17:44.085538922Z;gke-prod-seed-asia-n-gardener-node-po-31add738-37yl;monitoring;prometheus-operator-b4bf9fb96-s6hd4;prometheus-operator level=info ts=2022-10-16T21:17:44.085207501Z caller=operator.go:200 component=alertmanageroperator msg=\"connection established\" cluster-version=v1.21.14-gke.2700",
		"2022-10-16T21:17:44.085684571Z;gke-prod-seed-asia-n-gardener-node-po-31add738-37yl;monitoring;prometheus-operator-b4bf9fb96-s6hd4;prometheus-operator level=info ts=2022-10-16T21:17:44.085060987Z caller=operator.go:200 component=alertmanageroperator msg=\"connection established in 123ms\" cluster-version=v1.21.14-gke.2700",
		"2022-10-16T21:17:45.689906293Z;gke-prod-seed-asia-n-gardener-node-po-31add738-37yl;monitoring;prometheus-operator-b4bf9fb96-s6hd4;prometheus-operator level=info ts=2022-10-16T21:17:45.689485025Z caller=operator.go:171 component=alertmanageroperator msg=\"successfully synced all caches\"",
	}, "\n")
	store.SaveObject(ctx, config.BucketName(), keyPrefix+fmt.Sprintf("prometheus-operator/%d-a.log", logTs), strings.NewReader(logLines))

	assert.True(t, runClusterHourlyTask(ctx, store, clusterInfo, hourBegin))

	s3key := fmt.Sprintf("diagnosis/data/%s/log-stats/1h/%s.db", clusterInfo.StorageID, "2022101621")
	obj, err := store.Object(ctx, config.BucketName(), s3key)
	assert.NoError(t, err)
	assert.NotNil(t, obj)

	statsBytes, err := io.ReadAll(obj.Reader)
	assert.NoError(t, err)
	assert.Greater(t, len(statsBytes), 0)

	sqlDBPath := filepath.Join(storeDir, "unknown-logfmt.db")
	f, err := os.Create(sqlDBPath)
	assert.NoError(t, err)
	_, err = io.Copy(f, strings.NewReader(string(statsBytes)))
	assert.NoError(t, err)
	assert.NoError(t, f.Close())

	db, err := sql.Open("sqlite3", sqlDBPath)
	assert.NoError(t, err)

	res, err := db.Query(`
		select component, level, position, msg_digest, count
		from log_stats
		where component = 'prometheus-operator'
		order by count desc, position asc, msg_digest asc
	`)
	assert.NoError(t, err)

	type row struct {
		component string
		level     string
		position  string
		msgDigest string
		count     int64
	}
	var rows []row
	for res.Next() {
		var r row
		err = res.Scan(&r.component, &r.level, &r.position, &r.msgDigest, &r.count)
		assert.NoError(t, err)
		rows = append(rows, r)
	}
	assert.NoError(t, res.Close())
	assert.Len(t, rows, 2)

	assert.Equal(t, row{
		component: "prometheus-operator",
		level:     "INFO",
		position:  "operator.go:200",
		msgDigest: `level=info ts=?-?-?T?:?:?.?Z caller=operator.go:? component=alertmanageroperator msg="connection established" cluster-version=v?.?.?-gke.?`,
		count:     3,
	}, rows[0])
	assert.Equal(t, row{
		component: "prometheus-operator",
		level:     "INFO",
		position:  "operator.go:171",
		msgDigest: `level=info ts=?-?-?T?:?:?.?Z caller=operator.go:? component=alertmanageroperator msg="successfully synced all caches"`,
		count:     1,
	}, rows[1])

	assert.NoError(t, db.Close())
}

func TestLoadClusterInfo(t *testing.T) {
	ctx := context.Background()
	storeDir := t.TempDir()
	store := NewMockStorage2(storeDir)

	info := types.ClusterInfo{
		Type:           0,
		TenantID:       "0",
		ProjectID:      "1",
		ClusterID:      "2",
		StorageID:      "2",
		ClusterType:    types.ClusterTypeNextgen,
		SubClusterType: types.SubClusterTypePhysical,
	}

	key := "diagnosis/tasks/processor-dedicated/1.yaml"
	store.SaveObject(ctx, config.BucketName(), key, strings.NewReader(`
meta:
  tenantID: "0"
  projectID: "1"
  clusterID: "2"
  clusterType: nextgen
`))
	_, err := LoadClusterInfo(ctx, store, config.BucketName(), key)
	assert.Error(t, err)

	// check logical cluster dir first
	info.SubClusterType = types.SubClusterTypeLogical
	store.SaveObject(ctx, config.BucketName(), getLogPrefix(info)+"tidb/a.log", strings.NewReader("test"))
	info2, err := LoadClusterInfo(ctx, store, config.BucketName(), key)
	assert.NoError(t, err)
	assert.Equal(t, info2.SubClusterType, types.SubClusterTypeLogical)

	// check physical cluster dir next
	info.SubClusterType = types.SubClusterTypePhysical
	store.SaveObject(ctx, config.BucketName(), getLogPrefix(info)+"tidb/a.log", strings.NewReader("test"))
	info2, err = LoadClusterInfo(ctx, store, config.BucketName(), key)
	assert.NoError(t, err)
	assert.Equal(t, info2.SubClusterType, types.SubClusterTypePhysical)
}