// Copyright 2023 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package stmtsummary

import (
	"strings"
	"testing"
	"time"

	"github.com/pingcap/tidb/pkg/meta/model"
	pmodel "github.com/pingcap/tidb/pkg/parser/model"
	"github.com/pingcap/tidb/pkg/parser/mysql"
	"github.com/pingcap/tidb/pkg/types"
	"github.com/pingcap/tidb/pkg/util/chunk"
	"github.com/stretchr/testify/require"
)

func TestColumn(t *testing.T) {
	columns := []*model.ColumnInfo{
		{Name: pmodel.NewCIStr(ClusterTableInstanceColumnNameStr)},
		{Name: pmodel.NewCIStr(StmtTypeStr)},
		{Name: pmodel.NewCIStr(SchemaNameStr)},
		{Name: pmodel.NewCIStr(DigestStr)},
		{Name: pmodel.NewCIStr(DigestTextStr)},
		{Name: pmodel.NewCIStr(TableNamesStr)},
		{Name: pmodel.NewCIStr(IndexNamesStr)},
		{Name: pmodel.NewCIStr(SampleUserStr)},
		{Name: pmodel.NewCIStr(ExecCountStr)},
		{Name: pmodel.NewCIStr(SumLatencyStr)},
		{Name: pmodel.NewCIStr(MaxLatencyStr)},
		{Name: pmodel.NewCIStr(AvgIAReadSegmentCountStr)},
		{Name: pmodel.NewCIStr(MaxIAReadSegmentCountStr)},
		{Name: pmodel.NewCIStr(AvgIARemoteReadSegmentSizeStr)},
		{Name: pmodel.NewCIStr(MaxIARemoteReadSegmentSizeStr)},
		{Name: pmodel.NewCIStr(AvgIARemoteReadSegmentWaitTimeStr)},
		{Name: pmodel.NewCIStr(MaxIARemoteReadSegmentWaitTimeStr)},
		{Name: pmodel.NewCIStr(AvgTidbCPUTimeStr)},
		{Name: pmodel.NewCIStr(AvgTikvCPUTimeStr)},
	}
	factories := makeColumnFactories(columns)
	info := GenerateStmtExecInfo4Test("digest")
	record := NewStmtRecord(info)
	record.Add(info)
	for n, f := range factories {
		column := f(mockColumnInfo{}, record)
		switch columns[n].Name.O {
		case ClusterTableInstanceColumnNameStr:
			require.Equal(t, "instance_addr", column)
		case StmtTypeStr:
			require.Equal(t, record.StmtType, column)
		case SchemaNameStr:
			require.Equal(t, record.SchemaName, column)
		case DigestStr:
			require.Equal(t, record.Digest, column)
		case DigestTextStr:
			require.Equal(t, record.NormalizedSQL, column)
		case TableNamesStr:
			require.Equal(t, record.TableNames, column)
		case IndexNamesStr:
			require.Equal(t, strings.Join(record.IndexNames, ","), column)
		case SampleUserStr:
			require.Equal(t, info.User, column)
		case ExecCountStr:
			require.Equal(t, int64(1), column)
		case SumLatencyStr:
			require.Equal(t, int64(record.SumLatency), column)
		case MaxLatencyStr:
			require.Equal(t, int64(record.MaxLatency), column)
		case AvgIAReadSegmentCountStr:
			require.Equal(t, float64(record.SumIAReadSegmentCount), column)
		case MaxIAReadSegmentCountStr:
			require.Equal(t, record.MaxIAReadSegmentCount, column)
		case AvgIARemoteReadSegmentSizeStr:
			require.Equal(t, float64(record.SumIARemoteReadSegmentSize), column)
		case MaxIARemoteReadSegmentSizeStr:
			require.Equal(t, record.MaxIARemoteReadSegmentSize, column)
		case AvgIARemoteReadSegmentWaitTimeStr:
			require.Equal(t, int64(record.SumIARemoteReadSegmentWaitTime), column)
		case MaxIARemoteReadSegmentWaitTimeStr:
			require.Equal(t, int64(record.MaxIARemoteReadSegmentWaitTime), column)
		case AvgTidbCPUTimeStr:
			require.Equal(t, int64(record.SumTidbCPU), column)
		case AvgTikvCPUTimeStr:
			require.Equal(t, int64(record.SumTikvCPU), column)
		}
	}
}

