package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"sync/atomic"
	"syscall"
	"math/rand"
	"time"

	client "github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/pd/pkg/typeutil"
	"github.com/5kbpers/client-go/rawkv"
	"github.com/5kbpers/client-go/config"
)

var (
	readConcurrency = flag.Int("read_concurrency", 8, "read thread number")
	writeConcurrency = flag.Int("write_concurrency", 1, "write thread number")
	recordCount = flag.Uint64("record_count", 10000, "record count")
	mode client.Mode = client.ModeOnline
)

func main() {
	client.Info("[follower read]start follower read test")
	flag.Parse()

	ctx, cancel := context.WithCancel(context.Background())
	cli := client.New(ctx, mode)
	tikvCli, err := rawkv.NewClient(cli.PDs(), config.Default())
	if err != nil {
		panic(fmt.Sprintf("[follower read] create tikv client error: %v", err))
	}

	var wg sync.WaitGroup
	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	go func(wg *sync.WaitGroup) {
		wg.Add(1)
		defer wg.Done()
		sig := <-sc
		client.Infof("[follower read] got signal [%d] to exit.", sig)
		cancel()
		os.Exit(0)
	}(&wg)

	for i := uint64(0); i < *recordCount; i++ {
		key := fmt.Sprintf("k%d", i)
		value := fmt.Sprintf("%d", 0)
		err := tikvCli.Put([]byte(key), []byte(value))
		if err != nil {
			client.Infof("[follower read] init data error at key %s: %v", key, err)
		}
	}
	client.Info("[follower read] complete data loading")

	go func() {
		cli.SetNetworkWithRecover(ctx, "tikv", "fail", "100%", 3 * time.Minute)
		time.Sleep(time.Minute * 30)
	}()

	for i := 0; i < *readConcurrency; i++ {
		go func(wg *sync.WaitGroup) {
			wg.Add(1)
			defer wg.Done()
			cfg := config.Default()
			cfg.RPC.ReadTimeoutShort = time.Minute * 30
			tikvCli, err := rawkv.NewClient(cli.PDs(), cfg)
			if err != nil {
				panic(fmt.Sprintf("[follower read] create tikv client error: %v", err))
			}
			r := rand.New(rand.NewSource(time.Now().UnixNano()))
			for {
				select {
				case <-ctx.Done():
					return
				default:
					err := getRandData(tikvCli, r)
					if err != nil {
						client.Fatalf("[follower read] Get data error: %v", err)
						return
					}
				}
			}
		}(&wg)
	}

	var counter uint64 = 0
	for i := 0; i < *writeConcurrency; i++ {
		go func(wg *sync.WaitGroup) {
			wg.Add(1)
			defer wg.Done()
			tikvCli, err := rawkv.NewClient(cli.PDs(), config.Default())
			if err != nil {
				panic(fmt.Sprintf("[follower read] create tikv client error: %v", err))
			}
			r := rand.New(rand.NewSource(time.Now().UnixNano()))
			for {
				select {
				case <-ctx.Done():
					return
				default:
					putRandData(tikvCli, r, &counter)
				}
			}
		}(&wg)
	}

	wg.Wait()
}

type RegionInfo struct {
	ID          uint64              `json:"id"`
	StartKey    string              `json:"start_key"`
	EndKey      string              `json:"end_key"`
	RegionEpoch *metapb.RegionEpoch `json:"epoch,omitempty"`
	Peers       []*metapb.Peer      `json:"peers,omitempty"`

	Leader          *metapb.Peer      `json:"leader,omitempty"`
	DownPeers       []*pdpb.PeerStats `json:"down_peers,omitempty"`
	PendingPeers    []*metapb.Peer    `json:"pending_peers,omitempty"`
	WrittenBytes    uint64            `json:"written_bytes,omitempty"`
	ReadBytes       uint64            `json:"read_bytes,omitempty"`
	ApproximateSize int64             `json:"approximate_size,omitempty"`
	ApproximateKeys int64             `json:"approximate_keys,omitempty"`
}

type MetaStore struct {
	*metapb.Store
	StateName string `json:"state_name"`
}

// StoreStatus contains status about a store.
type StoreStatus struct {
	Capacity           typeutil.ByteSize  `json:"capacity,omitempty"`
	Available          typeutil.ByteSize  `json:"available,omitempty"`
	LeaderCount        int                `json:"leader_count,omitempty"`
	LeaderWeight       float64            `json:"leader_weight,omitempty"`
	LeaderScore        float64            `json:"leader_score,omitempty"`
	LeaderSize         int64              `json:"leader_size,omitempty"`
	RegionCount        int                `json:"region_count,omitempty"`
	RegionWeight       float64            `json:"region_weight,omitempty"`
	RegionScore        float64            `json:"region_score,omitempty"`
	RegionSize         int64              `json:"region_size,omitempty"`
	SendingSnapCount   uint32             `json:"sending_snap_count,omitempty"`
	ReceivingSnapCount uint32             `json:"receiving_snap_count,omitempty"`
	ApplyingSnapCount  uint32             `json:"applying_snap_count,omitempty"`
	IsBusy             bool               `json:"is_busy,omitempty"`
	StartTS            *time.Time         `json:"start_ts,omitempty"`
	LastHeartbeatTS    *time.Time         `json:"last_heartbeat_ts,omitempty"`
	Uptime             *typeutil.Duration `json:"uptime,omitempty"`
}

// StoreInfo contains information about a store.
type StoreInfo struct {
	Store  *MetaStore   `json:"store"`
	Status *StoreStatus `json:"status"`
}

func putRandData(cli *rawkv.Client, r *rand.Rand, counter *uint64) {
	key := generateRandKey(r)
	atomic.AddUint64(counter, 1)
	value := atomic.LoadUint64(counter)
	err := cli.Put([]byte(key), []byte(fmt.Sprintf("%d", value)))
	if err != nil {
		client.Errorf("[follower read] put data error at key %s: %v", key, err)
	}
}

func getRandData(cli *rawkv.Client, r *rand.Rand) error {
	key := generateRandKey(r)
	leaderResp, err := cli.Get([]byte(key), false)
	if err != nil {
		if err == context.DeadlineExceeded {
			client.Fatalf("[follower read] get leader value failed in key %s: %v", key, err)
		}
		return nil
	}
	leaderValue, err := strconv.ParseUint(string(leaderResp), 10, 64)
	if err != nil {
		client.Fatalf("[follower read] parse leader value failed in key %s: %v", key, err)
	}

	followerResp, err := cli.Get([]byte(key), true)
	if err != nil {
		client.Fatalf("[follower read] get follower value failed in key %s: %v", key, err)
	}
	followerValue, err := strconv.ParseUint(string(followerResp), 10, 64)
	if err != nil {
		client.Fatalf("[follower read] parse follower value failed in key %s: %v", key, err)
	}

	if leaderValue > followerValue {
		client.Fatalf("follower read older value: leader[%d], follower[%d]", leaderValue, followerValue)
	} else if leaderValue < followerValue {
		client.Infof("follower read newer value: leader[%d], follower[%d]", leaderValue, followerValue)
	}

	return nil
}

func generateRandKey(r *rand.Rand) string {
	return fmt.Sprintf("k%d", r.Uint64() % *recordCount)
}