package utils

import (
	"bytes"
	"database/sql"
	"fmt"
	"sort"

	pdclient "github.com/pingcap/endless/pkg/client/pd"
	. "github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/core"
	"github.com/pingcap/test-infra/sdk/resource"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/local"
	"github.com/tidwall/gjson"
	"go.uber.org/zap"
)

type SuiteContext struct {
	core.TestContext
	TC resource.TiDBCluster
	DB *sql.DB
	PD *pdclient.Client
}

func InitTiKVLabels(cluster resource.TiDBCluster, labels []string) {
	pd := Try(pdclient.NewPDClient(cluster)).(*pdclient.Client)

	// init placement labels
	stores := make([]uint64, 0)
	for _, storeJSON := range gjson.GetBytes(Try(pd.GetJsonBytes("/pd/api/v1/stores")).([]byte), "stores.#.store").Array() {
		id := storeJSON.Get("id")
		storeLabels := storeJSON.Get("labels")
		if storeLabels.Exists() {
			isTiflash := false
			for _, storeLabel := range storeLabels.Array() {
				kv := storeLabel.Map()
				key := kv["key"].String()
				value := kv["value"].String()
				if key == "engine" && value == "tiflash" {
					isTiflash = true
					break
				}
			}
			if isTiflash {
				continue
			}
		}
		stores = append(stores, id.Uint())
	}
	sort.Slice(stores, func(i, j int) bool {
		return stores[i] < stores[j]
	})
	for j, storeId := range stores {
		Try(pd.Post(fmt.Sprintf("%s/pd/api/v1/store/%d/label", pd.GetPDAddr(), storeId),
			"application/json", bytes.NewReader([]byte(labels[j]))))
	}
}

func SwitchPDLeader(tc resource.TiDBCluster) error {
	client, err := tc.PDClient()
	if err != nil {
		return err
	}

	members, err := client.GetMembers()
	if err != nil {
		log.Error("get pd members failed", zap.Error(err))
		return errors.Trace(err)
	}

	log.Info("pd leader", zap.String("leader-name", members.Leader.Name))
	var targetName string
	for _, member := range members.Members {
		if member.MemberID != members.Leader.MemberID {
			targetName = member.Name
			break
		}
	}

	if err := client.TransferPDLeader(targetName); err != nil {
		log.Error("transfer pd leader failed", zap.Error(err))
		return errors.Trace(err)
	}

	return nil
}