package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"time"

	"github.com/montanaflynn/stats"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/caselib/pkg/model/ticdc"
	"github.com/pingcap/test-infra/caselib/pkg/model/workload"
	"github.com/pingcap/test-infra/caselib/pkg/steps"
	"github.com/pingcap/test-infra/caselib/pkg/utils"
	"github.com/pingcap/test-infra/caselib/ticdc/perf/db"
	"github.com/pingcap/test-infra/caselib/ticdc/perf/workload/schema"
	apiresource "github.com/pingcap/test-infra/common/model/resource"
	sdkresource "github.com/pingcap/test-infra/sdk/resource"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"
	"github.com/pingcap/tiflow/cdc/model"
	prometheusapi "github.com/prometheus/client_golang/api/prometheus/v1"
	prometheusmodel "github.com/prometheus/common/model"
	"github.com/tikv/client-go/v2/oracle"
	"go.uber.org/zap"
)

type Operation string

const (
	addIndex          Operation = "add-index"
	reloadTiKV        Operation = "reload-tikv"
	reloadTiCDC       Operation = "reload-ticdc"
	restartChangeFeed Operation = "restart-changefeed"
)

var (
	dbName                 = "workload"
	resourceHandleSource   = &apiresource.ResourceHandle{}
	resourceHandleSink     = &apiresource.ResourceHandle{}
	resourceHandleWorkload = &apiresource.ResourceHandle{}
	feedIDs                []string

	metricsDBHost     string
	metricsDBPort     int
	metricsDBUser     string
	metricsDBPassword string
	metricsDBName     string

	option boundaryStepOption
)

type boundaryStepOption struct {
	name string
	// workload
	TableCount int   `json:"table_count"`
	Tps        int   `json:"tps"`
	Thread     int   `json:"thread"`
	RowCount   int64 `json:"row_count"`

	ChangeFeedCount int  `json:"change_feed_count"`
	CdcNodeCount    uint `json:"cdc_node_count"`

	// operation during sync
	Operation Operation `json:"operation"`

	SyncOldDataOnly bool `json:"sync_old_data_only"`

	SinkToKafka bool `json:"sink_to_kafka"`
}

func init() {
	flag.StringVar(&option.name, "case-name", "", "case name")
	flag.IntVar(&option.TableCount, "table-count", 1, "table count of the workload")
	flag.IntVar(&option.Tps, "tps", 20000, "Tps of the workload")
	flag.IntVar(&option.Thread, "thread", 100, "thread of the workload")
	flag.Int64Var(&option.RowCount, "row-count", 20000000, "row count of the workload")
	flag.IntVar(&option.ChangeFeedCount, "change-feed-count", 1, "change feed count")
	flag.UintVar(&option.CdcNodeCount, "cdc-count", 1, "cdc node count")
	flag.StringVar((*string)(&option.Operation), "operation", "add-index", "operation during sync, add-index, reload-tikv, reload-ticdc, restart-changefeed")
	flag.BoolVar(&option.SyncOldDataOnly, "old-data", false, "sync old data only")
	flag.BoolVar(&option.SinkToKafka, "sink-to-kafka", false, "sink data to kafka, otherwise sink to mysql")

	flag.StringVar(&metricsDBHost, "database-host", "172.16.4.36", "metrics database host")
	flag.StringVar(&metricsDBUser, "database-user", "cdc", "metrics database user")
	flag.StringVar(&metricsDBPassword, "database-password", "gI93XU1lMVMR", "metrics database password")
	flag.StringVar(&metricsDBName, "database-db-name", "cdc_perf", "metrics database db name")
	flag.IntVar(&metricsDBPort, "database-port", 4000, "metrics database port")
	flag.Parse()
}

type cdcBoundaryStep struct {
	tc            sdkresource.TiDBCluster
	prometheusAPI prometheusapi.API
	option        *boundaryStepOption
	Case          *db.Case
}

