package controller

import (
	"context"
	"fmt"
	"math"
	"math/rand"
	"sync"
	"time"

	"github.com/prometheus/client_golang/api"
	promv1 "github.com/prometheus/client_golang/api/prometheus/v1"
	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
	o11yv1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1"
	ctrl "sigs.k8s.io/controller-runtime"
)

const (
	vmSelectUrlTemplate = "http://ingress-nginx-controller.kube-system:80/internal/metrics/%s/"
	queryTimeout        = 32 * time.Second // 2s larger than vmselect.maxqueryduration
	defaultWorkers      = 3                // control qps spike
	maxDataPoints       = 2500
	rangeFactor         = 80
)

var (
	queriers map[string]*StressQuerier = map[string]*StressQuerier{}
	lock     sync.Mutex
	logger   = ctrl.Log.WithName("stress_querier")
	// refer to https://github.com/tidbcloud/metrics-server/blob/main/internal/datasource/central_thanos/metrics.go
	integrationQueries = []string{
		`node_memory_MemTotal_bytes{} - on(component, instance) node_memory_MemAvailable_bytes{}`,
		`sum(node_cpu_seconds_total{mode!= "idle"}) by (component, instance)`,
		`tikv_store_size_bytes{instance=~"db-.*",type="used"} or tiflash_system_current_metric_StoreSizeUsed{instance=~"db-.*"}`,
		`sum(pd_cluster_status{type="storage_capacity"})`,
		`tidb_executor_statement_total{instance=~"db-.*"}`,
		`tidb_server_execute_error_total{instance=~"db-.*"}`,
		`tidb_server_connections{instance=~"db-.*"}`,
		`{__name__=~"tidb_server_handle_query_duration_seconds.+",instance=~"db-.*"}`,
		`sum(tidb_server_connections{}) by (component, instance)`,
		`sum(tidb_server_tokens{}) by (component, instance)`,
		`sum(rate(tidb_server_disconnection_total{}[2m])) by (component, instance, result)`,
		`sum(rate(tidb_server_handle_query_duration_seconds_sum{sql_type!="internal"}[2m])) by (component, instance, sql_type)`,
		`sum(rate(tidb_executor_statement_total{}[2m])) by (component, instance, type)`,
		`sum(rate(tidb_server_execute_error_total{}[2m])) by (component, instance, type)`,
		`sum(rate(tidb_server_query_total{}[2m])) by (component, instance, type)`,
		`sum(rate(tidb_server_plan_cache_total{}[2m])) by (component, instance, type)`,
		`sum(rate(tidb_server_handle_query_duration_seconds_sum{sql_type!="internal"}[2m])) by (component, instance, sql_type) / sum(rate(tidb_server_handle_query_duration_seconds_count{sql_type!="internal"}[2m])) by (component, instance, sql_type) > 0`,
		`sum(rate(tidb_session_transaction_duration_seconds_count{}[2m])) by (component, instance, type, txn_mode)`,
	}
	// refer to https://github.com/tidbcloud/runbooks/blob/master/grafana/tidb-cluster-dashboards/overview.json & clinic overview dashboard
	pdGrafanaQueries = []string{
		"pd_tso_role{  dc=\"global\"}",
		"pd_cluster_status{ type=\"storage_capacity\"}",
		"pd_cluster_status{ type=\"storage_size\"}",
		"pd_cluster_status{  type=\"leader_count\"}",
		"sum(pd_cluster_status{  type=\"store_up_count\"})",
		"sum(pd_cluster_status{  type=\"store_disconnected_count\"})",
		"sum(pd_cluster_status{  type=\"store_unhealth_count\"})",
		"sum(pd_cluster_status{  type=\"store_low_space_count\"})",
		"sum(pd_cluster_status{  type=\"store_down_count\"})",
		"sum(pd_cluster_status{  type=\"store_offline_count\"})",
		"sum(pd_cluster_status{  type=\"store_tombstone_count\"})",
		"histogram_quantile(0.99, sum(rate(grpc_server_handling_seconds_bucket{ }[5m])) by (grpc_method, le))",
		"histogram_quantile(0.98, sum(rate(pd_client_request_handle_requests_duration_seconds_bucket{}[2m])) by (type, le))",
		"avg(rate(pd_client_request_handle_requests_duration_seconds_sum{}[2m])) by (type) /  avg(rate(pd_client_request_handle_requests_duration_seconds_count{}[2m])) by (type)",
		"pd_regions_status{ }",
		"sum(pd_regions_status{}) by (instance, type)",
		"pd_hotspot_status{ type=\"hot_write_region_as_leader\"}",
		"pd_hotspot_status{ type=\"hot_read_region_as_leader\"}",
		"sum(delta(pd_scheduler_region_heartbeat{  type=\"report\", status=\"ok\"}[2m])) by (store)",
		"histogram_quantile(0.99, sum(rate(pd_scheduler_region_heartbeat_latency_seconds_bucket{}[5m])) by (store, le))",
		"sum(rate(pd_client_cmd_handle_cmds_duration_seconds_count{ type=\"tso\"}[2m]))",
		"sum(rate(pd_client_request_handle_requests_duration_seconds_count{ type=\"tso\"}[2m]))",
		"histogram_quantile(0.999, sum(rate(pd_client_cmd_handle_cmds_duration_seconds_bucket{ type=\"tso\"}[2m])) by (le))",
		"histogram_quantile(0.99, sum(rate(pd_client_cmd_handle_cmds_duration_seconds_bucket{ type=\"tso\"}[2m])) by (le))",
		"histogram_quantile(0.90, sum(rate(pd_client_cmd_handle_cmds_duration_seconds_bucket{ type=\"tso\"}[2m])) by (le))",
	}

	tidbGrafanaQueries = []string{
		"sum(rate(tidb_executor_statement_total{}[2m])) by (type)",
		"histogram_quantile(0.999, tidb_server_handle_query_duration_seconds_bucket:le:sum:rate{})",
		"histogram_quantile(0.99, tidb_server_handle_query_duration_seconds_bucket:le:sum:rate{})",
		"histogram_quantile(0.95, tidb_server_handle_query_duration_seconds_bucket:le:sum:rate{})",
		"histogram_quantile(0.80, tidb_server_handle_query_duration_seconds_bucket:le:sum:rate{})",
		"rate(tidb_server_query_total{}[2m])",
		"sum(increase(tidb_server_execute_error_total{}[1m])) by (type)",
		"tidb_server_connections{}",
		"sum(tidb_server_connections{})",
		"process_resident_memory_bytes{ component=\"tidb\"}",
		"go_memory_classes_heap_objects_bytes{ component=\"tidb\"} + go_memory_classes_heap_unused_bytes{ component=\"tidb\"}",
		"sum(rate(tidb_session_transaction_duration_seconds_count{}[2m])) by (type)",
		"histogram_quantile(0.99, sum (tidb_session_transaction_duration_seconds_bucket:le_txn_mode:sum:rate) by (le))",
		"histogram_quantile(0.95, sum (tidb_session_transaction_duration_seconds_bucket:le_txn_mode:sum:rate) by (le))",
		"histogram_quantile(0.80, sum (tidb_session_transaction_duration_seconds_bucket:le_txn_mode:sum:rate) by (le))",
		"sum(rate(tidb_tikvclient_txn_cmd_duration_seconds_count{}[2m])) by (type)",
		"histogram_quantile(0.99, tidb_tikvclient_txn_cmd_duration_seconds_bucket:sum:rate)",
		"sum(rate(pd_client_cmd_handle_cmds_duration_seconds_count{ type=\"tso\"}[2m]))",
		"sum(rate(pd_client_request_handle_requests_duration_seconds_count{ type=\"tso\"}[2m]))",
		"sum(rate(tidb_tikvclient_region_err_total{}[2m])) by (type)",
		"sum(rate(tidb_tikvclient_lock_resolver_actions_total{}[2m])) by (type)",
		"sum(rate(tidb_tikvclient_backoff_seconds_count{}[2m])) by (type)",
		`sum(rate(tidb_server_handle_query_duration_seconds_sum{sql_type!="internal"}[2m]))`,
		`sum(rate(tidb_server_handle_query_duration_seconds_sum{sql_type!="internal"}[2m])) by (sql_type)`,
		`sum(rate(tidb_server_handle_query_duration_seconds_sum{sql_type!="internal"}[2m]))`,
		`sum(rate(tidb_session_parse_duration_seconds_sum{sql_type="general"}[2m]))`,
		`sum(rate(tidb_session_compile_duration_seconds_sum{sql_type="general"}[2m]))`,
		`sum(rate(tidb_session_execute_duration_seconds_sum{sql_type="general"}[2m]))`,
		`sum(rate(tidb_tikvclient_request_seconds_sum{store!="0"}[2m])) by (type)`,
		`sum(tidb_server_connections)`,
		`sum(tidb_server_tokens)`,
		`sum(rate(tidb_server_disconnection_total[2m])) by (instance, result)`,
		`sum(rate(tidb_executor_statement_total[2m]))`,
		`sum(rate(tidb_executor_statement_total[2m])) by (type)`,
		`sum(rate(tidb_server_execute_error_total[2m]))`,
		`sum(rate(tidb_server_query_total[2m])) by (type)`,
		`sum(rate(tidb_server_plan_cache_total[2m]))`,
		`sum(rate(tidb_server_plan_cache_miss_total[2m]))`,
		`sum(rate(tidb_server_handle_query_duration_seconds_sum{sql_type!="internal"}[2m])) / sum(rate(tidb_server_handle_query_duration_seconds_count{sql_type!="internal"}[2m]))`,
	}
	tikvGrafanaQueries = []string{
		"sum(tikv_raftstore_region_count{ type=\"leader\"}) by (instance)",
		"delta(tikv_raftstore_region_count{ type=\"leader\"}[2m]) < -10",
		"sum(tikv_raftstore_region_count{ type=\"region\"}) by (instance)",
		"sum(rate(tikv_thread_cpu_seconds_total{ component=\"tikv\"}[2m])) by (instance)",
		"avg(process_resident_memory_bytes{ component=\"tikv\"}) by (instance)",
		"sum(tikv_engine_size_bytes{}) by (instance)",
		"sum(tikv_engine_size_bytes{}) by (type)",
		"sum(rate(tikv_channel_full_total{}[2m])) by (instance, type)",
		"sum(rate(tikv_server_report_failure_msg_total{}[2m])) by (type,instance,store_id)",
		"sum(tikv_scheduler_contex_total{}) by (instance)",
		"sum(rate(tikv_coprocessor_executor_count{}[2m])) by (type)",
		"histogram_quantile(0.99, sum(rate(tikv_coprocessor_request_duration_seconds_bucket{}[2m])) by (le,req))",
		"histogram_quantile(0.95, sum(rate(tikv_coprocessor_request_duration_seconds_bucket{}[2m])) by (le,req))",
		" sum(rate(tikv_coprocessor_request_duration_seconds_sum{ req=\"select\"}[2m])) /  sum(rate(tikv_coprocessor_request_duration_seconds_count{ req=\"select\"}[2m]))",
		" sum(rate(tikv_coprocessor_request_duration_seconds_sum{ req=\"index\"}[2m])) /  sum(rate(tikv_coprocessor_request_duration_seconds_count{ req=\"index\"}[2m]))",
		"sum(rate(tikv_thread_cpu_seconds_total{ name=~\"raftstore_.*\"}[2m])) by (instance)",
		"sum(rate(tikv_thread_cpu_seconds_total{ name=~\"cop_.*\"}[2m])) by (instance)",
	}
	nodeGrafanaQueries = []string{
		"count(node_cpu_seconds_total{ mode=\"user\"}) by (instance)",
		"node_memory_MemTotal_bytes{}",
		"100 - avg by (instance) (irate(node_cpu_seconds_total{ mode=\"idle\"}[2m]) ) * 100",
		"node_load1{}",
		"node_memory_MemAvailable_bytes{}",
		"irate(node_network_receive_bytes_total{ device!=\"lo\"}[2m])",
		"irate(node_network_transmit_bytes_total{ device!=\"lo\"}[2m])",
		"irate(node_netstat_Tcp_RetransSegs{}[2m])",
		"irate(node_disk_reads_completed_total{}[2m]) + irate(node_disk_writes_completed_total{}[2m])",
		"irate(node_disk_read_bytes_total{}[2m]) + irate(node_disk_written_bytes_total{}[2m])",
		"irate(node_disk_io_time_seconds_total{}[2m])",
	}
	// refer to https://github.com/pingcap/tidb-dashboard/blob/master/ui/packages/tidb-dashboard-for-dbaas/src/apps/Monitoring/metricsQueries.ts
	serverConsoleQueries = []string{
		`avg by (cluster_id, instance) ((time() - process_start_time_seconds{component="tidb"}))`,
		`sum(rate(node_cpu_seconds_total{mode!="idle",component="tidb"}[2m])) by (instance)`,
		`avg(node_memory_MemTotal_bytes{component="tidb"} - node_memory_MemAvailable_bytes{component="tidb"}) by (instance)`,
		`avg by (cluster_id, instance) ((time() - process_start_time_seconds{component="tikv"}))`,
		`sum(rate(node_cpu_seconds_total{mode!="idle",component="tikv"}[2m])) by (instance)`,
		`avg(node_memory_MemTotal_bytes{component="tikv"} - node_memory_MemAvailable_bytes{component="tikv"}) by (instance)`,
		`sum(rate(tikv_engine_flow_bytes{db="kv", type="wal_file_bytes"}[2m])) by (instance) + (sum(rate(tikv_engine_flow_bytes{db="raft", type="wal_file_bytes"}[2m])) by (instance) or (0 * sum(rate(raft_engine_write_size_sum[2m])) by (instance))) + (sum(rate(raft_engine_write_size_sum[2m])) by (instance) or (0 * sum(rate(tikv_engine_flow_bytes{db="raft", type="wal_file_bytes"}[2m])) by (instance)))`,
		`sum(rate(tikv_engine_flow_bytes{db="kv", type=~"bytes_read|iter_bytes_read"}[2m])) by (instance)`,
		`sum(tikv_store_size_bytes{type="used"}) by (instance)`,
		`avg by (cluster_id, instance) (tiflash_system_asynchronous_metric_Uptime)`,
		`sum(rate(node_cpu_seconds_total{mode!="idle",component="tiflash"}[2m])) by (instance)`,
		`avg(node_memory_MemTotal_bytes{component="tiflash"} - node_memory_MemAvailable_bytes{component="tiflash"}) by (instance)`,
		`sum(rate(tiflash_system_profile_event_WriteBufferFromFileDescriptorWriteBytes[2m])) by (instance) + sum(rate(tiflash_system_profile_event_PSMWriteBytes[2m])) by (instance) + sum(rate(tiflash_system_profile_event_WriteBufferAIOWriteBytes[2m])) by (instance)`,
		`sum(rate(tiflash_system_profile_event_ReadBufferFromFileDescriptorReadBytes[2m])) by (instance) + sum(rate(tiflash_system_profile_event_PSMReadBytes[2m])) by (instance) + sum(rate(tiflash_system_profile_event_ReadBufferAIOReadBytes[2m])) by (instance)`,
		`sum(tiflash_system_current_metric_StoreSizeUsed) by (instance)`,
	}
	sdkQueries = []string{
		"tidb_executor_statement_total",
		"tikv_engine_size_bytes",
		`label_replace(time()-process_start_time_seconds{cluster_id=~".*",component=~"ti.*"},"name","$1","instance","(.*)")`,
	}
	grafanaSeries = []string{
		"tikv_engine_size_bytes",
		"tikv_storage_command_total",
		"tikv_engine_block_cache_size_bytes",
		"tidb_server_connections",
	}
)

