package main

import (
	"context"
	"database/sql"
	"fmt"
	"io/ioutil"
	"net/http"
	"sync"
	"time"
	"math/rand"
	"encoding/json"
	"bytes"

	"github.com/juju/errors"
	"github.com/pingcap/kvproto/pkg/debugpb"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"google.golang.org/grpc"
	client "github.com/pingcap/schrodinger/pkg/test_client"
)

var (
	CreateSQL = `DROP TABLE IF EXISTS test;
	CREATE TABLE test (
		id bigint primary key,
		timestamp bigint,
		data0 char(255),
		data1 char(255),
		data2 char(255),
		data3 char(243)
	);`
	InsertSQL   = `INSERT INTO test VALUES (?, ?, "", "", "", "") ON DUPLICATE KEY UPDATE id = ?, timestamp = ?;`
	DeleteSQL   = `DELETE FROM test WHERE id = ?;`
	PointGetSQL = `SELECT * FROM test WHERE id = ?;`
	ScanSQL     = `SELECT * FROM test WHERE id <= ?;`
)

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 RegionsInfo struct {
	Count   int           `json:"count"`
	Regions []*RegionInfo `json:"regions"`
}

type FrequentSchedulingTest struct {
	cli         *client.Client
	db          *sql.DB
	recordCount int
	workerCount int
}

func NewFrequentSchedulingTest(
	cli *client.Client,
	db *sql.DB,
	recordCount int,
	workerCount int,
) *FrequentSchedulingTest {
	return &FrequentSchedulingTest{
		cli:         cli,
		db:          db,
		recordCount: recordCount,
		workerCount: workerCount,
	}
}

func (t *FrequentSchedulingTest) Init(ctx context.Context) error {
	_, err := t.db.Exec(CreateSQL)
	if err != nil {
		return err
	} else {
		client.Info("[frequent_scheduling] Creating table success")
	}
	var wg sync.WaitGroup
	var insertErr error = nil
	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func(count int, round int) {
			for i := 0; i < count; i++ {
				select {
				case <-ctx.Done():
					return
				default:
					now := time.Now().UnixNano()
					id := i + count * round
					_, err := t.db.Exec(InsertSQL, id, now, id, now)
					if err != nil {
						insertErr = err
						return
					}
				}
			}
			wg.Done()
		}(t.recordCount / 100, i)
	}
	wg.Wait()
	if insertErr != nil {
		return err
	}
	client.Infof("[frequent_scheduling] Inserting records success")
	err = t.shuffleRegion()
	if err != nil {
		return err
	}
	client.Infof("[frequent_scheduling] Add shuffle region scheduler")
	return nil
}

func (t *FrequentSchedulingTest) Run(ctx context.Context) error {
	var wg sync.WaitGroup
	jobs := [4]func() (int, error){
		t.randDelete,
		t.randPointGet,
		t.randUpdate,
		t.randScan,
	}
	rand.Seed(time.Now().UnixNano())
	for i := 0; i < t.workerCount; i++ {
		go func(threadID int) {
			wg.Add(1)
			defer wg.Done()
			for {
				select {
				case <-ctx.Done():
					return
				default:
					jobNum := rand.Int() % len(jobs)
					id, err := jobs[jobNum]()
					if err != nil {
						if id > 0 {
							client.Fatalf("[frequent_scheduling] Execute job %d on record %d error: %v", jobNum, id, errors.Trace(err))
						} else {
							client.Fatalf("[frequent_scheduling] Execute job %d error: %v", jobNum, errors.Trace(err))
						}
					}
				}
			}
		}(i)
	}
	go func() {
		wg.Add(1)
		t.randKillServer(ctx, 30*time.Second, 3*time.Minute)
		wg.Done()
	}()
	client.Info("[frequent_scheduling] Running")
	wg.Wait()
	return nil
}

func (t *FrequentSchedulingTest) randPointGet() (int, error) {
	id := rand.Int() % t.recordCount
	err := withRetry(func() error {
		_, err := t.db.Exec(PointGetSQL, id)
		return err
	})
	return id, err
}

func (t *FrequentSchedulingTest) randUpdate() (int, error) {
	id := rand.Int() % t.recordCount
	err := withRetry(func() error {
		now := time.Now().UnixNano()
		_, err := t.db.Exec(InsertSQL, id, now, id, now)
		return err
	})
	return id, err
}

func (t *FrequentSchedulingTest) randDelete() (int, error) {
	id := rand.Int() % t.recordCount
	err := withRetry(func() error {
		_, err := t.db.Exec(DeleteSQL, id)
		return err
	})
	return id, err
}

func (t *FrequentSchedulingTest) randScan() (int, error) {
	id := rand.Int() % t.recordCount
	err := withRetry(func() error {
		_, err := t.db.Exec(ScanSQL, id)
		return err
	})
	return id, err
}

func (t *FrequentSchedulingTest) consistencyCheck() (int, error) {
	regions, err := t.getRegions()
	if err != nil {
		client.Error("[frequent_scheduling] Get RegionsInfo Failed")
		return -1, err
	}
	client.Infof("Get %d regions", regions.Count)
	tikvs := t.cli.TiKVs()
	for _, addr := range tikvs {
		conn, err := grpc.Dial(addr, grpc.WithInsecure())
		if err != nil {
			continue
		}
		c := debugpb.NewDebugClient(conn)
		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
		defer cancel()
		for i := 0; i < regions.Count; i++ {
			req := &debugpb.RegionConsistencyCheckRequest{
				RegionId: regions.Regions[i].ID,
			}
			c.CheckRegionConsistency(ctx, req)
		}
	}
	return -1, nil
}

func (t *FrequentSchedulingTest) randKillServer(
	ctx context.Context,
	recoverInterval time.Duration,
	waitInterval time.Duration,
) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
			time.Sleep(waitInterval)
			client.Info("[frequent_scheduling] Kill a random TiKV server.")
			t.cli.KillWithRecover(context.Background(), "tikv", recoverInterval)
			client.Info("[frequent_scheduling] TiKV server recovered.")
			t.consistencyCheck()
		}
	}
}

func (t *FrequentSchedulingTest) getRegions() (*RegionsInfo, error) {
	reqUrl := fmt.Sprintf("http://%s/pd/api/v1/regions", t.cli.PdService())
	resp, err := http.Get(reqUrl)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	regions := &RegionsInfo{}
	err = json.Unmarshal(body, regions)
	if err != nil {
		return nil, err
	}
	return regions, nil
}

func (t *FrequentSchedulingTest) shuffleRegion() error {
	type PdSchedulersRequest struct {
		Name string `json:"name"`
	}
	reqUrl := fmt.Sprintf("http://%s/pd/api/v1/schedulers", t.cli.PdService())
	req := PdSchedulersRequest{
		Name: "shuffle-region-scheduler",
	}
	reqData, err := json.Marshal(req)
	if err != nil {
		return err
	}
	_, err = http.Post(reqUrl, "application/json", bytes.NewReader(reqData))
	if err != nil {
		return err
	}

	return nil
}

func withRetry(f func() error) error {
	var err error
	for i := 0; i < 3; i++ {
		err = f()
		if err == nil {
			return nil
		}
	}
	return err
}