package scenario

import (
	"context"

	. "github.com/onsi/ginkgo"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"

	. "github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/pkg/workloadnode/cmdstr"
	"github.com/pingcap/endless/testcase/oltp/scenario/betting"
	"github.com/pingcap/endless/testcase/oltp/scenario/qhad"
)

type WorkloadOp string

var (
	WorkloadRun  WorkloadOp = "load"
	WorkloadLoad WorkloadOp = "run"
)

// TWorkloadLoad is the body of test case `scenario_workload_load`.
func TWorkloadLoad001(rootCtx context.Context, config Config, tc resource.TiDBCluster) {
	Try(WorkloadDoOp(rootCtx, config, tc, WorkloadLoad))
}

// TWorkloadLoad is the body of test case `scenario_workload_run`.
func TWorkloadRun001(rootCtx context.Context, config Config) {
	// TODO: check qps and latency.

	// Wait until workload finish.
	select {
	case <-rootCtx.Done():
	}
}

// WorkloadDoOp run workload with specified operation.
func WorkloadDoOp(rootCtx context.Context, config Config, tc resource.TiDBCluster, op WorkloadOp) error {
	By("Loading workload data")

	// Initialize common variables.
	workloadName := config.WorkloadName

	// Initialize variables based on workload op.
	var subcmd cmdstr.GoycsbSubcmd
	if op == WorkloadRun {
		subcmd = cmdstr.GoycsbRun
	} else if op == WorkloadLoad {
		subcmd = cmdstr.GoycsbLoad
	} else {
		panic("unimplemented")
	}

	// Do op.
	if workloadName == "betting" {
		return betting.RunWorkload(suiteTestCtx, tc, rootCtx, config.BettingOpts, subcmd)
	} else if workloadName == "qhad" {
		return qhad.RunWorkload(suiteTestCtx, tc, rootCtx, config.QHAdOpts, subcmd)
	} else {
		panic("unimplemented")
	}
	return nil
}

func WorkloadRunForever(rootCtx context.Context, config Config, tc resource.TiDBCluster) {
	workloadName := config.WorkloadName

	// Run workload continuously.
	go func() {
	RUNLOOP:
		for {
			select {
			case <-rootCtx.Done():
				break RUNLOOP
			default:
			}
			log.Info("Run workload",
				zap.String("workload", workloadName))

			err := WorkloadDoOp(rootCtx, config, tc, WorkloadRun)

			if err != nil {
				log.Error("Workload failed",
					zap.String("workload", workloadName),
					zap.Error(err))
			} else {
				log.Info("Workload finished",
					zap.String("workload", workloadName))
			}

		}
	}()
}