package betting

import (
	"context"
	"fmt"
	"time"

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

	"github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/endless/pkg/workloadnode"
	"github.com/pingcap/endless/pkg/workloadnode/cmdstr"
	"github.com/pingcap/endless/testcase/oltp/oltppkg"
	"github.com/pingcap/endless/testcase/oltp/scenario/params"
)

var dbToDrop = "betting_0"

func InitializeTruncate001Params() *params.Truncate {
	return &params.Truncate{
		RunTime:    2 * time.Hour,
		WarmUpTime: 30 * time.Second,
	}
}

func restoreDBToDrop(
	ctx context.Context,
	tc resource.TiDBCluster,
	br resource.BR,
	opts cmdstr.BettingOptions,
	tidbClient *tidb.Client,
) error {
	// Only restore db if it does not exist.
	if exists, err := tidbClient.IsDBExist(ctx, dbToDrop); err != nil {
		return err
	} else if exists {
		log.Info("The database to drop already exists")
		return nil
	}

	log.Info("Restore the database to drop from minio")
	// Initialize br options.
	var brOpts cmdstr.BROptions
	url, err := tc.ServiceURL(resource.PDAddr)
	if err != nil {
		return err
	}
	brOpts.PD = []string{fmt.Sprintf("%s:%s", url.Hostname(), url.Port())}
	// The size of 20m records is about 10GB (three replicas). If the record count
	// is large than 20m, we assume it is in prod env. Otherwise, it is in debug env.
	// Restore different dataset for different env.
	// TODO: set `storage` as test case parameter.
	if opts.RecordCount > 20*M {
		// Data set "s3://brie/cc-5.4/backupfull-2" is about 1.5TB
		brOpts.Storage = "s3://brie/cc-5.4/backupfull-2"
	} else {
		// Data set "s3://brie/cc-5.4/backupdb-betting_0-1G" is about 2.342GB per replica.
		// It is about 1.1GB after compressed.
		brOpts.Storage = "s3://brie/cc-5.4/backupdb-betting_0-1G"
	}
	brOpts.Storage += "?access-key=minioadmin&secret-access-key=minioadmin&endpoint=http%3a%2f%2fminio.pingcap.net%3a9000&force-path-style=true"
	cmdStr := cmdstr.NewBR("/br").RestoreFull().CommonOptions(&brOpts).String()
	cmd := workloadnode.CommandContext(ctx, cmdStr)
	return cmd.Run(br)
}

// DropData drop a db and check
func TruncateDB(
	ctx context.Context,
	tc resource.TiDBCluster,
	br resource.BR,
	opts cmdstr.BettingOptions,
	tidbClient *tidb.Client,
) error {
	if err := restoreDBToDrop(ctx, tc, br, opts, tidbClient); err != nil {
		return err
	}

	oltppkg.WaitFor(ctx, "wait till the database settle down", 10*time.Second)

	// Drop the data.
	log.Info("Drop database", zap.String("db", dbToDrop))
	query := fmt.Sprintf("drop database %s;", dbToDrop)
	_, err := tidbClient.ExecContext(ctx, query)
	if err != nil {
		return err
	}
	return nil
}
