package slowquery

import (
	"context"
	"fmt"
	"io"
	"strings"
	"testing"
	"time"

	"github.com/pingcap/errors"
	"github.com/stretchr/testify/require"

	"github.com/pingcap-inc/clinic/pkg/ngm/storage"
)

var testS3Data = map[string]string{
	// 2022-08-31T07:13:00Z
	"1661929980-9184af9a-1a13-47de-80f3-50767494d7d4.log.gz": `2022-08-31T07:12:54.233489375Z;ip-10-250-24-4.us-west-2.compute.internal;tidb61992;db-tidb-0;slowlog # Time: 2022-08-31T07:12:54.096687178Z
# Txn_start_ts: 0
# User@Host: root[root] @ 10.250.24.175 [10.250.24.175]
# Conn_ID: 6204412174144681481
# Query_time: 1.000883579
# Parse_time: 0.000046967
# Compile_time: 0.000166604
# Rewrite_time: 0.000063231
# Optimize_time: 0.000029576
# Wait_TS: 0
# Is_internal: false
# Digest: 926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb
# Num_cop_tasks: 0
# Prepared: false
# Plan_from_cache: false
# Plan_from_binding: false
# Has_more_results: false
# KV_total: 0
# PD_total: 0.000006361
# Backoff_total: 0
# Write_sql_response_total: 0.000000887
# Result_rows: 1
# Succ: true
# IsExplicitTxn: false
# Plan: tidb_decode_plan('ffBhMAkzXzMJMAkxCXNsZWVwKDEpLT5Db2x1bW4jMQkxCXRpbWU6MXMsIGxvb3BzOjIsIENvbmN1cnJlbmN5Ok9GRgkwIEJ5dGVzCU4vQQoxCTI1XzQJMAkxCXJvd3M6MQkxCXQBQgwyOThuGUUgCU4vQQlOL0EK')
# Plan_digest: a55eaee1d30304ad05500afbbfff1409ce22376e0f062c0d06ff34b8ed9c2b47
select sleep ( ? ) ;`,

	// 2022-08-31T07:14:00Z
	"1661930040-1619e367-00bc-403e-8856-169b816e8391.log.gz": `2022-08-31T07:13:28.236100388Z;ip-10-250-24-4.us-west-2.compute.internal;tidb61992;db-tidb-0;slowlog # Time: 2022-08-31T07:13:27.899863054Z
# Txn_start_ts: 0
# User@Host: root[root] @ 10.250.24.175 [10.250.24.175]
# Conn_ID: 6204412174144681481
# Query_time: 7.049241269
# Parse_time: 0.000062393
# Compile_time: 0.000217623
# Rewrite_time: 0.000103079
# Optimize_time: 0.000036333
# Wait_TS: 0
# Is_internal: false
# Digest: a1e5b2baf44eb501800e67b25555cf6722c416c1a49b8f500a27c5a32e8b231f
# Num_cop_tasks: 0
# Mem_max: 26316
# Prepared: false
# Plan_from_cache: false
# Plan_from_binding: false
# Has_more_results: false
# KV_total: 0
# PD_total: 0.000008436
# Backoff_total: 0
# Write_sql_response_total: 0.000001168
# Result_rows: 7
# Succ: true
# IsExplicitTxn: false
# Plan: tidb_decode_plan('kwHwWDAJM18zCTAJMTAwMDAJc2xlZXAoMSktPkNvbHVtbiM5CTcJdGltZTo3LjA1cywgbG9vcHM6MiwgQ29uY3VycmVuY3k6NQkyNS43IEtCCU4vQQoxCTExXzQJEVNEdGFibGU6Q0xVU1RFUl9JTkZPEVMUNDUuMW1zFVQgCU4vQQlOL0EK')
# Plan_digest: e6b8ba3530968656376af094d5f47d6ce6e73afd930a2f54699bf4ffc793f3af
select sleep(?) from information_schema.cluster_info;
2022-08-31T07:13:49.237672411Z;ip-10-250-24-4.us-west-2.compute.internal;tidb61992;db-tidb-0;slowlog # Time: 2022-08-31T07:13:48.322589791Z
# Txn_start_ts: 0
# User@Host: root[root] @ 10.250.24.175 [10.250.24.175]
# Conn_ID: 6204412174144681481
# Query_time: 1.00096328
# Parse_time: 0.000038883
# Compile_time: 0.000117665
# Rewrite_time: 0.000040561
# Optimize_time: 0.000025832
# Wait_TS: 0
# DB: mysql
# Is_internal: false
# Digest: 926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb
# Num_cop_tasks: 0
# Prepared: false
# Plan_from_cache: false
# Plan_from_binding: false
# Has_more_results: false
# KV_total: 0
# PD_total: 0.000006872
# Backoff_total: 0
# Write_sql_response_total: 0.000000831
# Result_rows: 1
# Succ: true
# IsExplicitTxn: false
# Plan: tidb_decode_plan('ffBhMAkzXzMJMAkxCXNsZWVwKDEpLT5Db2x1bW4jMQkxCXRpbWU6MXMsIGxvb3BzOjIsIENvbmN1cnJlbmN5Ok9GRgkwIEJ5dGVzCU4vQQoxCTI1XzQJMAkxCXJvd3M6MQkxCXQBQgw0MjZuGUUgCU4vQQlOL0EK')
# Plan_digest: a55eaee1d30304ad05500afbbfff1409ce22376e0f062c0d06ff34b8ed9c2b47
use mysql;
select sleep(?);
2022-08-31T07:13:54.238055795Z;ip-10-250-24-4.us-west-2.compute.internal;tidb61992;db-tidb-0;slowlog # Time: 2022-08-31T07:13:54.022015965Z
# Txn_start_ts: 0
# User@Host: root[root] @ 10.250.24.175 [10.250.24.175]
# Conn_ID: 6204412174144681481
# Query_time: 1.000614937
# Parse_time: 0.000048639
# Compile_time: 0.000167435
# Rewrite_time: 0.000047539
# Optimize_time: 0.000028722
# Wait_TS: 0
# DB: mysql
# Is_internal: false
# Digest: 926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb
# Num_cop_tasks: 0
# Prepared: false
# Plan_from_cache: false
# Plan_from_binding: false
# Has_more_results: false
# KV_total: 0
# PD_total: 0.000013826
# Backoff_total: 0
# Write_sql_response_total: 0.000000943
# Result_rows: 1
# Succ: true
# IsExplicitTxn: false
# Plan: tidb_decode_plan('ffBhMAkzXzMJMAkxCXNsZWVwKDEpLT5Db2x1bW4jMQkxCXRpbWU6MXMsIGxvb3BzOjIsIENvbmN1cnJlbmN5Ok9GRgkwIEJ5dGVzCU4vQQoxCTI1XzQJMAkxCXJvd3M6MQkxCXQBQgw0NjhuGUUgCU4vQQlOL0EK')
# Plan_digest: a55eaee1d30304ad05500afbbfff1409ce22376e0f062c0d06ff34b8ed9c2b47
select sleep(?);
2022-08-31T07:13:57.238285034Z;ip-10-250-24-4.us-west-2.compute.internal;tidb61992;db-tidb-0;slowlog # Time: 2022-08-31T07:13:56.26435671Z
# Txn_start_ts: 0
# User@Host: root[root] @ 10.250.24.175 [10.250.24.175]
# Conn_ID: 6204412174144681481
# Query_time: 1.000911511
# Parse_time: 0.00005975
# Compile_time: 0.000145095
# Rewrite_time: 0.000047735
# Optimize_time: 0.000023172
# Wait_TS: 0
# DB: mysql
# Is_internal: false
# Digest: 926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb
# Num_cop_tasks: 0
# Prepared: false
# Plan_from_cache: false
# Plan_from_binding: false
# Has_more_results: false
# KV_total: 0
# PD_total: 0.000007617
# Backoff_total: 0
# Write_sql_response_total: 0.00000082
# Result_rows: 1
# Succ: true
# IsExplicitTxn: false
# Plan: tidb_decode_plan('ffBhMAkzXzMJMAkxCXNsZWVwKDEpLT5Db2x1bW4jMQkxCXRpbWU6MXMsIGxvb3BzOjIsIENvbmN1cnJlbmN5Ok9GRgkwIEJ5dGVzCU4vQQoxCTI1XzQJMAkxCXJvd3M6MQkxCXQBQgw0NzluGUUgCU4vQQlOL0EK')
# Plan_digest: a55eaee1d30304ad05500afbbfff1409ce22376e0f062c0d06ff34b8ed9c2b47
select sleep(?);`,
}