func NewCdcBoundaryStep(ctx context.Context, option *boundaryStepOption) (*cdcBoundaryStep, error) {
	caseName, err := getCaseName(option)
	if err != nil {
		return nil, errors.Trace(err)
	}
	sinkName := "TI_PARAM_RES_downstream"
	if option.SinkToKafka {
		sinkName = "TI_PARAM_RES_kafka"
	}
	updateGcTime := "update mysql.tidb set VARIABLE_VALUE=\"2400h\" where VARIABLE_NAME=\"tikv_gc_life_time\";"
	builder := steps.NewBuilder().
		LoadEnv().
		LoadResourceFromEnv("TI_PARAM_RES_upstream", resourceHandleSource).
		LoadResourceFromEnv(sinkName, resourceHandleSink).
		LoadResourceFromEnv("TI_PARAM_RES_workload", resourceHandleWorkload).
		InitCDCComponents(resourceHandleSource, resourceHandleSink).
		InitWorkload(resourceHandleWorkload).
		CreateDatabase(resourceHandleSource, dbName).
		ExecuteSqlStatement(resourceHandleSource, &updateGcTime)
	if err := builder.Build().Execute(ctx); err != nil {
		return nil, errors.Trace(err)
	}
	tc, err := steps.GetInner(ctx).Board().TiDBCluster()
	if err != nil {
		return nil, errors.Trace(err)
	}
	prometheusAPI, err := tc.PrometheusAPI()
	if err != nil {
		return nil, errors.Trace(err)
	}
	boundaryCase, err := db.AddCase(caseName)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return &cdcBoundaryStep{
		tc:            tc,
		prometheusAPI: prometheusAPI,
		option:        option,
		Case:          boundaryCase,
	}, nil
}

func getCaseName(option *boundaryStepOption) (string, error) {
	name := option.name
	if name == "" {
		data, err := json.Marshal(option)
		if err != nil {
			return "", errors.Trace(err)
		}
		name = string(data)
	}
	return name, nil
}

// Execute runWorkload run the specified workload in a background goroutine.
func (s *cdcBoundaryStep) Execute(ctx context.Context) (err error) {
	// init case and run workload a while , then record the upstream tidb metrics
	builder := steps.NewBuilder().RunWorkload(s.getWorkloadOptions())
	startTime := time.Now()
	if err := builder.Build().Execute(ctx); err != nil {
		return err
	}
	endTime := time.Now()
	if err = s.recordMetrics(ctx, false, startTime, endTime); err != nil {
		return errors.Trace(err)
	}

	var sinkType ticdc.SinkType = ticdc.SinkTypeMySQL
	if option.SinkToKafka {
		sinkType = ticdc.SinkTypeKafka
	} else {
		if err := s.initSinkDatabase(ctx); err != nil {
			return err
		}
	}

	startTs := uint64(0)
	if s.option.SyncOldDataOnly {
		startTs, err = s.getGCSafePointFromPD(ctx)
		if err != nil {
			return err
		}
		physicalTs := oracle.ExtractPhysical(startTs)
		log.Info("sync from earliest time", zap.Time("ts", time.Unix(physicalTs/1000, physicalTs%1000*1000_000)))
	}
	if err := s.removeOldChangeFeed(ctx); err != nil {
		return err
	}

	builder = steps.NewBuilder().ScaleTiDBCluster(sdkresource.TiCDC, uint32(s.option.CdcNodeCount), time.Minute*10)
	for i := 0; i < s.option.ChangeFeedCount; i++ {
		feedID := fmt.Sprintf("cdc-boundary-test-%d", i)
		feedIDs = append(feedIDs, feedID)
		builder = builder.OperateChangeFeed(&ticdc.ChangeFeedCreateFlags{SinkType: sinkType, ID: feedID, StartTs: startTs, NoConfirm: true}).
			OperateChangeFeed(&ticdc.ChangeFeedAdminJobFlags{JobType: model.AdminStop, ID: feedID})
	}
	if err := builder.Build().Execute(ctx); err != nil {
		return err
	}

	// start to watch ticdc oom event
	tc, err := steps.GetInner(ctx).Board().TiDBCluster()
	if err != nil {
		return err
	}
	cancelCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	eventChan, err := tc.WatchOOM(cancelCtx, sdkresource.TiDBClusterOOMWatchOptions{Target: sdkresource.TiCDC})
	if err != nil {
		return err
	}
	oomCount := 0
	go func() {
		for {
			var oom sdkresource.OOMEvent
			select {
			case oom = <-eventChan:
				log.Warn("ticdc is oom", zap.String("node", oom.InstanceName))
				oomCount++
			case <-cancelCtx.Done():
				return
			}
		}
	}()

	builder = steps.NewBuilder()
	var cs []steps.Step
	if !s.option.SyncOldDataOnly {
		// run workload
		cs = append(cs, steps.NewRunWorkload(s.getWorkloadOptions()))
	}

	// resume change feed
	for _, changeFeedID := range feedIDs {
		cs = append(cs, steps.NewChangeFeedTask(&ticdc.ChangeFeedAdminJobFlags{JobType: model.AdminResume, ID: changeFeedID}))
	}
	s.startOperation(ctx)
	builder = builder.Parallel(cs...).SendFinishMark(dbName)
	startTime = time.Now()
	if err := builder.Build().Execute(ctx); err != nil {
		return errors.Trace(err)
	}
	if err = s.recordMetrics(ctx, true, startTime, time.Now()); err != nil {
		return errors.Trace(err)
	}

	if err := steps.NewBuilder().WaitFinishMark(resourceHandleSink, dbName, 2*time.Hour).Build().Execute(ctx); err != nil {
		return err
	}
	s.Case.OOMCount = oomCount
	s.Case.SyncDurationSec = time.Now().Unix() - startTime.Unix()
	if err = db.UpdateCase(s.Case); err != nil {
		return errors.Trace(err)
	}
	log.Info("sync finished", zap.Duration("elapsed", time.Since(startTime)))
	return nil
}

