package aggregate

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

type Aggregator interface {
	Add(*load.StmtRecord)
	Merge(Aggregator)
	Output(*api.Model)
}

type AggregatorFactory func() Aggregator

var StatementAggregatorFactories = map[string]AggregatorFactory{
	"summary_begin_time":                   AggSummaryBeginTime,
	"summary_end_time":                     AggSummaryEndTime,
	"digest_text":                          AggDigestText,
	"digest":                               AggDigest,
	"exec_count":                           AggExecCount,
	"stmt_type":                            AggStmtType,
	"sum_errors":                           AggSumErrors,
	"sum_warnings":                         AggSumWarnings,
	"sum_latency":                          AggSumLatency,
	"max_latency":                          AggMaxLatency,
	"min_latency":                          AggMinLatency,
	"avg_latency":                          AggAvgLatency,
	"avg_parse_latency":                    AggAvgParseLatency,
	"max_parse_latency":                    AggMaxParseLatency,
	"avg_compile_latency":                  AggAvgCompileLatency,
	"max_compile_latency":                  AggMaxCompileLatency,
	"sum_cop_task_num":                     AggSumCopTaskNum,
	"max_cop_process_time":                 AggMaxCopProcessTime,
	"max_cop_wait_time":                    AggMaxCopWaitTime,
	"avg_process_time":                     AggAvgProcessTime,
	"max_process_time":                     AggMaxProcessTime,
	"avg_wait_time":                        AggAvgWaitTime,
	"max_wait_time":                        AggMaxWaitTime,
	"avg_backoff_time":                     AggAvgBackoffTime,
	"max_backoff_time":                     AggMaxBackoffTime,
	"avg_total_keys":                       AggAvgTotalKeys,
	"max_total_keys":                       AggMaxTotalKeys,
	"avg_processed_keys":                   AggAvgProcessedKeys,
	"max_processed_keys":                   AggMaxProcessedKeys,
	"avg_prewrite_time":                    AggAvgPrewriteTime,
	"max_prewrite_time":                    AggMaxPrewriteTime,
	"avg_commit_time":                      AggAvgCommitTime,
	"max_commit_time":                      AggMaxCommitTime,
	"avg_get_commit_ts_time":               AggAvgGetCommitTsTime,
	"max_get_commit_ts_time":               AggMaxGetCommitTsTime,
	"avg_commit_backoff_time":              AggAvgCommitBackoffTime,
	"max_commit_backoff_time":              AggMaxCommitBackoffTime,
	"avg_resolve_lock_time":                AggAvgResolveLockTime,
	"max_resolve_lock_time":                AggMaxResolveLockTime,
	"avg_local_latch_wait_time":            AggAvgLocalLatchWaitTime,
	"max_local_latch_wait_time":            AggMaxLocalLatchWaitTime,
	"avg_write_keys":                       AggAvgWriteKeys,
	"max_write_keys":                       AggMaxWriteKeys,
	"avg_write_size":                       AggAvgWriteSize,
	"max_write_size":                       AggMaxWriteSize,
	"avg_prewrite_regions":                 AggAvgPrewriteRegions,
	"max_prewrite_regions":                 AggMaxPrewriteRegions,
	"avg_txn_retry":                        AggAvgTxnRetry,
	"max_txn_retry":                        AggMaxTxnRetry,
	"sum_backoff_times":                    AggSumBackoffTimes,
	"avg_mem":                              AggAvgMem,
	"max_mem":                              AggMaxMem,
	"avg_disk":                             AggAvgDisk,
	"max_disk":                             AggMaxDisk,
	"avg_affected_rows":                    AggAvgAffectedRows,
	"first_seen":                           AggFirstSeen,
	"last_seen":                            AggLastSeen,
	"sample_user":                          AggSampleUser,
	"query_sample_text":                    AggQuerySampleText,
	"prev_sample_text":                     AggPrevSampleText,
	"schema_name":                          AggSchemaName,
	"table_names":                          AggTableNames,
	"index_names":                          AggIndexNames,
	"plan_count":                           AggPlanCount,
	"plan":                                 AggPlan,
	"plan_digest":                          AggPlanDigest,
	"binary_plan":                          AggBinaryPlan,
	"plan_hint":                            AggPlanHint,
	"plan_in_cache":                        AggPlanInCache,
	"plan_cache_hits":                      AggPlanCacheHits,
	"plan_in_binding":                      AggPlanInBinding,
	"max_rocksdb_delete_skipped_count":     AggMaxRocksdbDeleteSkippedCount,
	"avg_rocksdb_delete_skipped_count":     AggAvgRocksdbDeleteSkippedCount,
	"max_rocksdb_key_skipped_count":        AggMaxRocksdbKeySkippedCount,
	"avg_rocksdb_key_skipped_count":        AggAvgRocksdbKeySkippedCount,
	"max_rocksdb_block_cache_hit_count":    AggMaxRocksdbBlockCacheHitCount,
	"avg_rocksdb_block_cache_hit_count":    AggAvgRocksdbBlockCacheHitCount,
	"max_rocksdb_block_read_count":         AggMaxRocksdbBlockReadCount,
	"avg_rocksdb_block_read_count":         AggAvgRocksdbBlockReadCount,
	"max_rocksdb_block_read_byte":          AggMaxRocksdbBlockReadByte,
	"avg_rocksdb_block_read_byte":          AggAvgRocksdbBlockReadByte,
	"sum_ru":                               AggSumRu,
	"sum_rru":                              AggSumRru,
	"sum_wru":                              AggSumWru,
	"sum_ru_wait_duration":                 AggSumRuWaitDuration,
	"avg_ru":                               AggAvgRu,
	"avg_rru":                              AggAvgRru,
	"avg_wru":                              AggAvgWru,
	"max_rru":                              AggMaxRru,
	"max_wru":                              AggMaxWru,
	"max_ru_wait_duration":                 AggMaxRuWaitDuration,
	"sum_ru_v2":                            AggSumRuv2,
	"max_ru_v2":                            AggMaxRuv2,
	"avg_ru_v2":                            AggAvgRuv2,
	"sum_ia_read_segment_count":            AggSumIaReadSegmentCount,
	"sum_ia_remote_read_segment_size":      AggSumIaRemoteReadSegmentSize,
	"sum_ia_remote_read_segment_wait_time": AggSumIaRemoteReadSegmentWaitTime,
	"avg_ia_read_segment_count":            AggAvgIaReadSegmentCount,
	"avg_ia_remote_read_segment_size":      AggAvgIaRemoteReadSegmentSize,
	"avg_ia_remote_read_segment_wait_time": AggAvgIaRemoteReadSegmentWaitTime,
	"max_ia_read_segment_count":            AggMaxIaReadSegmentCount,
	"max_ia_remote_read_segment_size":      AggMaxIaRemoteReadSegmentSize,
	"max_ia_remote_read_segment_wait_time": AggMaxIaRemoteReadSegmentWaitTime,
}

func AggregatorFactories(fs []AggregatorFactory) AggregatorFactory {
	return func() Aggregator {
		c := &AggregatorsChain{}
		for _, f := range fs {
			c.aggregators = append(c.aggregators, f())
		}
		return c
	}
}

func Aggregators(aggs ...Aggregator) Aggregator {
	return &AggregatorsChain{aggregators: aggs}
}

type AggregatorsChain struct {
	aggregators []Aggregator
}

func (a *AggregatorsChain) Add(record *load.StmtRecord) {
	for _, agg := range a.aggregators {
		agg.Add(record)
	}
}

func (a *AggregatorsChain) Merge(agg Aggregator) {
	if agg, ok := agg.(*AggregatorsChain); ok {
		if len(a.aggregators) == len(agg.aggregators) {
			for i := range a.aggregators {
				a.aggregators[i].Merge(agg.aggregators[i])
			}
		}
	}
}

func (a *AggregatorsChain) Output(model *api.Model) {
	for _, agg := range a.aggregators {
		agg.Output(model)
	}
}