func Test_s3SlowQuery_List(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 5)
}

func Test_s3SlowQuery_List_Partial(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661929979,
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 1)

	ms.getObjectCount = 0
	logs, err = sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 1661929981,
		EndTime:   1661930040,
		IsDesc:    true,
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 4)

	ms.getObjectCount = 0
	tm, err := time.Parse(time.RFC3339Nano, "2022-08-31T07:13:55.000000000Z")
	require.NoError(t, err)
	logs, err = sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: int(tm.Unix()),
		EndTime:   1661930040,
		IsDesc:    true,
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 1)
}

func Test_s3SlowQuery_List_Empty(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 1662452177,
		EndTime:   9999999999,
	})
	require.NoError(t, err)
	require.Nil(t, logs)
	require.Equal(t, 0, ms.getObjectCount)
}

func Test_s3SlowQuery_List_DB(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		DB:        []string{"mysql"},
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 3)
}

func Test_s3SlowQuery_List_Limit(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		Limit:     2,
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 2)
}

func Test_s3SlowQuery_List_Limit2(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		Limit:     2,
		IsDesc:    true,
	})
	require.NoError(t, err)
	require.Equal(t, 1, ms.getObjectCount)
	require.Len(t, logs, 2)
}

func Test_s3SlowQuery_List_Text(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		Text:      "information_schema.cluster_info",
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 1)
}