func (s *cdcBoundaryStep) startOperation(ctx context.Context) {
	workloadTotalTime := s.option.RowCount / int64(s.option.Tps)
	// add operation
	if s.option.Operation != "" {
		go func() {
			<-time.After(time.Duration(workloadTotalTime/2) * time.Second)
			var err error
			switch s.option.Operation {
			case addIndex:
				err = s.addIndex()
			case reloadTiKV:
				err = s.restartNode(sdkresource.TiKV)
			case reloadTiCDC:
				err = s.restartNode(sdkresource.TiCDC)
			case restartChangeFeed:
				err = s.restartChangeFeed(ctx)
			}
			if err != nil {
				panic(err)
			}
		}()
	}
}

func (s *cdcBoundaryStep) addIndex() error {
	upstreamDB, err := s.tc.GetDB(dbName)
	if err != nil {
		return err
	}
	for i := 0; i < s.option.TableCount; i++ {
		_, err = upstreamDB.Exec(schema.GetAddIndexStatement(i))
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *cdcBoundaryStep) restartNode(compType sdkresource.TiDBCompType) error {
	nodes, err := s.tc.GetNodes(compType)
	if err != nil {
		return err
	}
	for _, node := range nodes {
		if err := node.Restart(); err != nil {
			return err
		}
	}
	return nil
}

func (s *cdcBoundaryStep) restartChangeFeed(ctx context.Context) error {
	cdcHost, err := steps.GetCDCHost(ctx)
	if err != nil {
		return err
	}
	resp, err := cdcHost.ChangeFeed(&ticdc.ChangeFeedListFlags{})
	if err != nil {
		return err
	}
	changeFeedList := resp.(ticdc.ChangeFeedList)
	for _, feed := range changeFeedList {
		log.Info("stop changed feed", zap.String("id", feed.ID))
		if _, err = cdcHost.ChangeFeed(&ticdc.ChangeFeedAdminJobFlags{JobType: model.AdminStop, ID: feed.ID}); err != nil {
			return errors.Annotate(err, "failed to stop change feed")
		}
	}
	for _, feed := range changeFeedList {
		log.Info("resume changed feed", zap.String("id", feed.ID))
		if _, err = cdcHost.ChangeFeed(&ticdc.ChangeFeedAdminJobFlags{JobType: model.AdminResume, ID: feed.ID}); err != nil {
			return errors.Annotate(err, "failed to resume change feed")
		}
	}
	return nil
}

func (s *cdcBoundaryStep) removeOldChangeFeed(ctx context.Context) error {
	cdcHost, err := steps.GetCDCHost(ctx)
	if err != nil {
		return err
	}
	resp, err := cdcHost.ChangeFeed(&ticdc.ChangeFeedListFlags{})
	if err != nil {
		return err
	}
	changeFeedList := resp.(ticdc.ChangeFeedList)
	for _, feed := range changeFeedList {
		log.Info("removing old changed feed", zap.String("id", feed.ID))
		if _, err = cdcHost.ChangeFeed(&ticdc.ChangeFeedAdminJobFlags{JobType: model.AdminRemove, ID: feed.ID, ForceRemove: true}); err != nil {
			return errors.Annotate(err, "failed to remove change feed")
		}
	}
	return nil
}

func (s *cdcBoundaryStep) getGCSafePointFromPD(ctx context.Context) (uint64, error) {
	tc, err := steps.GetInner(ctx).Board().TiDBCluster()
	if err != nil {
		return 0, errors.Trace(err)
	}
	pdClient, err := tc.PDClient()
	if err != nil {
		return 0, errors.Trace(err)
	}
	gcSafePoint, err := pdClient.GetGCSafePoint()
	if err != nil {
		return 0, errors.Trace(err)
	}
	return gcSafePoint.GCSafePoint, nil
}

func (s *cdcBoundaryStep) initSinkDatabase(ctx context.Context) error {
	sink, err := steps.GetInner(ctx).Board().Sink()
	if err != nil {
		return errors.Trace(err)
	}
	sinkURL, err := sink.SinkEndpoint()
	if err != nil {
		return errors.Trace(err)
	}
	sinkDB, err := sql.Open("mysql", sdkresource.BuildDSN(sinkURL, "", "timeout=30m", "readTimeout=30m", "writeTimeout=30m"))
	if err != nil {
		return errors.Trace(err)
	}
	if _, err = sinkDB.Exec("create database if not exists " + dbName); err != nil {
		return errors.Trace(err)
	}
	_ = sinkDB.Close()
	sinkDB, err = sql.Open("mysql", sdkresource.BuildDSN(sinkURL, dbName))
	if err != nil {
		return errors.Trace(err)
	}
	for i := 0; i < s.option.TableCount; i++ {
		if _, err = sinkDB.Exec(schema.GetCreateTableStatement(i)); err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

func (s *cdcBoundaryStep) String() string {
	return "cdcBoundaryStep"
}

func getMatrixStats(matrix prometheusmodel.Matrix) (float64, float64, float64, error) {
	var inputs []float64
	for _, stream := range matrix {
		for _, sample := range stream.Values {
			inputs = append(inputs, float64(sample.Value))
		}
	}
	if len(inputs) == 0 {
		log.Info("no metrics value is returned")
		return 0, 0, 0, nil
	}
	p80, err := stats.Percentile(inputs, 80)
	if err != nil {
		return 0, 0, 0, err
	}
	p95, err := stats.Percentile(inputs, 95)
	if err != nil {
		return 0, 0, 0, err
	}
	max, err := stats.Max(inputs)
	if err != nil {
		return 0, 0, 0, err
	}
	return p80, p95, max, nil
}

func (s *cdcBoundaryStep) recordMetric(ctx context.Context, cdcRunning bool, metricName, query string, start, end time.Time) (err error) {
	// tidb db duration
	values, _, err := s.prometheusAPI.QueryRange(ctx, query,
		prometheusapi.Range{
			Start: start,
			End:   end,
			Step:  time.Second * 30,
		})
	if err != nil {
		return errors.Trace(err)
	}
	bytes, err := json.Marshal(values)
	if err != nil {
		return errors.Trace(err)
	}
	jsonStr := string(bytes)
	p80, p95, max, err := getMatrixStats(values.(prometheusmodel.Matrix))
	if err != nil {
		return errors.Trace(err)
	}
	_, err = db.AddMetric(s.Case.ID, cdcRunning, metricName, p80, p95, max, jsonStr)
	if err != nil {
		return errors.Trace(err)
	}
	return nil
}

func (s *cdcBoundaryStep) recordMetrics(ctx context.Context, cdcRunning bool, start, end time.Time) (err error) {
	metrics := map[string]string{
		"histogram_quantile":            "histogram_quantile(0.95, sum(rate(tidb_server_handle_query_duration_seconds_bucket{}[1m])) by (le))",
		"tidb_executor_statement_total": "sum(rate(tidb_executor_statement_total{type=\"insert\"}[1m]))",
		"ticdc_owner_checkpoint_ts_lag": "max(ticdc_owner_checkpoint_ts_lag) by (changefeed)",
	}
	for metricsName, query := range metrics {
		if err := s.recordMetric(ctx, cdcRunning, metricsName, query, start, end); err != nil {
			return err
		}
	}
	return nil
}

func (s *cdcBoundaryStep) getWorkloadOptions() workload.Option {
	return &workload.CdcBoundaryOption{
		TableCount: s.option.TableCount,
		Tps:        s.option.Tps,
		RowCount:   s.option.RowCount,
		Thread:     s.option.Thread,
		DBName:     dbName,
		Duration:   time.Hour * 24,
	}
}

func main() {
	_, logger := utils.GetCtxAndLogger()
	if err := db.InitDataBase(metricsDBHost, metricsDBName, metricsDBUser, metricsDBPassword, metricsDBPort); err != nil {
		logger.Fatal(" failed to init database", zap.Error(err))
	}
	ctx := steps.New(context.Background())
	boundary, err := NewCdcBoundaryStep(ctx, &option)
	if err != nil {
		logger.Fatal(" ailed to init boundary test step", zap.Error(err))
	}
	if err := boundary.Execute(ctx); err != nil {
		logger.Fatal("failed", zap.Error(err))
	}
	logger.Info("passed")
}