package oltppkg

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

	"github.com/pingcap/endless/pkg/client/tidb"
	. "github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
)

type LightningRunCmd struct {
	lightningPath string
	tidbPort      string
	tidbHost      string
	pdAddr        string
	sortedKvDir   string
	storageUrl    string
	s3Parms       S3Parms
}

func NewLightning() *LightningRunCmd {
	b := &LightningRunCmd{}
	return b
}
func (l *LightningRunCmd) TidbPort(tidbPort string) *LightningRunCmd {
	l.tidbPort = tidbPort
	return l
}
func (l *LightningRunCmd) TidbHost(tidbHost string) *LightningRunCmd {
	l.tidbHost = tidbHost
	return l
}
func (l *LightningRunCmd) LightningPath(lightningPath string) *LightningRunCmd {
	l.lightningPath = lightningPath
	return l
}
func (l *LightningRunCmd) PdAddr(pdAddr string) *LightningRunCmd {
	l.pdAddr = pdAddr
	return l
}
func (l *LightningRunCmd) SortedKvDir(sortedKvDir string) *LightningRunCmd {
	l.sortedKvDir = sortedKvDir
	return l
}
func (l *LightningRunCmd) StorageUrl(storageUrl string) *LightningRunCmd {
	l.storageUrl = storageUrl
	return l
}
func (l *LightningRunCmd) S3Parms(s3Parms S3Parms) *LightningRunCmd {
	l.s3Parms = s3Parms
	return l
}
func (l *LightningRunCmd) Build() string {
	cmdTemp := `%s/tidb-lightning --tidb-host=%s --tidb-port=%s --backend=local --sorted-kv-dir=%s -d "%s?access-key=%s&secret-access-key=%s&endpoint=%s&force-path-style=true"`
	cmd := fmt.Sprintf(cmdTemp, l.lightningPath, l.tidbHost, l.tidbPort, l.sortedKvDir, l.storageUrl,
		l.s3Parms.S3AccessKeyID, l.s3Parms.S3SecretAccessKey, l.s3Parms.S3Endpoint)
	return cmd
}

type BrRunCmd struct {
	brPath    string
	options   []string
	command   string
	accessid  string
	accesskey string
}

func NewBr(brPath string, command string, accessid string, accesskey string) *BrRunCmd {
	b := &BrRunCmd{command: command, accessid: accessid, accesskey: accesskey, brPath: brPath}
	b.options = append(b.options, "--send-credentials-to-tikv=true")
	//b.options = append(b.options, "---check-requirements=false")
	return b
}
func (b *BrRunCmd) Pd(pd string) *BrRunCmd {
	b.options = append(b.options, fmt.Sprintf("--pd %s", pd))
	return b
}
func (b *BrRunCmd) Storage(storage string) *BrRunCmd {
	b.options = append(b.options, fmt.Sprintf("--storage %s", storage))
	return b
}
func (b *BrRunCmd) S3endpoint(s3endpoint string) *BrRunCmd {
	b.options = append(b.options, fmt.Sprintf("--s3.endpoint  %s", s3endpoint))
	return b
}
func (t *BrRunCmd) Command(command string) *BrRunCmd {
	t.command = command
	return t
}
func (t *BrRunCmd) CheckReq(checkReq bool) *BrRunCmd {
	t.options = append(t.options, fmt.Sprintf("--check-requirements=%t", checkReq))
	return t
}
func (s *BrRunCmd) Build() string {
	return fmt.Sprintf("export AWS_ACCESS_KEY_ID=%s;export AWS_SECRET_ACCESS_KEY=%s;%s/br %s %s",
		s.accessid, s.accesskey, s.brPath, s.command, strings.Join(s.options[:], " "))
}

func BrRestoreData(ctx context.Context, tc resource.MySQLCompatible, br resource.BR, s3Parms S3Parms, caseCmdParms CaseCmdParms) {
	// Try(ExecPdCtl(tc, "log debug"))
	// defer func() {
	// 	Try(ExecPdCtl(tc, "log info"))
	// }()
	ctxbr, cancel := context.WithTimeout(ctx, 6*time.Hour)
	defer cancel()
	db := Try(tc.GetDB("")).(*sql.DB)
	tidbClient := tidb.NewClient(db)
	defer db.Close()
	importFlagTpcc := false
	importFlagSysbench := false

	pdURL := Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)

	LogAndBy("BR restore data for tpcc")
	brcmd := NewBr("", "restore full", s3Parms.S3AccessKeyID, s3Parms.S3SecretAccessKey)
	brTpccCmd := brcmd.Pd(pdURL.String()).Storage(caseCmdParms.TpccParms.Storage).S3endpoint(s3Parms.S3Endpoint).CheckReq(false).Build()
	exist := Try(tidbClient.IsTableExist(ctxbr, caseCmdParms.TpccParms.Db, "new_order")).(bool)
	if !exist {
		Try(ExecWorkloadCmd(ctxbr, br, brTpccCmd, true, nil))
		importFlagTpcc = true
	}

	LogAndBy("BR restore data for sysbench")
	brcmd = NewBr("", "restore full", s3Parms.S3AccessKeyID, s3Parms.S3SecretAccessKey)
	exist = Try(tidbClient.IsTableExist(ctxbr, caseCmdParms.SysbenchParms.Db, "sbtest1")).(bool)
	brSysbenchCmd := brcmd.Pd(pdURL.String()).Storage(caseCmdParms.SysbenchParms.Storage).S3endpoint(s3Parms.S3Endpoint).CheckReq(false).Build()
	if !exist {
		Try(ExecWorkloadCmd(ctxbr, br, brSysbenchCmd, true, nil))
		importFlagSysbench = true
	}

	if caseCmdParms.WorkloadType.Type[0] == "ycsb" {
		LogAndBy("BR restore data for ycsb")
		brcmd := NewBr("", "restore full", s3Parms.S3AccessKeyID, s3Parms.S3SecretAccessKey)
		exist := Try(tidbClient.IsTableExist(ctxbr, caseCmdParms.YcsbParms.Db, "usertable")).(bool)
		brSysbenchCmd := brcmd.Pd(pdURL.String()).Storage(caseCmdParms.YcsbParms.Storage).S3endpoint(s3Parms.S3Endpoint).CheckReq(false).Build()
		if !exist {
			Try(ExecWorkloadCmd(ctxbr, br, brSysbenchCmd, true, nil))
		}
	}

	if !CheckIfHeterogeneousDisk(tc) && importFlagSysbench && importFlagTpcc {
		log.Info("After br import succeeded, now waiting for balance if this is not heterogeneous cluster.")
		Try(BasicClusterExecPdCtl(tc, "store limit all 200"))
		WaitForBalance(ctx, tc, 3*time.Hour, 5*time.Minute, 0.2)
		Try(BasicClusterExecPdCtl(tc, "store limit all 15"))
	}
}