package main

import (
	"database/sql"
	"flag"
	"fmt"
	"net/url"

	"github.com/joho/godotenv"
	"go.uber.org/zap"

	. "github.com/pingcap/test-infra/sdk/core"
	. "github.com/pingcap/test-infra/sdk/resource"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"
)

var (
	testDSN      = flag.String("dsn", "", "test dsn")
	testHost     = flag.String("host", "127.0.0.1", "test host")
	testPort     = flag.Int("port", 3306, "test port")
	testUser     = flag.String("user", "", "test username")
	testPassword = flag.String("password", "", "test password")
)

func main() {
	flag.Parse()
	godotenv.Load()
	ctx := try(BuildContext()).(TestContext)
	tc := ctx.Resource("tc").(TiDBCluster)
	goTpc := ctx.Resource("go-tpc-workload").(WorkloadNode)
	sugar := ctx.Logger().Sugar()
	defer sugar.Sync()

	sugar.Infof("dsn=%s,host=%s,port=%d,user=%s,password=%s", *testDSN, *testHost, *testPort, *testUser, *testPassword)

	openGeneralLog(ctx, sugar, tc, "test")
	prepare(ctx, sugar, tc, goTpc)
	runTPCC(ctx, sugar, tc, goTpc)
}

func openGeneralLog(ctx TestContext, logger *zap.SugaredLogger, tc TiDBCluster, dbName string) {
	dbURL := try(tc.ServiceURL(DBAddr)).(*url.URL)
	logger.Infof("the tidb URL is: %s", dbURL)
	db := try(sql.Open("mysql", BuildDSN(dbURL, dbName))).(*sql.DB)
	try(db.Exec("set @@tidb_general_log=on"))
}

func prepare(ctx TestContext, sugar *zap.SugaredLogger, tc TiDBCluster, goTpc WorkloadNode) {
	dbURL := try(tc.ServiceURL(DBAddr)).(*url.URL)
	sugar.Infof("the tidb URL is: %s", dbURL)

	stdout, stderr, ec, err := goTpc.Exec(WorkloadNodeExecOptions{Command: fmt.Sprintf(`/go-tpc tpcc --host %s -P4000 --warehouses 4 prepare -T 4`, dbURL.Hostname())})
	if err != nil {
		Fail(fmt.Sprintf("prepare data error: %s\n\tstdout: %s\n\tstderr: %s", err, stdout, stderr))
	}
	if ec != 0 {
		Fail(fmt.Sprintf("prepare data failed, exit code: %d\n\tstdout: %s\n\tstderr: %s", ec, stdout, stderr))
	}
}

func runTPCC(ctx TestContext, sugar *zap.SugaredLogger, tc TiDBCluster, goTpc WorkloadNode) {
	dbURL := try(tc.ServiceURL(DBAddr)).(*url.URL)
	sugar.Infof("the tidb URL is: %s", dbURL)

	stdout, stderr, ec, err := goTpc.Exec(WorkloadNodeExecOptions{Command: fmt.Sprintf(`/go-tpc tpcc --host %s -P4000 --warehouses 4 run -T 4 --time 5m`, dbURL.Hostname())})
	if err != nil {
		Fail(fmt.Sprintf("run error: %s\n\tstdout: %s\n\tstderr: %s", err, stdout, stderr))
	}
	if ec != 0 {
		Fail(fmt.Sprintf("run failed, exit code: %d\n\tstdout: %s\n\tstderr: %s", ec, stdout, stderr))
	}
	sugar.Infof("stdout: %s\nstderr: %s", stdout, stderr)
}

func try(xs ...interface{}) interface{} {
	if len(xs) == 0 {
		return nil
	}
	if err, ok := xs[len(xs)-1].(error); ok && err != nil {
		Fail(err.Error())
	}
	return xs[0]
}