func TestIAAvgColumns(t *testing.T) {
	columns := []*model.ColumnInfo{
		{Name: pmodel.NewCIStr(AvgIAReadSegmentCountStr)},
		{Name: pmodel.NewCIStr(MaxIAReadSegmentCountStr)},
		{Name: pmodel.NewCIStr(AvgIARemoteReadSegmentSizeStr)},
		{Name: pmodel.NewCIStr(MaxIARemoteReadSegmentSizeStr)},
		{Name: pmodel.NewCIStr(AvgIARemoteReadSegmentWaitTimeStr)},
		{Name: pmodel.NewCIStr(MaxIARemoteReadSegmentWaitTimeStr)},
	}
	factories := makeColumnFactories(columns)

	info1 := GenerateStmtExecInfo4Test("digest")
	info2 := GenerateStmtExecInfo4Test("digest")
	info2.ExecDetail.ScanDetail.IaRemoteReadSegmentCount = 5
	info2.ExecDetail.ScanDetail.IaRemoteReadSegmentBytes = 8192
	info2.ExecDetail.ScanDetail.IaRemoteReadSegmentDuration = 9 * time.Millisecond

	record := NewStmtRecord(info1)
	record.Add(info1)
	record.Add(info2)

	require.Equal(t, 4.0, factories[0](mockColumnInfo{}, record))
	require.Equal(t, uint64(5), factories[1](mockColumnInfo{}, record))
	require.Equal(t, 6144.0, factories[2](mockColumnInfo{}, record))
	require.Equal(t, uint64(8192), factories[3](mockColumnInfo{}, record))
	require.Equal(t, int64(7*time.Millisecond), factories[4](mockColumnInfo{}, record))
	require.Equal(t, int64(9*time.Millisecond), factories[5](mockColumnInfo{}, record))
}

func TestIAAvgColumnsChunkRoundTrip(t *testing.T) {
	columns := []*model.ColumnInfo{
		{Name: pmodel.NewCIStr(AvgIAReadSegmentCountStr)},
		{Name: pmodel.NewCIStr(MaxIAReadSegmentCountStr)},
		{Name: pmodel.NewCIStr(AvgIARemoteReadSegmentSizeStr)},
		{Name: pmodel.NewCIStr(MaxIARemoteReadSegmentSizeStr)},
		{Name: pmodel.NewCIStr(AvgIARemoteReadSegmentWaitTimeStr)},
		{Name: pmodel.NewCIStr(MaxIARemoteReadSegmentWaitTimeStr)},
	}
	factories := makeColumnFactories(columns)

	info1 := GenerateStmtExecInfo4Test("digest")
	info1.ExecDetail.ScanDetail.IaRemoteReadSegmentCount = 3
	info1.ExecDetail.ScanDetail.IaRemoteReadSegmentBytes = 4096
	info1.ExecDetail.ScanDetail.IaRemoteReadSegmentDuration = 5 * time.Millisecond

	info2 := GenerateStmtExecInfo4Test("digest")
	info2.ExecDetail.ScanDetail.IaRemoteReadSegmentCount = 5
	info2.ExecDetail.ScanDetail.IaRemoteReadSegmentBytes = 8192
	info2.ExecDetail.ScanDetail.IaRemoteReadSegmentDuration = 9 * time.Millisecond

	record := NewStmtRecord(info1)
	record.Add(info1)
	record.Add(info2)

	rowDatums := make([]types.Datum, len(factories))
	for i, factory := range factories {
		rowDatums[i] = types.NewDatum(factory(mockColumnInfo{}, record))
	}

	maxUnsignedType := types.NewFieldType(mysql.TypeLonglong)
	maxUnsignedType.SetFlag(mysql.UnsignedFlag)
	retTypes := []*types.FieldType{
		types.NewFieldType(mysql.TypeDouble),
		maxUnsignedType,
		types.NewFieldType(mysql.TypeDouble),
		maxUnsignedType.Clone(),
		types.NewFieldType(mysql.TypeLonglong),
		types.NewFieldType(mysql.TypeLonglong),
	}
	mutRow := chunk.MutRowFromTypes(retTypes)
	mutRow.SetDatums(rowDatums...)
	row := mutRow.ToRow()

	require.Equal(t, 4.0, row.GetFloat64(0))
	require.Equal(t, uint64(5), row.GetUint64(1))
	require.Equal(t, 6144.0, row.GetFloat64(2))
	require.Equal(t, uint64(8192), row.GetUint64(3))
	require.Equal(t, int64(7*time.Millisecond), row.GetInt64(4))
	require.Equal(t, int64(9*time.Millisecond), row.GetInt64(5))
}

type mockColumnInfo struct{}

func (mockColumnInfo) getInstanceAddr() string {
	return "instance_addr"
}

func (mockColumnInfo) getTimeLocation() *time.Location {
	loc, _ := time.LoadLocation("Asia/Shanghai")
	return loc
}