package main

import (
	"context"
	"database/sql"
	"github.com/pingcap/errors"
	"github.com/pingcap/parser/mysql"
	"github.com/pingcap/schrodinger-test/util"
	"math/rand"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	mysqlDriver "github.com/go-sql-driver/mysql"
	client "github.com/pingcap/schrodinger/pkg/test_client"
)

type TxnKeepAliveCase struct {
	dbDsn     string
	tableSize int
	threads   int
	db        *sql.DB

	// To make sure different threads runs tests on exclusive rows
	rowLock []sync.Mutex

	// Used to assign each single execution a unique id, so that it will be printed to the log and helps debugging.
	executionCount int32
}

const initialBalance = 10000

func NewTxnKeepAliveCase(dbDsn string, tableSize int, threads int) *TxnKeepAliveCase {
	return &TxnKeepAliveCase{
		dbDsn:     dbDsn,
		tableSize: tableSize,
		threads:   threads,
		rowLock:   nil,
	}
}

func (c *TxnKeepAliveCase) Initialize(ctx context.Context) error {
	client.Info("[txn_keep_alive] Initializing...")
	db, err := c.openDB(1)
	defer db.Close()
	if err != nil {
		return errors.Trace(err)
	}

	client.Info("[txn_keep_alive] Set txn mode to pessimistic mode...")
	_, err = db.Exec(`set @@global.tidb_txn_mode = "pessimistic"`)
	if err != nil {
		return errors.Trace(err)
	}

	updateTxnModeTime := time.Now()

	// Initialize the table
	_, err = db.Exec("drop table if exists accounts")
	if err != nil {
		return errors.Trace(err)
	}
	_, err = db.Exec("create table accounts(id integer, balance integer, primary key (id));")
	if err != nil {
		return errors.Trace(err)
	}

	// Batch insertion to reduce time to load the data
	client.Info("[txn_keep_alive] Loading data to the test table...")
	const insertBatchSize = 256

	txn, err := db.Begin()
	if err != nil {
		return errors.Trace(err)
	}

	currentBatchSize := 0
	for i := 0; i < c.tableSize; i++ {
		_, err = txn.ExecContext(ctx, "insert into accounts values(?, ?)", i, initialBalance)
		if err != nil {
			return errors.Trace(err)
		}
		currentBatchSize++
		if currentBatchSize >= insertBatchSize {
			err = txn.Commit()
			if err != nil {
				return errors.Trace(err)
			}
			currentBatchSize = 0
			txn, err = db.Begin()
			if err != nil {
				return errors.Trace(err)
			}
		}
	}
	if currentBatchSize > 0 {
		err = txn.Commit()
		if err != nil {
			return errors.Trace(err)
		}
	} else {
		_ = txn.Rollback()
	}

	// The updated global variable will need some time to take effect. Wait for it.
	remainWaitTime := time.Second*5 - time.Since(updateTxnModeTime)
	if remainWaitTime > 0 {
		time.Sleep(remainWaitTime)
	}

	c.rowLock = make([]sync.Mutex, c.tableSize)
	c.db, err = c.openDB(c.threads * 4)
	client.Info("[txn_keep_alive] Finished initializing")
	return nil
}

func (c *TxnKeepAliveCase) Execute(ctx context.Context) error {
	defer c.db.Close()

	client.Info("[txn_keep_alive] Start testing...")
	errChan := make(chan error, c.threads)
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	for i := 0; i < c.threads; i++ {
		go func() { errChan <- c.run(ctx) }()
	}

	go func() { errChan <- c.checkConsistency(ctx) }()

	for i := 0; i < c.threads; i++ {
		err := <-errChan
		if err != nil {
			return errors.Trace(err)
		}
	}

	return nil
}

func (c *TxnKeepAliveCase) run(ctx context.Context) error {

	for {
		select {
		case <-ctx.Done():
			return nil
		default:
		}

		err := c.runOnce(ctx)
		if err != nil {
			return errors.Trace(err)
		}
	}
}

