package main

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

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

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 SUM(id) FROM test;`
)

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 HotspotTest struct {
	cli         *client.Client
	db          *sql.DB
	recordCount int
	workerCount int
}

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

func (t *HotspotTest) Init() error {
	_, err := t.db.Exec(CreateSQL)
	if err != nil {
		return err
	} else {
		client.Info("[hotspot] Creating table success")
	}
	for i := 0; i < t.recordCount; i++ {
		now := time.Now().UnixNano()
		_, err := t.db.Exec(InsertSQL, i, now, i, now)
		if err != nil {
			return err
		}
	}
	client.Infof("[hotspot] Inserting records success")
	return nil
}

func (t *HotspotTest) Run(ctx context.Context) {
	var wg sync.WaitGroup
	jobs := [4]func() (int, error){
		t.randDelete,
		t.randPointGet,
		t.randUpdate,
		t.scan,
	}

	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("[hotspot] Execute job %d on record %d error: %v", jobNum, id, errors.Trace(err))
						} else {
							client.Fatalf("[hotspot] 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("[hotspot] Running")

	wg.Wait()
}

func (t *HotspotTest) 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 *HotspotTest) 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 *HotspotTest) 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 *HotspotTest) scan() (int, error) {
	err := withRetry(func() error {
		_, err := t.db.Exec(ScanSQL)
		return err
	})
	return -1, err
}

func (t *HotspotTest) consistencyCheck() (int, error) {
	regions, err := t.getRegions()
	if err != nil {
		client.Error("[hotspot] 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 *HotspotTest) randKillServer(
	ctx context.Context,
	recoverInterval time.Duration,
	waitInterval time.Duration,
) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
			time.Sleep(waitInterval)
			client.Info("[hotspot] Kill a random TiKV server.")
			t.cli.KillWithRecover(context.Background(), "tikv", recoverInterval)
			client.Info("[hotspot] TiKV server recovered.")
			t.consistencyCheck()
		}
	}
}

func (t *HotspotTest) getRegions() (*RegionsInfo, error) {
	req_url := fmt.Sprintf("http://%s/pd/api/v1/regions", t.cli.PdService())
	resp, err := http.Get(req_url)
	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 withRetry(f func() error) error {
	var err error
	for i := 0; i < 3; i++ {
		err = f()
		if err == nil {
			return nil
		}
	}
	return err
}