func Test_s3SlowQuery_List_OrderByTime(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 5)
	require.True(t, logs[0].Time.Unix() <= logs[1].Time.Unix())
	require.True(t, logs[1].Time.Unix() <= logs[2].Time.Unix())
	require.True(t, logs[2].Time.Unix() <= logs[3].Time.Unix())
	require.True(t, logs[3].Time.Unix() <= logs[4].Time.Unix())
}

func Test_s3SlowQuery_List_OrderByTimeDesc(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		IsDesc:    true,
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 5)
	require.True(t, logs[0].Time.Unix() >= logs[1].Time.Unix())
	require.True(t, logs[1].Time.Unix() >= logs[2].Time.Unix())
	require.True(t, logs[2].Time.Unix() >= logs[3].Time.Unix())
	require.True(t, logs[3].Time.Unix() >= logs[4].Time.Unix())
}

func Test_s3SlowQuery_List_OrderByQueryTimeDesc(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		OrderBy:   "query_time",
		IsDesc:    true,
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 5)
	require.True(t, logs[0].QueryTime >= logs[1].QueryTime)
	require.True(t, logs[1].QueryTime >= logs[2].QueryTime)
	require.True(t, logs[2].QueryTime >= logs[3].QueryTime)
	require.True(t, logs[3].QueryTime >= logs[4].QueryTime)
}

func Test_s3SlowQuery_List_OrderByMemMax(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		OrderBy:   "query_time",
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 5)
	require.True(t, logs[0].MemMax <= logs[1].MemMax)
	require.True(t, logs[1].MemMax <= logs[2].MemMax)
	require.True(t, logs[2].MemMax <= logs[3].MemMax)
	require.True(t, logs[3].MemMax <= logs[4].MemMax)
}

func Test_s3SlowQuery_List_Plans(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		Plans:     []string{"a55eaee1d30304ad05500afbbfff1409ce22376e0f062c0d06ff34b8ed9c2b47"},
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 4)
}

func Test_s3SlowQuery_List_Digest(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		Digest:    "926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb",
	})
	require.NoError(t, err)
	require.Equal(t, 2, ms.getObjectCount)
	require.Len(t, logs, 4)
}

func Test_s3SlowQuery_List_DigestLimit(t *testing.T) {
	ms := &mockStorage{data: testS3Data}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}
	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		Digest:    "926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb",
		Limit:     1,
	})
	require.NoError(t, err)
	require.Equal(t, 1, ms.getObjectCount)
	require.Len(t, logs, 1)
}

