package bank

import (
	"context"
	"database/sql"
	"fmt"
	"math/rand"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/juju/errors"
	"github.com/pingcap/log"
	"github.com/rogpeppe/fastuuid"
	"go.uber.org/zap"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

const (
	initialBalance    = 1000
	insertConcurrency = 100
	insertBatchSize   = 100
	maxTransfer       = 100
	systemAccountID   = 0
)

var (
	stmtsCreate = []string{
		`CREATE TABLE IF NOT EXISTS bank2_accounts (
		id INT,
		balance INT NOT NULL,
		name VARCHAR(32),
		remark VARCHAR(2048),
		PRIMARY KEY (id),
		UNIQUE INDEX byName (name),
		KEY byBalance (balance)
	);`,
		`CREATE TABLE IF NOT EXISTS bank2_transaction (
		id INT,
		booking_date TIMESTAMP DEFAULT NOW(),
		txn_date TIMESTAMP DEFAULT NOW(),
		txn_ref VARCHAR(32),
		remark VARCHAR(2048),
		PRIMARY KEY (id),
		UNIQUE INDEX byTxnRef (txn_ref)
	);`,
		`CREATE TABLE IF NOT EXISTS bank2_transaction_leg (
		id INT AUTO_INCREMENT,
		account_id INT,
		amount INT NOT NULL,
		running_balance INT NOT NULL,
		txn_id INT,
		remark VARCHAR(2048),
		PRIMARY KEY (id)
	);`,
		`TRUNCATE TABLE bank2_accounts;`,
		`TRUNCATE TABLE bank2_transaction;`,
		`TRUNCATE TABLE bank2_transaction_leg;`,
	}
	stmtsTruncate = []string{
		`TRUNCATE TABLE bank2_transaction;`,
		`TRUNCATE TABLE bank2_transaction_leg;`,
	}
	tableNames = []string{
		"bank2_accounts",
		"bank2_transaction",
		"bank2_transaction_leg",
	}
	remark = strings.Repeat("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXVZabcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXVZlkjsanksqiszndqpijdslnnq", 16)

	// PadLength returns a random padding length for `remark` field within user
	// specified bound.
	baseLen      = [3]int{36, 48, 16}
	tidbDatabase = true
)

// Config ...
type Config struct {
	// NumAccounts is total accounts
	NumAccounts         int           `toml:"num_accounts"`
	Interval            time.Duration `toml:"interval"`
	Concurrency         int           `toml:"concurrency"`
	RetryLimit          int           `toml:"retry_limit"`
	EnableLongTxn       bool          `toml:"enable_long_txn"`
	Contention          string        `toml:"contention"`
	Pessimistic         bool          `toml:"pessimistic"`
	MinLength           int           `toml:"min_length"`
	MaxLength           int           `toml:"max_length"`
	ReplicaRead         string
	DbName              string
	TiFlashDataReplicas int
}

// ClientCreator ...
type ClientCreator struct {
	Cfg *Config
}

type Bank2Client struct {
	*Config
	wg    sync.WaitGroup
	stop  int32
	txnID int32
}

func (c *Bank2Client) padLength(table int) int {
	minLen := 0
	maxLen := len(remark)
	if table >= 0 && table < 3 {
		if c.Config.MinLength > baseLen[table] {
			minLen = c.Config.MinLength - baseLen[table]
		}
		if c.Config.MaxLength < maxLen+baseLen[table] {
			maxLen = c.Config.MaxLength - baseLen[table]
		}
	}
	return minLen + rand.Intn(maxLen-minLen)
}

//SetUp with only once account verify after finished.
func (c *Bank2Client) SetUp(ctx context.Context, dsn string) error {
	var err error
	defer func() {
		log.Info("[%s] Setup end...", zap.Any("client", c))
	}()

	log.Info("start to prepare count data...")
	db, err := OpenDB(dsn, 1)
	if err != nil {
		log.Fatal("[Bank2Client] create db client error")
	}
	defer db.Close()

	_, err = db.Exec("set @@global.tidb_txn_mode = 'pessimistic';")
	if err != nil {
		log.Fatal("[Bank2Client] set txn_mode failed", zap.Error(err))
	}
	time.Sleep(5 * time.Second)

	defer func() {
		log.Info("init end...")
	}()

	for _, stmt := range stmtsCreate {
		if _, err := db.Exec(stmt); err != nil {
			log.Fatal("execute statement error",
				zap.String("stmt", stmt),
				zap.Error(err))
		}
	}
	if c.TiFlashDataReplicas > 0 {
		// create tiflash replica
		maxSecondsBeforeTiFlashAvail := 1000
		for _, tableName := range tableNames {
			if err := SetAndWaitTiFlashReplica(ctx, db, c.DbName, tableName, c.TiFlashDataReplicas, maxSecondsBeforeTiFlashAvail); err != nil {
				return errors.Trace(err)
			}
		}
	}

	var wg sync.WaitGroup
	type Job struct {
		begin, end int
	}

	ch := make(chan Job)
	for i := 0; i < insertConcurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for job := range ch {
				select {
				case <-ctx.Done():
					return
				default:
				}
				start := time.Now()
				args := make([]string, 0, insertBatchSize)
				for i := job.begin; i < job.end; i++ {
					args = append(args, fmt.Sprintf(`(%d, %d, "account %d", "%s")`, i, initialBalance, i, remark[:c.padLength(0)]))
				}

				query := fmt.Sprintf("INSERT IGNORE INTO bank2_accounts (id, balance, name, remark) VALUES %s", strings.Join(args, ","))
				err := RunWithRetry(ctx, c.Config.RetryLimit, 5*time.Second, func() error {
					_, err := db.Exec(query)
					if IsErrDupEntry(err) {
						return nil
					}
					return err
				})
				if err != nil {
					log.Fatal("exec query err",
						zap.Any("client", c),
						zap.String("query", query),
						zap.Error(err))
				}
				log.Debug("insert accounts",
					zap.Any("client", c),
					zap.Int("number", job.end-job.begin),
					zap.Duration("takes", time.Since(start)))
			}
		}()
	}

	var begin, end int
	for begin = 1; begin <= c.Config.NumAccounts; begin = end {
		end = begin + insertBatchSize
		if end > c.Config.NumAccounts {
			end = c.Config.NumAccounts + 1
		}
		select {
		case <-ctx.Done():
			return nil
		case ch <- Job{begin: begin, end: end}:
		}
	}
	close(ch)
	wg.Wait()

	select {
	case <-ctx.Done():
		log.Warn("initialize is cancelled",
			zap.Any("client", c))
		return nil
	default:
	}
	query := fmt.Sprintf(`INSERT IGNORE INTO bank2_accounts (id, balance, name) VALUES (%d, %d, "system account")`, systemAccountID, int64(c.Config.NumAccounts*initialBalance))
	err = RunWithRetry(ctx, c.Config.RetryLimit, 3*time.Second, func() error {
		_, e := db.Exec(query)
		if IsErrDupEntry(e) {
			return nil
		}
		return e
	})
	if err != nil {
		log.Fatal("insert system account err",
			zap.Any("client", c),
			zap.Error(err))
	}
	c.Verify(ctx, db)
	return nil
}

