package main

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

	"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"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"
	prometheus_api "github.com/prometheus/client_golang/api/prometheus/v1"
)

func main() {
	godotenv.Load()
	ctx := try(BuildContext()).(TestContext)
	tc := ctx.Resource("tc").(TiDBCloudCluster)
	dbName := ctx.Param("db-name", "test")
	goTpc := ctx.Resource("go-tpc-workload").(WorkloadNode)

	sugar := ctx.Logger().Sugar()
	defer sugar.Sync()

	opOnTiDB(ctx, tc, dbName)
	queryPDUptime(ctx, tc)

	prepare(ctx, sugar, tc, goTpc)
	runTPCC(ctx, sugar, tc, goTpc)

}

func opOnTiDB(ctx TestContext, tc TiDBCloudCluster, dbName string) {
	sugar := ctx.Logger().Sugar()
	defer sugar.Sync()
	db := try(tc.GetDB(dbName)).(*sql.DB)
	try(db.Exec("create table IF NOT EXISTS sbtest (id int(11) not null primary key, k int(11) not null, c char(120) not null default '', pad char(255) not null default '', key idx(k))"))
	sugar.Infof("create table succeed\n")
}
func queryPDUptime(ctx TestContext, tc TiDBCloudCluster) {
	sugar := ctx.Logger().Sugar()
	defer sugar.Sync()
	promAPI := try(tc.PrometheusAPI()).(prometheus_api.API)
	vs, ws, err := promAPI.Query(ctx, "(time() - process_start_time_seconds{component=~\".*pd.*\"})", time.Now())
	if err != nil {
		Fail(err.Error())
	}
	sugar.Infof("result: %s, warnings: %s", vs, ws)
}

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

	stdout, stderr, ec, err := goTpc.Exec(WorkloadNodeExecOptions{Command: fmt.Sprintf(`/go-tpc tpcc --host %s -P %s -p %s --warehouses 1 prepare -T 1`, dbURL.Hostname(), dbURL.Port(), password)})
	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 TiDBCloudCluster, goTpc WorkloadNode) {
	dbURL := try(tc.ServiceURL(resource.DBAddr)).(*url.URL)
	password, _ := dbURL.User.Password()
	sugar.Infof("the tidb URL is: %s", dbURL)

	stdout, stderr, ec, err := goTpc.Exec(WorkloadNodeExecOptions{Command: fmt.Sprintf(`/go-tpc tpcc --host %s -P %s -p %s --warehouses 1 run -T 1 --time 1m`, dbURL.Hostname(), dbURL.Port(), password)})
	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]
}