package arsenal

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"time"

	"github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/schrodinger/util"
)

// Config configures the PDShuffle function
type Config struct {
	PDAddrs []string

	ShuffleLeader bool
	ShuffleRegion bool

	ShufflePDLeader      bool
	ShuffleInterval      time.Duration
	ShufflePerRound      int
	ShuffleRoundInterval time.Duration
}

// PDShuffle is a blocking function to shuffle region and leader
func PDShuffle(ctx context.Context, cfg Config) {
	if len(cfg.PDAddrs) == 0 {
		client.Warnf("[nemesis] empty pd address cfg: %v", cfg)
		return
	}

	var schedulerNames []string
	if cfg.ShuffleLeader {
		schedulerNames = append(schedulerNames, "shuffle-leader-scheduler")
	}
	if cfg.ShuffleRegion {
		schedulerNames = append(schedulerNames, "shuffle-region-scheduler")
	}

	if len(schedulerNames) == 0 {
		client.Warnf("[nemesis] empty pd shuffle scheduler cfg: %v", cfg)
	}

AddSchedulers:
	for _, name := range schedulerNames {
		data := map[string]string{"name": name}
		b, _ := json.Marshal(data)
		// Make sure that schedulers are added to the PD leader.
		for {
			for _, addr := range cfg.PDAddrs {
				err := requestPD(
					util.GenURL(addr+"/pd/api/v1/schedulers"),
					"application/json",
					bytes.NewBuffer(b),
				)

				if err == nil || strings.Contains(fmt.Sprintf("%s", err), "scheduler existed") {
					client.Infof("[nemesis] add %s to pd %v", name, addr)
					continue AddSchedulers
				}

				client.Errorf(
					"[nemesis] fail to add %s to pd %s, err: %v",
					name, addr, err)
			}
			// Backoff, cluster may not bootstrapped.
			time.Sleep(time.Second)
		}
	}

	length := len(cfg.PDAddrs)
	for cfg.ShufflePDLeader {
		for i := 0; i < cfg.ShufflePerRound; i++ {
			addr := cfg.PDAddrs[i%length]
			err := requestPD(
				util.GenURL(addr+"/pd/api/v1/leader/resign"),
				"",
				strings.NewReader(""),
			)
			if err != nil {
				client.Errorf(
					"[nemesis] fail to shuffle pd leader, err %s", err)
			}

			wait, _ := context.WithTimeout(ctx, cfg.ShuffleInterval)
			select {
			case <-ctx.Done():
				return
			case <-wait.Done():
			}
		}

		wait, _ := context.WithTimeout(ctx, cfg.ShuffleRoundInterval)
		select {
		case <-ctx.Done():
			return
		case <-wait.Done():
		}
	}
}

func requestPD(url string, contentType string, body io.Reader) error {
	resp, err := http.Post(url, contentType, body)
	if err != nil ||
		(resp != nil && resp.StatusCode/100 != 2) {
		var reason []byte
		if resp != nil {
			reason, _ = ioutil.ReadAll(resp.Body)
		}
		return fmt.Errorf("err %v, reason: %s, resp: %+v", err, reason, resp)
	}
	return nil
}