package aggregate

import (
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/pingcap-inc/clinic/pkg/ngm/statements/api"
	"github.com/pingcap-inc/clinic/pkg/ngm/statements/load"
)

func TestIAReadSegmentAggregators(t *testing.T) {
	record1 := &load.StmtRecord{
		SumIaReadSegmentCount:          3,
		SumIaRemoteReadSegmentSize:     100,
		SumIaRemoteReadSegmentWaitTime: 0.1,
		ExecCount:                      2,
		MaxIaReadSegmentCount:          2,
		MaxIaRemoteReadSegmentSize:     80,
		MaxIaRemoteReadSegmentWaitTime: 0.05,
	}
	record2 := &load.StmtRecord{
		SumIaReadSegmentCount:          5,
		SumIaRemoteReadSegmentSize:     200,
		SumIaRemoteReadSegmentWaitTime: 0.4,
		ExecCount:                      4,
		MaxIaReadSegmentCount:          6,
		MaxIaRemoteReadSegmentSize:     300,
		MaxIaRemoteReadSegmentWaitTime: 0.7,
	}

	aggs := Aggregators(
		AggSumIaReadSegmentCount(),
		AggSumIaRemoteReadSegmentSize(),
		AggSumIaRemoteReadSegmentWaitTime(),
		AggAvgIaReadSegmentCount(),
		AggAvgIaRemoteReadSegmentSize(),
		AggAvgIaRemoteReadSegmentWaitTime(),
		AggMaxIaReadSegmentCount(),
		AggMaxIaRemoteReadSegmentSize(),
		AggMaxIaRemoteReadSegmentWaitTime(),
	)
	aggs.Add(record1)
	aggs.Add(record2)

	model := &api.Model{}
	aggs.Output(model)
	require.Equal(t, int64(8), model.SumIaReadSegmentCount)
	require.Equal(t, int64(300), model.SumIaRemoteReadSegmentSize)
	require.InDelta(t, 0.5, model.SumIaRemoteReadSegmentWaitTime, 1e-9)
	require.InDelta(t, float64(8)/float64(6), model.AvgIaReadSegmentCount, 1e-9)
	require.InDelta(t, float64(300)/float64(6), model.AvgIaRemoteReadSegmentSize, 1e-9)
	require.InDelta(t, 0.5/float64(6), model.AvgIaRemoteReadSegmentWaitTime, 1e-9)
	require.Equal(t, int64(6), model.MaxIaReadSegmentCount)
	require.Equal(t, int64(300), model.MaxIaRemoteReadSegmentSize)
	require.InDelta(t, 0.7, model.MaxIaRemoteReadSegmentWaitTime, 1e-9)
}

func TestIAReadSegmentAggregatorsRegistered(t *testing.T) {
	require.Contains(t, StatementAggregatorFactories, "sum_ia_read_segment_count")
	require.Contains(t, StatementAggregatorFactories, "sum_ia_remote_read_segment_size")
	require.Contains(t, StatementAggregatorFactories, "sum_ia_remote_read_segment_wait_time")
	require.Contains(t, StatementAggregatorFactories, "avg_ia_read_segment_count")
	require.Contains(t, StatementAggregatorFactories, "avg_ia_remote_read_segment_size")
	require.Contains(t, StatementAggregatorFactories, "avg_ia_remote_read_segment_wait_time")
	require.Contains(t, StatementAggregatorFactories, "max_ia_read_segment_count")
	require.Contains(t, StatementAggregatorFactories, "max_ia_remote_read_segment_size")
	require.Contains(t, StatementAggregatorFactories, "max_ia_remote_read_segment_wait_time")
}