func (c *TxnKeepAliveCase) runOnce(ctx context.Context) error {
	execId := atomic.AddInt32(&c.executionCount, 1)
	client.Infof("[txn_keep_alive %v] Start execution", execId)
	defer client.Infof("[txn_keep_alive %v] Exited", execId)

	client.Infof("[txn_keep_alive %v] Allocate random rows for test", execId)
	// There are a common key in the two transactions.
	// t1Ids: [a, b]
	// t2Ids: [a, c]
	allIds := c.randIDs(3)

	t1Ids := []int{allIds[0], allIds[1]}
	t2Ids := []int{allIds[0], allIds[2]}

	// Exclude other threads from these affected rows.
	c.lockRows(allIds)
	defer c.unlockRows(allIds)

	client.Infof("[txn_keep_alive %v] Begin transactions", execId)

	t1, err := c.db.Begin()
	if err != nil {
		return errors.Trace(err)
	}
	defer t1.Rollback()

	t2, err := c.db.Begin()
	if err != nil {
		return errors.Trace(err)
	}
	defer t2.Rollback()

	var t1ConnId, t2ConnId int
	err = t1.QueryRow("select connection_id()").Scan(&t1ConnId)
	if err != nil {
		return errors.Trace(err)
	}
	err = t2.QueryRow("select connection_id()").Scan(&t2ConnId)
	if err != nil {
		return errors.Trace(err)
	}

	shuffle(t1Ids)
	shuffle(t2Ids)

	client.Infof("[txn_keep_alive %v] t1(conn %v) runs on %v, t2(conn %v) runs on %v", execId, t1ConnId, t1Ids, t2ConnId, t2Ids)

	// Execute transfer on transaction t1. rand(1000) + 1 to make sure it transfers at least 1.
	errCh := c.asyncTransfer(ctx, execId, t1, t1Ids[0], t1Ids[1], rand.Intn(1000)+1)
	select {
	case err = <-errCh:
		if err != nil {
			return errors.Trace(err)
		}
	case <-time.After(time.Second * 5):
		return errors.Errorf("Transaction T1 which shouldn't meet lock lasts for over 5s")
	}
	// errCh consumed here

	// Execute transfer on transaction t2, however one of the keys is locked by t1.
	errCh = c.asyncTransfer(ctx, execId, t2, t2Ids[0], t2Ids[1], rand.Intn(1000)+1)
	// Wait for t1 for a random time (1m to 5m, greater than default TTL), and t1 will enlarge it's TTL continuously
	waitTime := time.Second * time.Duration(rand.Intn(540)+60)
	select {
	case err = <-errCh:
		if err != nil {
			return errors.Trace(err)
		}
	case <-time.After(waitTime):
	}

	const (
		actionCommit = iota
		actionRollback
		actionKill
	)

	availableActions := []int{actionCommit, actionRollback, actionKill}
	// Random select what action to do on transaction t1
	action := availableActions[rand.Intn(len(availableActions))]

	switch action {
	case actionCommit:
		client.Infof("[txn_keep_alive %v] Committing t1", execId)
		err = t1.Commit()
		if err != nil {
			return errors.Trace(err)
		}
	case actionRollback:
		client.Infof("[txn_keep_alive %v] Rolling back t1", execId)
		err = t1.Rollback()
		if err != nil {
			return errors.Trace(err)
		}
	case actionKill:
		client.Infof("[txn_keep_alive %v] Killing t1", execId)
		_, err = c.db.Exec("kill tidb " + strconv.FormatInt(int64(t1ConnId), 10))
		if err != nil {
			return errors.Trace(err)
		}
	default:
		return errors.Errorf("invalid action type: %+v", action)
	}

	// Then T2 should continue.
	select {
	case err = <-errCh:
		if err != nil {
			return errors.Trace(err)
		}
	case <-time.After(time.Second * 30):
		return errors.Errorf("%v: transaction t2 didn't finish", execId)
	}

	client.Infof("[txn_keep_alive %v] Committing t2", execId)
	err = t2.Commit()
	if err != nil {
		return errors.Trace(err)
	}

	client.Infof("[txn_keep_alive %v] Success", execId)
	return nil
}

