package copr_paging

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

	. "github.com/onsi/gomega"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
)

func prepareIndexLookUp(tc resource.TiDBCluster, table string) {
	log.Info("preparing")
	db := getDB(tc)
	defer db.Close()

	MustExec(db, "drop table if exists "+table)
	MustExec(db, "create table "+table+"(id int primary key, c1 int, c2 int, key i(c1))")

	var eg errgroup.Group
	seedBase := time.Now().UnixNano()

	for i := 0; i < prepareConc; i++ {
		seed := seedBase + int64(i)
		i := i
		cnt := prepareCnt / prepareConc
		eg.Go(func() error {
			db := getDB(tc)
			defer func() {
				db.Close()
			}()
			rng := rand.New(rand.NewSource(seed))
			var id int
			for j := 0; j < cnt; j++ {
				id = j*prepareConc + i
				if _, err := db.Exec(fmt.Sprintf("insert into "+table+" values(%d, %d, %d)", id, id, rng.Intn(stateCnt))); err != nil {
					return err
				}
			}
			return nil
		})
	}

	Expect(eg.Wait()).Should(Succeed())
}

func getCurrentRows(tc resource.TiDBCluster, table string) uint64 {
	db := getDB(tc)
	defer db.Close()
	res, err := db.Query("SELECT * FROM " + table + " ORDER BY id DESC LIMIT 1")
	Expect(err).Should(Succeed())
	data, err := ReadAll(res)
	if err != nil {
		panic(err)
	}
	rows, err := strconv.Atoi(data[0][0])
	Expect(err).Should(Succeed())
	return uint64(rows)
}

func testIndexLookUpPaging(caseName string, tc resource.TiDBCluster, table string, write, regionErr, networkErr, lockedErr, podFailure bool) {
	log.Info("starting testing")

	errCh := make(chan error, 2*testConc)
	finishCh := make(chan struct{})
	var wg sync.WaitGroup
	rows := getCurrentRows(tc, table)
	ctx, cancel := context.WithCancel(context.Background())

	wg.Add(testConc)
	for i := 0; i < testConc; i++ {
		go indexLookUpReadVerify(ctx, tc, table, &rows, errCh, finishCh, &wg)
	}

	if write {
		wg.Add(2)
		go indexLookUpWrite(ctx, tc, table, &rows, errCh, finishCh, &wg)
		go indexLookUpWriteLargeTxn(ctx, tc, table, &rows, errCh, finishCh, &wg)
	}

	if regionErr {
		wg.Add(1)
		go doRegionError(tc, errCh, finishCh, &wg)
	}

	if networkErr {
		wg.Add(1)
		go doNetworkError(tc, errCh, finishCh, &wg)
	}

	if lockedErr {
		wg.Add(1)
		go doLockedError(tc, errCh, finishCh, &wg)
	}

	if podFailure {
		wg.Add(1)
		go doPodFailure(tc, errCh, finishCh, &wg)
	}

	closeAndWait := func() {
		log.Info("close and wait for quit")
		close(finishCh)
		cancel()
		wg.Wait()
	}

	timer := time.NewTimer(testDuration)
	select {
	case <-timer.C:
		closeAndWait()
	case err := <-errCh:
		closeAndWait()
		// should fail the test
		Expect(err).ShouldNot(Succeed())
		Expect(err).Should(Succeed())
		return
	}
	log.Info("exit")
}

func indexLookUpWriteStmt(table string, count *uint64) string {
	id := atomic.AddUint64(count, 1)
	return fmt.Sprintf("insert into %s values(%d, %d, %d)", table, id, id, rand.Intn(stateCnt))
}

func indexLookUpReadStmt(table string, count *uint64) string {
	id := atomic.LoadUint64(count)
	var rg int64
	if rand.Intn(5) == 0 {
		// very large range
		rg = rand.Int63n(int64(id))
	} else if rand.Intn(5) == 0 {
		// large range but can go paging
		rg = 100 + rand.Int63n(860)
	} else {
		// small range
		rg = rand.Int63n(100)
	}
	order := "ASC"
	if rand.Intn(2) == 0 {
		// desc order, will meet lock
		order = "DESC"
	}
	return fmt.Sprintf("SELECT * FROM %s WHERE id < %d AND c1 >= %d AND c1 < %d AND c2 in (0, 2, 4, 6, 8) ORDER BY c1 %s LIMIT %d",
		table, id, id-uint64(rg), id, order, rg)
}

func indexLookUpReadVerify(ctx context.Context, tc resource.TiDBCluster, table string, count *uint64, errCh chan<- error, finishCh <-chan struct{}, wg *sync.WaitGroup) {
	db := getDB(tc)
	defer func() {
		db.Close()
		wg.Done()
	}()
	db.SetMaxOpenConns(1)
	for {
		select {
		case <-finishCh:
			return
		default:
		}
		txn, err := db.BeginTx(ctx, nil)
		if err != nil {
			errCh <- err
			return
		}
		for i := 0; i < checkCnt; i++ {
			stmt := indexLookUpReadStmt(table, count)
			if err := execAndCompare(ctx, txn, stmt); err != nil {
				log.Info(fmt.Sprintf("exec stmt failure, stmt: %s", stmt))
				errCh <- errors.Wrapf(err, "stmt: %s", stmt)
				return
			}
		}
		if err := txn.Rollback(); err != nil {
			errCh <- err
			return
		}
	}
}

func indexLookUpWrite(ctx context.Context, tc resource.TiDBCluster, table string, count *uint64, errCh chan<- error, finishCh <-chan struct{}, wg *sync.WaitGroup) {
	db := getDB(tc)
	defer func() {
		db.Close()
		wg.Done()
	}()
	db.SetMaxOpenConns(1)
	db.Exec("set tidb_enable_1pc=off")
	for {
		select {
		case <-finishCh:
			return
		default:
		}
		if _, err := db.ExecContext(ctx, indexLookUpWriteStmt(table, count)); err != nil {
			err = filterError(err)
			if err != nil {
				log.Info("write error", zap.Error(err))
			}
		}
	}
}

func indexLookUpWriteLargeTxn(ctx context.Context, tc resource.TiDBCluster, table string, count *uint64, errCh chan<- error, finishCh <-chan struct{}, wg *sync.WaitGroup) {
	db := getDB(tc)
	defer func() {
		db.Close()
		wg.Done()
	}()
	db.SetMaxOpenConns(1)
	for {
		select {
		case <-finishCh:
			return
		default:
		}
		txn, err := db.BeginTx(ctx, nil)
		if err != nil {
			continue
		}
		for i := 0; i < 1000; i++ {
			if _, err := txn.ExecContext(ctx, indexLookUpWriteStmt(table, count)); err != nil {
				if txn != nil {
					txn.Rollback()
				}
				continue
			}
		}
		err = txn.Commit()
		if err != nil {
			err = filterError(err)
			if err != nil {
				log.Info("write large-txn error", zap.Error(err))
			}
		}
	}
}