package main

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/juju/errors"
	"github.com/pingcap/schrodinger/pkg/test_client"
)

const deleteBatch = 100

func checkSplitTableRegion(db *DB) error {
	tableID, err := db.CreateTable()
	if err != nil {
		return errors.Trace(err)
	}

	tableRegions, err := FindTableRegions(tableID)
	if err != nil {
		return errors.Trace(err)
	}

	if len(tableRegions) != 1 {
		return fmt.Errorf("unexpected the number(%d) of regions after create table", len(tableRegions))
	}

	origin := tableRegions[0]

	err = SplitRegion(tableRegions[0].ID)
	if err != nil {
		return errors.Trace(err)
	}

	// wait regions to merge
	for {
		time.Sleep(60 * time.Second)

		now, err := GetRegionInfo(origin.ID)
		if err != nil {
			return errors.Trace(err)
		}
		if now.RegionEpoch.GetVersion() == origin.RegionEpoch.GetVersion()+1 {
			tableRegions, err := FindTableRegions(tableID)
			if err != nil {
				return errors.Trace(err)
			}
			if len(tableRegions) != 2 {
				return fmt.Errorf("unexpected the number(%d) of regions after split", len(tableRegions))
			}
		}

		if now.RegionEpoch.GetVersion() >= origin.RegionEpoch.GetVersion()+2 {
			break
		}
	}

	tableRegions, err = FindTableRegions(tableID)
	if err != nil {
		return errors.Trace(err)
	}

	if len(tableRegions) != 1 {
		return fmt.Errorf("unexpected the number(%d) of regions after split and merge", len(tableRegions))
	}

	now := tableRegions[0]
	if origin.RegionEpoch.GetVersion()+2 != now.RegionEpoch.GetVersion() {
		return fmt.Errorf("unexpected region epoch version after split and merge, before: %d, after: %d",
			origin.RegionEpoch.GetVersion(), now.RegionEpoch.GetVersion())
	}

	return nil
}

func cleanTableData(db *DB, fn func([]*RegionInfo) error) error {
	tableID, err := db.CreateTable()
	if err != nil {
		return errors.Trace(err)
	}

	err = db.InsertUntilSplitToRegions(tableID, 10)
	if err != nil {
		return errors.Trace(err)
	}

	time.Sleep(time.Second * 60)

	tableRegions, err := FindTableRegions(tableID)
	if err != nil {
		return errors.Trace(err)
	}

	err = fn(tableRegions)
	if err != nil {
		return errors.Trace(err)
	}
	return nil
}

func checkDropTable(db *DB) error {
	callback := func(regions []*RegionInfo) error {
		for _, region := range regions {
			nowRegion, err := GetRegionInfo(region.ID)
			if err != nil {
				return errors.Trace(err)
			}
			client.Infof("[region %d] approximate size is %d before drop", nowRegion.ID, nowRegion.ApproximateSize)
		}

		err := db.DropTable()
		if err != nil {
			return errors.Trace(err)
		}

		time.Sleep(time.Second * 60)

		for _, region := range regions {
			nowRegion, err := GetRegionInfo(region.ID)
			if err != nil {
				return errors.Trace(err)
			}
			client.Infof("[region %d] approximate size is %d after drop", nowRegion.ID, nowRegion.ApproximateSize)
		}
		client.Info("------------------------------------------------")

		// wait to merge
		err = db.WaitUntilMerge(regions)
		if err != nil {
			return errors.Trace(err)
		}
		return nil
	}

	return cleanTableData(db, callback)
}

func checkTruncateTable(db *DB) error {
	callback := func(regions []*RegionInfo) error {
		for _, region := range regions {
			nowRegion, err := GetRegionInfo(region.ID)
			if err != nil {
				return errors.Trace(err)
			}
			client.Infof("[region %d] approximate size is %d before truncate", nowRegion.ID, nowRegion.ApproximateSize)
		}

		err := db.TruncateTable()
		if err != nil {
			return errors.Trace(err)
		}

		time.Sleep(time.Second * 60)

		for _, region := range regions {
			nowRegion, err := GetRegionInfo(region.ID)
			if err != nil {
				return errors.Trace(err)
			}
			client.Infof("[region %d] approximate size is %d after truncate", nowRegion.ID, nowRegion.ApproximateSize)
		}
		client.Info("------------------------------------------------")

		// wait to merge
		err = db.WaitUntilMerge(regions)
		if err != nil {
			return errors.Trace(err)
		}
		return nil
	}

	return cleanTableData(db, callback)
}