func Test_s3SlowQuery_List_FiltersInternalFlagWithUser(t *testing.T) {
	ms := &mockStorage{data: map[string]string{
		"1661930040-internal.log.gz": strings.Join([]string{
			slowQueryListTestEntry("2022-08-31T07:13:27.899863054Z", "external_digest", false, "select external;"),
			slowQueryListTestEntry("2022-08-31T07:13:54.022015965Z", "internal_digest", true, "select internal;"),
		}, "\n"),
	}}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}

	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		IsDesc:    true,
	})
	require.NoError(t, err)
	require.Len(t, logs, 1)
	require.Equal(t, "external_digest", logs[0].Digest)

	ms.getObjectCount = 0
	logs, err = sq.List(context.Background(), &ListParams{
		Provider:     "aws",
		Region:       "us-west-2",
		BeginTime:    0,
		EndTime:      1661930040,
		IsDesc:       true,
		ShowInternal: true,
	})
	require.NoError(t, err)
	require.Len(t, logs, 2)
	require.Equal(t, "internal_digest", logs[0].Digest)
	require.Equal(t, "external_digest", logs[1].Digest)
}

func Test_s3SlowQuery_List_OrderByFiltersInternalFlagWithUser(t *testing.T) {
	ms := &mockStorage{data: map[string]string{
		"1661930040-internal.log.gz": strings.Join([]string{
			slowQueryListTestEntry("2022-08-31T07:13:27.899863054Z", "external_digest", false, "select external;"),
			slowQueryListTestEntry("2022-08-31T07:13:54.022015965Z", "internal_digest", true, "select internal;"),
		}, "\n"),
	}}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}

	logs, err := sq.List(context.Background(), &ListParams{
		Provider:  "aws",
		Region:    "us-west-2",
		BeginTime: 0,
		EndTime:   1661930040,
		OrderBy:   "query_time",
		IsDesc:    true,
	})
	require.NoError(t, err)
	require.Len(t, logs, 1)
	require.Equal(t, "external_digest", logs[0].Digest)
}

func Test_s3SlowQuery_Detail_ForwardRetry(t *testing.T) {
	const (
		digest = "926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb"
		connID = uint64(6204412174144681481)
	)
	tm, err := time.Parse(time.RFC3339Nano, "2022-08-31T07:13:54.022015965Z")
	require.NoError(t, err)
	futureTm := tm.Add(10 * time.Minute)
	timestamp := timeToFloat(tm)
	keyTs := int(timestamp) + detailLookupForwardRetryStep*5
	ms := &mockStorage{data: map[string]string{
		fmt.Sprintf("%d-delayed.log.gz", keyTs): strings.Join([]string{
			slowQueryDetailTestEntry(futureTm, "a1e5b2baf44eb501800e67b25555cf6722c416c1a49b8f500a27c5a32e8b231f", connID+1),
			slowQueryDetailTestEntry(tm, digest, connID),
		}, "\n"),
	}}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}

	sl, err := sq.Detail(context.Background(), &DetailParams{
		Provider:  "aws",
		Region:    "us-west-2",
		Timestamp: timestamp,
		ConnID:    connID,
		Digest:    digest,
	})
	require.NoError(t, err)
	require.Equal(t, digest, sl.Digest)
	require.Equal(t, connID, sl.ConnID)
	require.Equal(t, timestamp, timeToFloat(sl.Time))
	require.Equal(t, 1, ms.getObjectCount)
}

func Test_s3SlowQuery_Detail_ForwardRetryLimit(t *testing.T) {
	const (
		digest = "926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb"
		connID = uint64(6204412174144681481)
	)
	tm, err := time.Parse(time.RFC3339Nano, "2022-08-31T07:13:54.022015965Z")
	require.NoError(t, err)
	timestamp := timeToFloat(tm)
	keyTs := int(timestamp) + detailLookupForwardRetryStep*(detailLookupMaxForwardRetries+2)
	ms := &mockStorage{data: map[string]string{
		fmt.Sprintf("%d-too-late.log.gz", keyTs): slowQueryDetailTestEntry(tm, digest, connID),
	}}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}

	_, err = sq.Detail(context.Background(), &DetailParams{
		Provider:  "aws",
		Region:    "us-west-2",
		Timestamp: timestamp,
		ConnID:    connID,
		Digest:    digest,
	})
	require.EqualError(t, err, errSlowQueryDetailNotFound.Error())
	require.Equal(t, 0, ms.getObjectCount)
}

