package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/pingcap/schrodinger-test/pd/hot_scheduler/types"
	"github.com/pingcap/schrodinger-test/pd/hot_scheduler/workload"
	client "github.com/pingcap/schrodinger/pkg/test_client"
	"golang.org/x/net/context"
	"gonum.org/v1/gonum/stat"
)

var (
	pwd      = flag.String("pwd", "", "password")
	database = flag.String("db", "test", "database")
	runMode  = flag.String("run-mode", "online", "case mode, support values: online / dev, default value: online")
)

func main() {
	client.Info("[hot scheduler] starts the hot scheduler test")
	flag.Parse()

	ctx, cancel := context.WithCancel(context.Background())
	cli := client.New(ctx, func(mode string) (clientMode client.Mode) {
		switch client.Mode(mode) {
		case client.ModeOnline:
			clientMode = client.ModeOnline
		case client.ModeDev:
			clientMode = client.ModeDev
		default:
			client.Fatalf("[hot scheduler] Invalid mode: %v,support values: online / dev", mode)
		}
		client.Info("Mode:", clientMode)
		return clientMode
	}(*runMode))

	dbAddr := cli.TiDBService()
	go workload.StartWriteAndRead(dbAddr, *pwd, *database)

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	go func() {
		sig := <-sc
		client.Infof("[hot scheduler] Got signal [%s] to exist.", sig)
		cancel()
		os.Exit(0)
	}()
	ticker := time.NewTicker(15 * time.Second)
	defer ticker.Stop()
	start := time.Now()
	for {
		<-ticker.C
		if time.Since(start) < 10*time.Minute {
			probeSchedule(cli.PdService())
			continue
		}

		if probeSchedule(cli.PdService()) {
			client.Info("[hot scheduler] hot schedule works fine!")
			break
		}

		if time.Since(start) > time.Hour {
			client.Fatal("[hot scheduler] hot schedule not works")
		}
	}
}

const (
	hotPrefix   = "pd/api/v1/hotspot/regions"
	storePrefix = "pd/api/v1/stores"
)

func probeSchedule(pdAddr string) bool {
	storesUrl := fmt.Sprintf("http://%s/%s", pdAddr, storePrefix)
	res, err := http.Get(storesUrl)
	if err != nil {
		client.Fatal(err)
	}
	defer res.Body.Close()
	rawRes, err := ioutil.ReadAll(res.Body)
	stores := types.StoresInfo{}
	json.Unmarshal(rawRes, &stores)

	// write
	hotWriteUrl := fmt.Sprintf("http://%s/%s/write", pdAddr, hotPrefix)
	res, err = http.Get(hotWriteUrl)
	if err != nil {
		client.Fatal(err)
	}
	defer res.Body.Close()
	rawRes, err = ioutil.ReadAll(res.Body)
	stats := types.StoreHotRegionInfos{}
	json.Unmarshal(rawRes, &stats)
	var countDist []float64
	for _, s := range stats.AsPeer {
		countDist = append(countDist, float64(s.RegionsCount))
	}
	client.Infof("[hot scheduler] got write stats: %v", countDist)
	if !isBalanced(countDist, stores.Count) {
		return false
	}

	// read
	hotReadUrl := fmt.Sprintf("http://%s/%s/read", pdAddr, hotPrefix)
	res, err = http.Get(hotReadUrl)
	if err != nil {
		client.Fatal(err)
	}
	defer res.Body.Close()
	rawRes, err = ioutil.ReadAll(res.Body)
	stats = types.StoreHotRegionInfos{}
	json.Unmarshal(rawRes, &stats)
	countDist = countDist[:0]
	for _, s := range stats.AsLeader {
		countDist = append(countDist, float64(s.RegionsCount))
	}
	client.Infof("[hot scheduler] got read stats: %v", countDist)
	return isBalanced(countDist, stores.Count)
}

func isBalanced(dist []float64, total int) bool {
	if len(dist) != total {
		return false
	}
	for _, v := range dist {
		// hotspot should more than 2.
		if v < 2 {
			return false
		}
	}
	vari := stat.Variance(dist, nil)
	client.Infof("[hot scheduler] got variance: %.4f", vari)
	return vari < 2
}