func checkDropIndex(db *DB) error {
	tableID, err := db.CreateTable()
	if err != nil {
		return errors.Trace(err)
	}

	err = db.InsertUntilSplitToRegions(tableID, 10)
	if err != nil {
		return errors.Trace(err)
	}

	time.Sleep(time.Second * 60)

	client.Infof("add index...")
	err = db.AddIndex()
	if err != nil {
		return errors.Trace(err)
	}
	client.Infof("add index finish...")

	time.Sleep(time.Second * 60)

	tableRegions, err := FindTableRegions(tableID)
	if err != nil {
		return errors.Trace(err)
	}

	// get index regions from table regions
	indexRegions := GetIndexRegions(tableRegions)

	for _, region := range indexRegions {
		nowRegion, err := GetRegionInfo(region.ID)
		if err != nil {
			return errors.Trace(err)
		}
		client.Infof("[region %d] approximate size is %d before drop index", nowRegion.ID, nowRegion.ApproximateSize)
	}

	client.Infof("drop index...")
	err = db.DropIndex()
	if err != nil {
		return errors.Trace(err)
	}
	client.Infof("drop index finish...")

	time.Sleep(time.Second * 60)

	for _, region := range indexRegions {
		nowRegion, err := GetRegionInfo(region.ID)
		if err != nil {
			return errors.Trace(err)
		}
		client.Infof("[region %d] approximate size is %d after drop index", nowRegion.ID, nowRegion.ApproximateSize)
	}
	client.Info("------------------------------------------------")

	// wait to merge
	err = db.WaitUntilMerge(indexRegions)
	if err != nil {
		return errors.Trace(err)
	}
	return nil

}

func checkInsertAndDeleteAll(db *DB) error {
	callback := func(regions []*RegionInfo) error {
		for _, region := range regions {
			nowRegion, err := GetRegionInfo(region.ID)
			if err != nil {
				return errors.Trace(err)
			}
			client.Infof("[region %d] approximate size is %d before delete", nowRegion.ID, nowRegion.ApproximateSize)
		}

		err := db.DeleteAll()
		if err != nil {
			return errors.Trace(err)
		}

		time.Sleep(time.Second * 60)

		for _, region := range regions {
			nowRegion, err := GetRegionInfo(region.ID)
			if err != nil {
				return errors.Trace(err)
			}
			client.Infof("[region %d] approximate size is %d after delete", nowRegion.ID, nowRegion.ApproximateSize)
		}
		client.Info("------------------------------------------------")

		// wait to merge
		err = db.WaitUntilMerge(regions)
		if err != nil {
			return errors.Trace(err)
		}
		return nil
	}

	return cleanTableData(db, callback)
}

func checkInsertAndDeletePartial(db *DB) error {
	callback := func(regions []*RegionInfo) error {
		for _, region := range regions {
			nowRegion, err := GetRegionInfo(region.ID)
			if err != nil {
				return errors.Trace(err)
			}
			client.Infof("[region %d] approximate size is %d before delete", nowRegion.ID, nowRegion.ApproximateSize)
		}

		err := db.DeletePartial()
		if err != nil {
			return errors.Trace(err)
		}

		time.Sleep(time.Second * 60)

		for _, region := range regions {
			nowRegion, err := GetRegionInfo(region.ID)
			if err != nil {
				return errors.Trace(err)
			}
			client.Infof("[region %d] approximate size is %d after delete", nowRegion.ID, nowRegion.ApproximateSize)
		}
		client.Info("------------------------------------------------")

		// wait to merge
		err = db.WaitUntilTimeout(regions, 60*time.Minute)
		if err != nil {
			return errors.Trace(err)
		}

		return nil
	}

	return cleanTableData(db, callback)
}