func Test_s3SlowQuery_Detail_ForwardRetryDeduplicatesKeys(t *testing.T) {
	const (
		requestDigest = "926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb"
		actualDigest  = "a1e5b2baf44eb501800e67b25555cf6722c416c1a49b8f500a27c5a32e8b231f"
		connID        = uint64(6204412174144681481)
	)
	tm, err := time.Parse(time.RFC3339Nano, "2022-08-31T07:13:54.022015965Z")
	require.NoError(t, err)
	timestamp := timeToFloat(tm)
	ms := &mockStorage{data: map[string]string{
		fmt.Sprintf("%d-overlap.log.gz", int(timestamp)): slowQueryDetailTestEntry(tm, actualDigest, connID),
	}}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}

	_, err = sq.Detail(context.Background(), &DetailParams{
		Provider:  "aws",
		Region:    "us-west-2",
		Timestamp: timestamp,
		ConnID:    connID,
		Digest:    requestDigest,
	})
	require.EqualError(t, err, errSlowQueryDetailNotFound.Error())
	require.Equal(t, 1, ms.getObjectCount)
}

func Test_s3SlowQuery_Detail_ForwardRetryStopsOnObjectError(t *testing.T) {
	const (
		digest = "926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb"
		connID = uint64(6204412174144681481)
	)
	tm, err := time.Parse(time.RFC3339Nano, "2022-08-31T07:13:54.022015965Z")
	require.NoError(t, err)
	timestamp := timeToFloat(tm)
	ms := &mockStorage{
		data: map[string]string{
			fmt.Sprintf("%d-error.log.gz", int(timestamp)): slowQueryDetailTestEntry(tm, digest, connID),
		},
		objectErr: errors.New("boom"),
	}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}

	_, err = sq.Detail(context.Background(), &DetailParams{
		Provider:  "aws",
		Region:    "us-west-2",
		Timestamp: timestamp,
		ConnID:    connID,
		Digest:    digest,
	})
	require.EqualError(t, err, "boom")
	require.Equal(t, 1, ms.getObjectCount)
}

func Test_s3SlowQuery_Detail_FiltersInternalFlagWithUser(t *testing.T) {
	const (
		digest = "internal_digest"
		connID = uint64(6204412174144681481)
	)
	tm, err := time.Parse(time.RFC3339Nano, "2022-08-31T07:13:54.022015965Z")
	require.NoError(t, err)
	timestamp := timeToFloat(tm)
	ms := &mockStorage{data: map[string]string{
		fmt.Sprintf("%d-internal.log.gz", int(timestamp)): slowQueryDetailTestEntryWithInternal(tm, digest, connID, true),
	}}
	sq := &slowQuery{stores: newMockStorageFactory(ms)}

	_, err = sq.Detail(context.Background(), &DetailParams{
		Provider:  "aws",
		Region:    "us-west-2",
		Timestamp: timestamp,
		ConnID:    connID,
		Digest:    digest,
	})
	require.EqualError(t, err, errSlowQueryDetailNotFound.Error())

	sl, err := sq.Detail(context.Background(), &DetailParams{
		Provider:     "aws",
		Region:       "us-west-2",
		Timestamp:    timestamp,
		ConnID:       connID,
		Digest:       digest,
		ShowInternal: true,
	})
	require.NoError(t, err)
	require.Equal(t, digest, sl.Digest)
	require.True(t, sl.IsInternal)
}

func slowQueryListTestEntry(tm, digest string, internal bool, query string) string {
	return fmt.Sprintf(`# Time: %s
# Txn_start_ts: 0
# User@Host: root[root] @ 127.0.0.1 [127.0.0.1]
# Conn_ID: 6204412174144681481
# Query_time: 1.000000
# Is_internal: %t
# Digest: %s
%s`, tm, internal, digest, query)
}