func (c *Bank2Client) TearDown(ctx context.Context) error {
	atomic.StoreInt32(&c.stop, 1)
	return nil
}

//Start with bank2_transaction and bank2_transaction_leg truncated firstly, so than Start can be invoked more than once after SetUp.
func (c *Bank2Client) Start(ctx context.Context, dsn string) error {
	log.Info("[%s] start to move money...", zap.Any("client", c))
	defer func() {
		log.Info("[%s] test end...", zap.Any("client", c))
	}()
	var err error
	db, err := OpenDB(dsn, c.Config.Concurrency)
	if err != nil {
		log.Fatal("[Bank2Client] create db client error")
	}
	defer db.Close()
	RandomlyChangeReplicaRead(c.String(), c.ReplicaRead, db)

	for _, stmt := range stmtsTruncate {
		if _, err := db.Exec(stmt); err != nil {
			log.Fatal("execute statement error",
				zap.String("stmt", stmt),
				zap.Error(err))
		}
	}
	var wg sync.WaitGroup
	for i := 0; i < c.Config.Concurrency; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			for {
				select {
				case <-ctx.Done():
					return
				default:
				}
				if atomic.LoadInt32(&c.stop) != 0 {
					log.Error("stopped", zap.Any("client", c))
					return
				}
				c.wg.Add(1)
				c.moveMoney(db)
				c.wg.Done()
			}
		}(i)
	}
	wg.Wait()
	return nil
}

