package main

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/pingcap/schrodinger-test/util"
	"github.com/pingcap/schrodinger/agent/core"
	"github.com/pingcap/schrodinger/pkg/test_client"
	"golang.org/x/net/context"
)

func checkAllRegionsAvailable(cli *client.Client) {
	// Check all regions are available or not with a SQL client.
	for {
		time.Sleep(10 * time.Second)

		dbAddr := cli.TiDBService()
		if len(strings.Trim(dbAddr, " ")) == 0 {
			client.Fatalf("[scale tikv] TiDBService returns a blank string, skip")
		} else {
			client.Infof("[scale tikv] get tidb service: %s", dbAddr)
		}

		dbDSN := fmt.Sprintf("root:@tcp(%s)/%s", dbAddr, "test")
		db, err := util.OpenDB(dbDSN, 2)
		if err != nil {
			client.Fatalf("[scale tikv] open db fail: %v", err)
		}

		rows, err := db.Query("select count(*) from sbtest1")
		if err != nil {
			client.Fatalf("[scale tikv] select count(*) from sbtest1 fail: %v", err)
		}
		if rows.Next() {
			var data string
			if err = rows.Scan(&data); err != nil {
				client.Fatalf("[scale tikv] get row data fail: %v", err)
			}
			client.Infof("[scale tikv] count(*): %s", data)
		}
		rows.Close()
		db.Close()
	}
}

func main() {
	client.Info("[scale tikv] start...")

	ctx, cancel := context.WithCancel(context.Background())
	cli := client.New(ctx, client.ModeOnline)

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	go func() {
		sig := <-sc
		client.Infof("[scale tikv] got signal [%s] to exist.", sig)
		cancel()
		os.Exit(0)
	}()

	tikvs := cli.TiKVs()
	client.Infof("tikvs: %v", tikvs)

	go checkAllRegionsAvailable(cli)

	pdService := cli.PdService()
	client.Infof("[scale tikv] get pd: %s", pdService)
	for {
		storeIds := getStoreIds(pdService)
		client.Infof("[scale tikv] get store ids: %v", storeIds)
		if len(storeIds) <= 3 {
			client.Infof("[scale tikv] skip test because there is only %d tikvs", len(storeIds))
			return
		}

		var toDelete = storeIds[0 : len(storeIds)-3]
		client.Infof("[scale tikv] stores to delete: %v", toDelete)

		// Delete stores (including their data) one by one.
		var deletedStores []string
		for _, storeId := range toDelete {
			var agentClient = deleteStore(pdService, storeId)
			deletedStores = append(deletedStores, agentClient)
		}

		// Restart those stores with a empty data directory.
		for _, ip := range deletedStores {
			if err := util.StartPresetProcess(ip, "tikv"); err != nil {
				client.Fatalf("[scale tikv] can't restart tikv: %v", err)
			} else {
				client.Infof("[scale tikv] restart tikv success")
			}
			time.Sleep(5 * time.Minute)
		}
	}
}

func pdCtlRequest(pdService string, command ...string) bytes.Buffer {
	cmdString := []string{"/pd-ctl", "-u", fmt.Sprintf("http://%s", pdService)}
	cmdString = append(cmdString, command...)

	cmd := exec.Command(cmdString[0], cmdString[1:]...)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		client.Errorf("[scale tikv] %s stderr: %s", strings.Join(cmdString, " "), stderr.String())
		client.Fatalf("[scale tikv] pd-ctl fail: %v", err)
	}
	return stdout
}

func getStoreIds(pdService string) []string {
	var stdout bytes.Buffer
	stdout = pdCtlRequest(pdService, "-d", "store")
	reader := bytes.NewReader(stdout.Bytes())
	return extractStoreIds(reader)
}

func getStoreAddress(pdService string, storeId string) (string, string) {
	var stdout bytes.Buffer
	stdout = pdCtlRequest(pdService, "-d", "store", storeId)
	reader := bytes.NewReader(stdout.Bytes())
	return extractStoreAddress(reader)
}

// deleteStore deletes the given store, and returns the store's IP.
func deleteStore(pdService string, storeId string) string {
	client.Infof("[scale tikv] deleting tikv %s from PD", storeId)

	var stdout bytes.Buffer
	stdout = pdCtlRequest(pdService, "-d", "store", "delete", storeId)
	if !strings.HasPrefix(stdout.String(), "Success") {
		client.Fatalf("[scale tikv] delete store %s from PD fail: %s", storeId, stdout.String())
	}

	var tombstone = false
	for i := 0; i < 30; i++ {
		if tombstone {
			break
		}

		time.Sleep(1 * time.Minute)
		stdout = pdCtlRequest(pdService, "-d", "store", storeId)
		reader := bytes.NewReader(stdout.Bytes())
		switch extractStoreState(reader) {
		case "Up":
			client.Fatalf("[scale tikv] the store %s shouldn't be up", storeId)
		case "Offline":
			client.Infof("[scale tikv] the store %s has been deleted, current state: Offline", storeId)
		case "Tombstone":
			client.Infof("[scale tikv] the store %s has been deleted, current state: Tombstone", storeId)
			tombstone = true
		}
	}

	if !tombstone {
		client.Warnf("[scale tikv] store %s is still offline, should be tombstone", storeId)
	}

	ip, _ := getStoreAddress(pdService, storeId)

	// Stop	the process.
	if err := util.StopPresetProcess(ip, "tikv"); err != nil {
		client.Fatalf("[scale tikv] can't stop tikv: %v", err)
	} else {
		client.Infof("[scale tikv] stop tikv success")
	}

	// Delete the data dir.
	var agentClient = util.GenPersistClients([]string{ip})[0]
	res, err := agentClient.ExecCommand(&core.ShellCommand{Command: "rm -rf /var/lib/tikv/data"})
	if err != nil {
		client.Errorf("[scale tikv] rm -rf /var/lib/tikv/data fail: %v", err)
	} else {
		client.Infof("[scale tikv] rm -rf /var/lib/tikv/data success: %s", res)
	}
	return ip
}