func checkShuffle(db *DB) error {
	tableID, err := db.CreateTable()
	if err != nil {
		return errors.Trace(err)
	}

	err = db.InsertUntilSplitToRegions(tableID, 10)
	if err != nil {
		return errors.Trace(err)
	}

	tableRegions, err := FindTableRegions(tableID)
	if err != nil {
		return errors.Trace(err)
	}

	// add shuffle scheduler
	err = StartShuffleRegions()
	if err != nil {
		return errors.Trace(err)
	}
	for retry := 0; retry < 5; retry++ {
		ok := false
		for _, region := range tableRegions {
			nowRegion, err := GetRegionInfo(region.ID)
			if err != nil {
				return errors.Trace(err)
			}

			if !IsSameLocation(region.Peers, nowRegion.Peers) {
				// when ensure not in same location, stop shuffle scheduler
				err = StopShuffleRegions()
				if err != nil {
					return errors.Trace(err)
				}
				ok = true
				break
			}
		}

		if ok {
			break
		}
		// wait region peers to move
		time.Sleep(60 * time.Second)
	}

	err = db.DropTable()
	if err != nil {
		return errors.Trace(err)
	}

	// wait to merge
	err = db.WaitUntilMerge(tableRegions)
	if err != nil {
		return errors.Trace(err)
	}

	return nil
}

func checkSpeed(db *DB) error {
	var wg sync.WaitGroup

	for i := 1; i <= 10; i++ {
		tableID, err := db.CreateTable()
		if err != nil {
			continue
		}

		ctx, cancel := context.WithCancel(context.Background())
		go func() {
			time.Sleep(time.Minute * 5)
			cancel()
		}()

		db.InsertIntoNewTable(ctx, 200)

		time.Sleep(time.Second * 60)

		tableRegions, err := FindTableRegions(tableID)
		if err != nil {
			return errors.Trace(err)
		}

		err = db.DropTable()
		if err != nil {
			return errors.Trace(err)
		}

		begin := time.Now()

		go func() {
			wg.Add(1)
			defer wg.Done()

			for {
				time.Sleep(time.Second * 60)

				prevCount := len(tableRegions)
				// remove the merged regions from tableRegions
				for i := len(tableRegions) - 1; i >= 0; i-- {
					region := tableRegions[i]
					nowRegion, err := GetRegionInfo(region.ID)
					if err != nil {
						client.Fatal(err)
					}

					if nowRegion.ID == 0 {
						tableRegions = append(tableRegions[:i], tableRegions[i+1:]...)
						continue
					}

					tableRegions[i] = nowRegion
				}
				if len(tableRegions) <= 1 {
					client.Infof("table %d has merged all regions into one, %v seconds after drop table", tableID, time.Since(begin).Seconds())
					return
				}

				var err error
				tableRegions, err = CheckRegionsTableID(tableID, tableRegions)
				if err != nil {
					client.Fatal(err)
				}

				if len(tableRegions) != prevCount {
					client.Infof("table %d has %d regions, %v seconds after drop table", tableID, len(tableRegions), time.Since(begin).Seconds())
				}
			}
		}()
	}

	wg.Wait()

	return nil
}

func checkCorrectness(db *DB) error {
	initRegions, err := GetRegions()
	if err != nil {
		return errors.Trace(err)
	}

	time.Sleep(time.Minute)

	ctx, cancel := context.WithTimeout(context.Background(), time.Hour*8)
	defer cancel()
	rand.Seed(time.Now().UnixNano())

	var monitorRegions []uint64
	for {
		time.Sleep(time.Second * time.Duration(rand.Intn(60)))

		select {
		case <-ctx.Done():
			return nil
		default:
		}
		nowRegions, err := GetRegions()
		if err != nil {
			client.Fatal(err)
		}

		for i := len(monitorRegions) - 1; i >= 0; i-- {
			if !IsRegionIDExist(nowRegions, monitorRegions[i]) {
				client.Infof("region %v has been merged", monitorRegions[i])
				monitorRegions = append(monitorRegions[:i], monitorRegions[i+1:]...)
				continue
			}
		}

		newRegions, err := GetNewTableRegions(initRegions, nowRegions)
		if err != nil {
			return errors.Trace(err)
		}
		if len(newRegions) == 0 {
			client.Info("no new regions")
			continue
		}
		source := newRegions[rand.Intn(len(newRegions))]
		target := FindAdjacentRegion(nowRegions, source)
		if target == nil {
			client.Warningf("can't find adjacent region of region %v", source.ID)
			continue
		}

		client.Infof("try to add operator merge region %v into region %v", source.ID, target.ID)
		err = MergeRegion(source, target)
		if err != nil {
			client.Warning(err)
		} else {
			monitorRegions = append(monitorRegions, source.ID)
		}
	}
}