func NewStressQuerier(cluster o11yv1.Cluster) *StressQuerier {
	return &StressQuerier{
		cluster: cluster,
		stopCh:  make(chan struct{}),
		client:  newQueryClient(fmt.Sprintf(vmSelectUrlTemplate, cluster.Name)),
	}
}

func registerStressQuerier(ctx context.Context, cluster *o11yv1.Cluster) {
	// add stress querier to queriers if not exists
	if _, ok := queriers[cluster.Name]; !ok {
		lock.Lock()
		defer lock.Unlock()
		q := NewStressQuerier(*cluster)
		queriers[cluster.Name] = q
		q.start(ctx)
	}
}

func unregisterStressQuerier(cluster o11yv1.Cluster) {
	// remove stress querier from queriers if exists
	if _, ok := queriers[cluster.Name]; ok {
		queriers[cluster.Name].stop()
		delete(queriers, cluster.Name)
	}
}

type StressQuerier struct {
	cluster o11yv1.Cluster
	client  promv1.API
	stopCh  chan struct{}
}

func (s *StressQuerier) start(ctx context.Context) {
	tasks := []task{
		{s.instantQuery, Instant, integrationQueries, 1 * time.Minute, 0 * time.Second},
		// grafana queries appears to be executed together during troubleshooting.
		{s.randRangeQuery(true, 5), Range, append(append(append(pdGrafanaQueries, tidbGrafanaQueries...), tikvGrafanaQueries...), nodeGrafanaQueries...), 1 * time.Hour, 24 * time.Hour},
		// {s.randRangeQuery, tidbGrafanaQueries, 15 * time.Minute, 24 * time.Hour},
		// {s.randRangeQuery, tikvGrafanaQueries, 15 * time.Minute, 24 * time.Hour},
		{s.seriesQuery, Range, grafanaSeries, 1 * time.Hour, 1 * time.Hour},
		{s.randRangeQuery(false, defaultWorkers), Range, serverConsoleQueries, 5 * time.Minute, 12 * time.Hour},
		{s.rangeQuery, Instant, sdkQueries, 3 * time.Minute, 3 * time.Minute},
	}

	for _, t := range tasks {
		go t.run(ctx, s.stopCh)
	}
	logger.Info(fmt.Sprintf("started querier task of cluster %s", s.cluster.Name))
}