//StartVerify add support for read with some session parameters, so that you can query the bank account with any setting needed.
func (c *Bank2Client) StartVerify(ctx context.Context, dsn string, sessionParams ...string) {
	db, err := OpenDB(dsn, 1)
	if err != nil {
		log.Fatal("[Bank2Client] create db client error")
	}
	defer db.Close()

	c.Verify(ctx, db)
	go func(ctx context.Context) {
		for {
			select {
			case <-ctx.Done():
				return
			case <-time.After(c.Config.Interval):
				c.Verify(ctx, db, sessionParams...)
			}
		}
	}(ctx)
}

//Verify add support for read with some session parameters, so that you can query the bank account with any setting needed.
func (c *Bank2Client) Verify(ctx context.Context, db *sql.DB, sessionParams ...string) {
	log.Debug("Begin to verify")
	tx, err := db.Begin()
	if err != nil {
		_ = errors.Trace(err)
		return
	}
	defer func() {
		_ = tx.Rollback()
	}()

	if tidbDatabase {
		var tso uint64
		if err = tx.QueryRow("SELECT @@tidb_current_ts").Scan(&tso); err == nil {
			log.Info("[%s] SELECT SUM(balance) to verify use tso %d",
				zap.Any("client", c),
				zap.Uint64("tso", tso))
		}
	}

	var total int64
	expectTotal := (int64(c.Config.NumAccounts) * initialBalance) * 2

	// read from tikv
	_, err = tx.Exec("set @@session.tidb_isolation_read_engines='tikv'")
	if err != nil {
		log.Warn("", zap.Error(err))
		return
	}

	for _, param := range sessionParams {
		_, err = tx.Exec(param)
		if err != nil {
			log.Warn("", zap.Error(err))
			return
		}
	}

	// query with IndexScan
	uuid := fastuuid.MustNewGenerator().Hex128()
	query := fmt.Sprintf("SELECT SUM(balance) AS total, '%s' as uuid FROM bank2_accounts", uuid)
	if err = tx.QueryRow(query).Scan(&total, &uuid); err != nil {
		_ = errors.Trace(err)
		log.Error(err.Error())
		return
	}
	if total != expectTotal {
		log.Error("[%s] bank2_accounts total should be %d, but got %d, query uuid is %s",
			zap.Any("client", c),
			zap.Int64("expected", expectTotal),
			zap.Int64("total", total),
			zap.String("uuid", uuid))
		atomic.StoreInt32(&c.stop, 1)
		c.wg.Wait()
		log.Error("[%s] bank2_accounts total should be %d, but got %d, query uuid is %s",
			zap.Any("client", c),
			zap.Int64("expected", expectTotal),
			zap.Int64("total", total),
			zap.String("uuid", uuid))
	} else {
		log.Info(fmt.Sprintf("check succeeded,total equal to expectTotal %d", total))
	}

	// query with TableScan
	uuid = fastuuid.MustNewGenerator().Hex128()
	query = fmt.Sprintf("SELECT SUM(balance) AS total, '%s' as uuid FROM bank2_accounts ignore index(byBalance)", uuid)
	if err = tx.QueryRow(query).Scan(&total, &uuid); err != nil {
		_ = errors.Trace(err)
		return
	}
	if total != expectTotal {
		log.Error("[%s] bank2_accounts total should be %d, but got %d, query uuid is %s",
			zap.Any("client", c),
			zap.Int64("expected", expectTotal),
			zap.Int64("total", total),
			zap.String("uuid", uuid))
		atomic.StoreInt32(&c.stop, 1)
		c.wg.Wait()
		log.Error("[%s] bank2_accounts total should be %d, but got %d, query uuid is %s",
			zap.Any("client", c),
			zap.Int64("expected", expectTotal),
			zap.Int64("total", total),
			zap.String("uuid", uuid))
	}

	// query with tiflash
	if c.TiFlashDataReplicas > 0 {
		_, err = tx.Exec("set @@session.tidb_isolation_read_engines='tiflash'")
		if err != nil {
			log.Warn("", zap.Error(err))
			return
		}

		readTiFlash := func() error {
			uuid := fastuuid.MustNewGenerator().Hex128()
			query := fmt.Sprintf("SELECT SUM(balance) AS total, '%s' as uuid FROM bank2_accounts", uuid)
			if err = tx.QueryRow(query).Scan(&total, &uuid); err != nil {
				log.Error("[%s] select sum error %v",
					zap.Any("client", c),
					zap.Error(err))
				return errors.Trace(err)
			}
			return nil
		}
		err = RunWithRetry(ctx, c.RetryLimit, time.Second, readTiFlash)
		if err != nil {
			_ = err
			return
		}
		if total != expectTotal {
			log.Error("[%s] bank2_accounts total should be %d, but tiflash got %d, query uuid is %s",
				zap.Any("client", c),
				zap.Int64("expected", expectTotal),
				zap.Int64("total", total),
				zap.String("uuid", uuid))
			atomic.StoreInt32(&c.stop, 1)
			c.wg.Wait()
			log.Error("[%s] bank2_accounts total should be %d, but tiflash got %d, query uuid is %s",
				zap.Any("client", c),
				zap.Int64("expected", expectTotal),
				zap.Int64("total", total),
				zap.String("uuid", uuid))
		}

		_, err = tx.Exec("set @@session.tidb_isolation_read_engines='tikv'")
		if err != nil {
			log.Warn("", zap.Error(err))
			return
		}
	}

	if _, err := tx.Exec("ADMIN CHECK TABLE bank2_accounts"); err != nil {
		log.Error("[%s] ADMIN CHECK TABLE bank2_accounts fails: %v",
			zap.Any("client", c),
			zap.Error(err))
		errStr := err.Error()
		if strings.Contains(errStr, "1105") &&
			!(strings.Contains(errStr, "cancelled DDL job") ||
				strings.Contains(errStr, "Information schema is changed") ||
				strings.Contains(errStr, "TiKV server timeout") ||
				strings.Contains(errStr, "redirect failed") ||
				strings.Contains(errStr, "no leader") ||
				strings.Contains(errStr, "injected") ||
				strings.Contains(errStr, "context deadline exceeded") ||
				strings.Contains(errStr, "redirect to not leader") ||
				strings.Contains(errStr, "raft: stopped") ||
				strings.Contains(errStr, "raft proposal dropped") ||
				strings.Contains(errStr, "no available connections") ||
				(errStr == "Error 1105: ")) {
			atomic.StoreInt32(&c.stop, 1)
			c.wg.Wait()
			log.Fatal("[%s] ADMIN CHECK TABLE bank2_accounts fails: %v",
				zap.Any("client", c),
				zap.Error(err))
		}
	}
}

