package main

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

	"github.com/juju/errors"
	"github.com/pingcap/schrodinger-test/util"
	client "github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/rogpeppe/fastuuid"
	"golang.org/x/net/context"
)

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

// remark length is 2KB
var remark = strings.Repeat("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXVZabcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXVZlkjsanksqiszndqpijdslnnq", 16)

// Bank2CaseConfig is for bank2 test case.
type Bank2CaseConfig struct {
	// NumAccounts is total accounts
	NumAccounts int           `toml:"num_accounts"`
	Interval    time.Duration `toml:"interval"`
	Contention  string        `toml:"contention"`
	Concurrency int           `toml:"concurrency"`
	MinLength   int           `toml:"min_length"`
	MaxLength   int           `toml:"max_length"`
}

// Bank2Case transfers money between accounts, creating new ledger transactions
// in the form of a transaction record and two transaction "legs" per database
// transaction.
type Bank2Case struct {
	cfg   *Bank2CaseConfig
	wg    sync.WaitGroup
	stop  int32
	txnID int32
}

// NewBank2Case returns a Bank2Case.
func NewBank2Case(cfg *Bank2CaseConfig) *Bank2Case {
	b := &Bank2Case{
		cfg: cfg,
	}
	return b
}

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

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

// Initialize creates the tables and inserts initial balances.
func (c *Bank2Case) Initialize(ctx context.Context, db *sql.DB) error {
	client.Infof("[%s] start to init...", c)
	defer func() {
		client.Infof("[%s] init end...", c)
	}()
	util.MustExec(db, `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)
);`)

	util.MustExec(db, `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)
);`)

	util.MustExec(db, `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)
);`)

	util.MustExec(db, "TRUNCATE TABLE bank2_accounts;")
	util.MustExec(db, "TRUNCATE TABLE bank2_transaction;")
	util.MustExec(db, "TRUNCATE TABLE bank2_transaction_leg;")

	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 := util.RunWithRetry(ctx, *retryLimit, 5*time.Second, func() error {
					_, err := db.Exec(query)
					if util.IsErrDupEntry(err) {
						return nil
					}
					return err
				})
				if err != nil {
					client.Fatalf("[%s] exec %s err %s", c, query, err)
				}
				client.Infof("[%s] insert %d accounts, takes %s", c, job.end-job.begin, time.Since(start))
			}
		}()
	}

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

	select {
	case <-ctx.Done():
		client.Warnf("[%s] initialize is cancelled", c)
		return nil
	default:
	}
	query := fmt.Sprintf(`INSERT IGNORE INTO bank2_accounts (id, balance, name) VALUES (%d, %d, "system account")`, systemAccountID, int64(c.cfg.NumAccounts*initialBalance))
	err := util.RunWithRetry(ctx, *retryLimit, 3*time.Second, func() error {
		_, err := db.Exec(query)
		if util.IsErrDupEntry(err) {
			return nil
		}
		return err
	})
	if err != nil {
		client.Fatalf("[%s] insert system account err: %v", c, err)
	}

	c.startVerify(ctx, db)
	return nil
}

func (c *Bank2Case) startVerify(ctx context.Context, db *sql.DB) {
	c.verify(db)

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case <-time.After(c.cfg.Interval):
				c.verify(db)
			}
		}
	}()
}

func (c *Bank2Case) verify(db *sql.DB) {
	start := time.Now()

	tx, err := db.Begin()
	if err != nil {
		bank2VerifyFailedCounter.Inc()
		return
	}
	defer tx.Rollback()

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

	util.RandomlyChangeReplicaRead("bank2", *replicaRead, db)

	uuid := fastuuid.MustNewGenerator().Hex128()
	query := fmt.Sprintf("SELECT SUM(balance) AS total, '%s' as uuid FROM bank2_accounts", uuid)
	var total int64
	err = tx.QueryRow(query).Scan(&total, &uuid)
	if err != nil {
		bank2VerifyFailedCounter.Inc()
		return
	}

	bank2VerifyDuration.Observe(time.Since(start).Seconds())

	expectTotal := (int64(c.cfg.NumAccounts) * initialBalance) * 2
	if total != expectTotal {
		client.Errorf("[%s] bank2_accounts total should be %d, but got %d, query uuid is %s", c, expectTotal, total, uuid)
		atomic.StoreInt32(&c.stop, 1)
		c.wg.Wait()
		client.Fatalf("[%s] bank2_accounts total should be %d, but got %d, query uuid is %s", c, expectTotal, total, uuid)
	}
}

// Execute implements Case Execute interface.
func (c *Bank2Case) Execute(ctx context.Context, db *sql.DB) error {
	client.Infof("[%s] start to test...", c)
	defer func() {
		client.Infof("[%s] test end...", c)
	}()
	var wg sync.WaitGroup
	for i := 0; i < c.cfg.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 {
					client.Errorf("[%s] stopped", c)
					return
				}
				c.wg.Add(1)
				c.moveMoney(db)
				c.wg.Done()
			}
		}(i)
	}
	wg.Wait()
	return nil
}

func (c *Bank2Case) moveMoney(db *sql.DB) {
	from, to := rand.Intn(c.cfg.NumAccounts), rand.Intn(c.cfg.NumAccounts)
	if c.cfg.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.cfg.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()
		client.Errorf("[%s] move money err %v", c, err)
		return
	}
	bank2VerifyDuration.Observe(time.Since(start).Seconds())
}

func (c *Bank2Case) execTransaction(db *sql.DB, from, to int, amount int) error {
	tx, err := db.Begin()
	if err != nil {
		return errors.Trace(err)
	}
	defer 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:
			client.Fatalf("[%s] got unexpected account %d", c, id)
		}
		count++
	}

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

	if count != 2 {
		client.Fatalf("[%s] select %d(%d) -> %d(%d) invalid count %d", c, from, fromBalance, to, toBalance, 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 {
		client.Infof("[%s] %d(%d) failed to transfer %d to %d(%d), err: %v", c, from, fromBalance, amount, to, toBalance, err)
	} else {
		client.Infof("[%s] %d(%d) transfer %d to %d(%d)", c, from, fromBalance, amount, to, toBalance)
	}
	return nil
}

func (c *Bank2Case) String() string {
	return "bank2"
}