func (s *StressQuerier) stop() {
	logger.Info(fmt.Sprintf("stop querier of cluster %s", s.cluster.Name))
	close(s.stopCh) // Send stop signal to the goroutine
}

func newQueryClient(url string) promv1.API {
	apiClient, err := api.NewClient(api.Config{
		Address: url,
	})

	if err != nil {
		logger.Error(err, "failed to create prometheus client")
		return nil
	}
	return promv1.NewAPI(apiClient)
}

type task struct {
	Fn           func(context.Context, []string, time.Duration)
	Type         TaskType
	queries      []string
	tickDuration time.Duration
	lookback     time.Duration
}

type TaskType string

const (
	Instant TaskType = "instant"
	Range   TaskType = "range"
)

func (t *task) run(ctx context.Context, stopCh <-chan struct{}) {
	// sleep a random time for the first run to avoid query storm after controller restart
	switch t.Type {
	case Instant:
		time.Sleep(time.Duration(rand.Intn(180)) * time.Second)
		t.Fn(ctx, t.queries, t.lookback)
	case Range:
		time.Sleep(time.Duration(rand.Intn(60)) * time.Minute)
		t.Fn(ctx, t.queries, t.lookback)
	}

	ticker := time.NewTicker(t.tickDuration)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			t.Fn(ctx, t.queries, t.lookback)
		case <-stopCh:
			return
		case <-ctx.Done():
			return
		}
	}
}