func (c *Bank2Client) moveMoney(db *sql.DB) {
	from, to := rand.Intn(c.Config.NumAccounts), rand.Intn(c.Config.NumAccounts)
	if c.Config.Contention == "high" {
		// Use the first account number we generated as a coin flip to
		// determine whether we're transferring money into or out of
		// the system account.
		if from > c.Config.NumAccounts/2 {
			from = systemAccountID
		} else {
			to = systemAccountID
		}
	}
	if from == to {
		return
	}
	amount := rand.Intn(maxTransfer)
	// start := time.Now()
	if err := c.execTransaction(db, from, to, amount); err != nil {
		// bank2VerifyFailedCounter.Inc()
		log.Error("[%s] move money err %v",
			zap.Any("client", c),
			zap.Error(err))
		return
	}
	// bank2VerifyDuration.Observe(time.Since(start).Seconds())
}

func (c *Bank2Client) execTransaction(db *sql.DB, from, to int, amount int) error {
	tx, err := db.Begin()
	if err != nil {
		return errors.Trace(err)
	}
	defer func() {
		_ = tx.Rollback()
	}()

	rows, err := tx.Query(fmt.Sprintf("SELECT id, balance FROM bank2_accounts WHERE id IN (%d, %d) FOR UPDATE", from, to))
	if err != nil {
		return errors.Trace(err)
	}
	defer rows.Close()

	var (
		fromBalance int
		toBalance   int
		count       int
	)

	for rows.Next() {
		var id, balance int
		if err = rows.Scan(&id, &balance); err != nil {
			return errors.Trace(err)
		}
		switch id {
		case from:
			fromBalance = balance
		case to:
			toBalance = balance
		default:
			log.Fatal("[%s] got unexpected account %d",
				zap.Any("client", c),
				zap.Int("id", id))
		}
		count++
	}

	if err = rows.Err(); err != nil {
		return errors.Trace(err)
	}

	if count != 2 {
		log.Fatal("[%s] select %d(%d) -> %d(%d) invalid count %d",
			zap.Any("client", c),
			zap.Int("from", from),
			zap.Int("fromBalance", fromBalance),
			zap.Int("to", to),
			zap.Int("toBalance", toBalance),
			zap.Int("count", count))
	}

	if fromBalance < amount {
		return nil
	}

	insertTxn := `INSERT INTO bank2_transaction (id, txn_ref, remark) VALUES (?, ?, ?)`
	insertTxnLeg := `INSERT INTO bank2_transaction_leg (account_id, amount, running_balance, txn_id, remark) VALUES (?, ?, ?, ?, ?)`
	updateAcct := `UPDATE bank2_accounts SET balance = ? WHERE id = ?`
	txnID := atomic.AddInt32(&c.txnID, 1)
	if _, err := tx.Exec(insertTxn, txnID, fmt.Sprintf("txn %d", txnID), remark[:c.padLength(1)]); err != nil {
		_ = tx.Rollback()
		return errors.Trace(err)
	}
	if _, err = tx.Exec(insertTxnLeg, from, -amount, fromBalance-amount, txnID, remark[:c.padLength(2)]); err != nil {
		_ = tx.Rollback()
		return errors.Trace(err)
	}
	if _, err = tx.Exec(insertTxnLeg, to, amount, toBalance+amount, txnID, remark[:c.padLength(2)]); err != nil {
		_ = tx.Rollback()
		return errors.Trace(err)
	}
	if _, err = tx.Exec(updateAcct, toBalance+amount, to); err != nil {
		_ = tx.Rollback()
		return errors.Trace(err)
	}
	if _, err = tx.Exec(updateAcct, fromBalance-amount, from); err != nil {
		_ = tx.Rollback()
		return errors.Trace(err)
	}

	err = tx.Commit()
	if err != nil {
		log.Info("[%s] %d(%d) failed to transfer %d to %d(%d), err: %v",
			zap.Any("client", c),
			zap.Int("from", from),
			zap.Int("fromBalance", fromBalance),
			zap.Int("amount", amount),
			zap.Int("to", to),
			zap.Int("toBalance", toBalance),
			zap.Error(err))
	} else {
		log.Debug("[%s] %d(%d) transfer %d to %d(%d)",
			zap.Any("client", c),
			zap.Int("from", from),
			zap.Int("fromBalance", fromBalance),
			zap.Int("amount", amount),
			zap.Int("to", to),
			zap.Int("toBalance", toBalance))
	}
	return nil
}

// String ...
func (c *Bank2Client) String() string {
	return "bank2"
}