func slowQueryDetailTestEntry(tm time.Time, digest string, connID uint64) string {
	return slowQueryDetailTestEntryWithInternal(tm, digest, connID, false)
}

func slowQueryDetailTestEntryWithInternal(tm time.Time, digest string, connID uint64, internal bool) string {
	return fmt.Sprintf(`# Time: %s
# Txn_start_ts: 0
# User@Host: root[root] @ 127.0.0.1 [127.0.0.1]
# Conn_ID: %d
# Query_time: 1.000000
# Is_internal: %t
# Digest: %s
select 1;`, tm.Format(time.RFC3339Nano), connID, internal, digest)
}

type mockStorageFactory struct {
	store *mockStorage
}

func (m *mockStorageFactory) Storage(ctx context.Context, provider, region string) (storage.Storage, error) {
	return m.store, nil
}

func newMockStorageFactory(store *mockStorage) storage.StorageFactory {
	return &mockStorageFactory{
		store: store,
	}
}

type mockStorage struct {
	data           map[string]string
	getObjectCount int
	objectErr      error
}

func (m *mockStorage) Object(ctx context.Context, bucket, key string) (*storage.Object, error) {
	m.getObjectCount++
	if m.objectErr != nil {
		return nil, m.objectErr
	}
	v, ok := m.data[key]
	if !ok {
		return nil, errors.New("not found")
	}
	return &storage.Object{
		Reader: io.NopCloser(strings.NewReader(v)),
	}, nil
}

func (m *mockStorage) Objects(ctx context.Context, bucket, prefix string) ([]storage.ObjectMeta, error) {
	var metas []storage.ObjectMeta
	for key := range m.data {
		metas = append(metas, storage.ObjectMeta{
			Key:  key,
			Size: int64(len(m.data[key])),
		})
	}
	return metas, nil
}

func (m *mockStorage) ObjectsCallback(
	ctx context.Context,
	bucket, prefix string,
	callback func(storage.ObjectMeta) bool,
) error {
	panic("unreachable")
}

func (m *mockStorage) ObjectsCallbackWithStartKey(
	ctx context.Context,
	bucket, prefix, startKey string,
	callback func(storage.ObjectMeta) bool,
) error {
	panic("unreachable")
}

func TestRangeTsPrefix(t *testing.T) {
	require.Equal(t, "16696312", rangeTsPrefix(1669631230, 1669631231))
	require.Equal(t, "1669631", rangeTsPrefix(1669631131, 1669631231))
	require.Equal(t, "166963", rangeTsPrefix(1669630231, 1669631231))
	require.Equal(t, "16696", rangeTsPrefix(1669621231, 1669631231))
	require.Equal(t, "1669", rangeTsPrefix(1669531231, 1669631231))
	require.Equal(t, "166", rangeTsPrefix(1668631231, 1669631231))
	require.Equal(t, "16", rangeTsPrefix(1659631231, 1669631231))
	require.Equal(t, "1", rangeTsPrefix(1234567890, 1999999999))
	require.Equal(t, "", rangeTsPrefix(1, 1999999999))
	require.Equal(t, "", rangeTsPrefix(0, 1999999999))
	require.Equal(t, "", rangeTsPrefix(-1, 1999999999))
	require.Equal(t, "", rangeTsPrefix(1, 0))
	require.Equal(t, "", rangeTsPrefix(1, -1))
}

func TestSingleTsPrefix(t *testing.T) {
	require.Equal(t, "166963", singleTsPrefix(1669631231))
	require.Equal(t, "", singleTsPrefix(0))
	require.Equal(t, "", singleTsPrefix(-1))
}

func TestStorageInfoLocationOverridesDefaultPath(t *testing.T) {
	require.Equal(t, "bucket-from-o11y", slowQueryBucket("prod", "aws", "us-east-1", true, false, "bucket-from-o11y"))
	require.Equal(t,
		"path/from/o11y/",
		slowQueryKeyPrefix("prod", "aws", "us-east-1", "7777", "8888", "9999", true, false, "path/from/o11y"),
	)
}