func (c *TxnKeepAliveCase) asyncTransfer(ctx context.Context, execId int32, txn *sql.Tx, from int, to int, amount int) chan error {
	ch := make(chan error, 1)
	go func() {
		for {
			_, err := txn.ExecContext(ctx, "update accounts set balance = balance - ? where id = ?", amount, from)
			if err != nil {
				if isResolveLockTimeout(err) {
					client.Debugf("[txn_keep_alive %v] Resolve lock timeout. Retry. err: %v", execId, err)
					continue
				}
				ch <- errors.Trace(err)
				return
			}
			break
		}

		for {
			_, err := txn.ExecContext(ctx, "update accounts set balance = balance + ? where id = ?", amount, to)
			if err != nil {
				if isResolveLockTimeout(err) {
					client.Debugf("[txn_keep_alive %v] Resolve lock timeout. Retry. err: %v", execId, err)
					continue
				}
				ch <- errors.Trace(err)
				return
			}
			break
		}

		ch <- nil
	}()
	return ch
}

func (c *TxnKeepAliveCase) checkConsistency(ctx context.Context) error {
	tick := time.NewTicker(time.Minute * 5)
	for {
		select {
		case <-ctx.Done():
			return nil
		case <-tick.C:
		}

		err := c.checkConsistencyOnce(ctx)
		if err != nil {
			return errors.Trace(err)
		}
	}
}

func (c *TxnKeepAliveCase) checkConsistencyOnce(ctx context.Context) error {
	var totalBalance int
	err := c.db.QueryRow("select sum(balance) from accounts").Scan(&totalBalance)
	if err != nil {
		client.Errorf("[txn_keep_alive] Check consistency failed on executing query: %+v", err)
		return errors.Trace(err)
	}

	expectBalance := c.tableSize * initialBalance
	if totalBalance != c.tableSize*initialBalance {
		client.Errorf("[txn_keep_alive] Check consistency failed: expected total balance %v, got %v", expectBalance, totalBalance)
		return errors.Errorf("check consistency failed: lost consistency")
	}

	client.Infof("[txn_keep_alive] Check consistency passed")
	return nil
}

func (c *TxnKeepAliveCase) openDB(conns int) (*sql.DB, error) {
	db, err := sql.Open("mysql", c.dbDsn)
	if err != nil {
		return nil, err
	}

	db.SetMaxOpenConns(conns)

	client.Infof("[txn_keep_alive] Open DB: %v", c.dbDsn)
	return db, nil
}

func (c *TxnKeepAliveCase) lockRows(rows []int) {
	sort.Slice(rows, func(i, j int) bool {
		return rows[i] < rows[j]
	})
	for _, row := range rows {
		c.rowLock[row].Lock()
	}
}

func (c *TxnKeepAliveCase) unlockRows(rows []int) {
	sort.Slice(rows, func(i, j int) bool {
		return rows[i] < rows[j]
	})
	for _, row := range rows {
		c.rowLock[row].Unlock()
	}
}

// randIDs picks `count` distinct ids in the range `[0, c.tableSize)`
func (c *TxnKeepAliveCase) randIDs(count int) []int {
	return util.RandDistinct(count, c.tableSize)
}

func shuffle(values []int) {
	rand.Shuffle(2, func(i, j int) {
		values[i], values[j] = values[j], values[i]
	})
}

func isResolveLockTimeout(err error) bool {
	if mysqlErr, ok := err.(*mysqlDriver.MySQLError); ok {
		if mysqlErr.Number == mysql.ErrResolveLockTimeout {
			return true
		}
	}
	return false
}