func (s *StressQuerier) instantQuery(ctx context.Context, queries []string, _ time.Duration) {
	for _, promql := range queries {
		_, _, err := s.client.Query(ctx, promql, time.Now())
		if err != nil {
			logger.Error(err, fmt.Sprintf("failed instant query '%s' for cluster %s", promql, s.cluster.Name))
		}
	}
}

func (s *StressQuerier) randRangeQuery(chained bool, workers int) func(context.Context, []string, time.Duration) {
	return func(ctx context.Context, queries []string, lookback time.Duration) {
		max := time.Now()
		smallRange := time.Duration(rand.Intn(int(lookback.Minutes()))) * time.Minute
		bigRange := time.Duration((rand.Intn(14) + 1)) * lookback
		workerCh := make(chan struct{}, workers)

		// Run the queries concurrently in separate goroutines
		for i, promql := range queries {
			r, step := smallRange, 30*time.Second
			if rand.Intn(100) > rangeFactor {
				r = bigRange
				step = time.Duration(math.Ceil(r.Minutes()/maxDataPoints)) * time.Minute
			}
			q := promql
			// Sleep a random time to avoid query storm
			if chained && i > 0 && i%(2*workers) == 0 {
				time.Sleep(time.Duration(rand.Intn(5)) * time.Second)
			}
			// Acquire a worker slot (blocking if the limit is reached)
			workerCh <- struct{}{}
			go func() {
				opCtx, cancel := context.WithTimeout(ctx, queryTimeout)
				defer func() {
					<-workerCh
					cancel()
				}()
				_, _, err := s.client.QueryRange(opCtx, q, v1.Range{
					Start: max.Add(-r),
					End:   max,
					Step:  step,
				})
				if err != nil {
					logger.Error(err, fmt.Sprintf("failed range query '%s' for cluster %s", q, s.cluster.Name))
				}
			}()
		}
	}
}

func (s *StressQuerier) rangeQuery(ctx context.Context, queries []string, t time.Duration) {
	max := time.Now()
	min := max.Add(-t)
	workerCh := make(chan struct{}, defaultWorkers)
	for _, promql := range queries {
		q := promql
		// Acquire a worker slot (blocking if the limit is reached)
		workerCh <- struct{}{}
		go func() {
			opCtx, cancel := context.WithTimeout(ctx, queryTimeout)
			defer func() {
				<-workerCh
				cancel()
			}()
			_, _, err := s.client.QueryRange(opCtx, q, v1.Range{
				Start: min,
				End:   max,
				Step:  30 * time.Second,
			})
			if err != nil {
				logger.Error(err, fmt.Sprintf("failed range query %s for cluster %s", q, s.cluster.Name))
			}
		}()
	}
}

func (s *StressQuerier) seriesQuery(ctx context.Context, matches []string, t time.Duration) {
	end := time.Now()
	start := end.Add(-t)
	for _, match := range matches {
		_, _, err := s.client.Series(ctx, []string{match}, start, end)
		if err != nil {
			logger.Error(err, fmt.Sprintf("failed series query %s for cluster %s", match, s.cluster.Name))
		}
	}
}