package main

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

	"github.com/rogpeppe/fastuuid"
	log "github.com/sirupsen/logrus"
)

const initBalance = 1000

// bankCase is for concurrent balance transfer.
type bankCase struct {
	cfg *config
}

// config is config for bank test.
type config struct {
	Accounts int
	Interval time.Duration
	Threads  int
}

// newBankCase returns the bankCase.
func newBankCase(cfg *config) *bankCase {
	return &bankCase{
		cfg: cfg,
	}
}

// initialize implements Case initialize interface.
func (c *bankCase) initialize(db *sql.DB, create, insert bool) error {
	log.Info("start to init...")
	defer func() {
		log.Info("init end...")
	}()
	if create {
		MustExec(db, "drop table if exists accounts, record")
		MustExec(db, `create table accounts (
			id BIGINT PRIMARY KEY,
			balance BIGINT NOT NULL,
			remark VARCHAR(128),
			tso BIGINT DEFAULT NULL,
			UNIQUE KEY(tso, id))`)
		MustExec(db, `create table record (
        from_id BIGINT NOT NULL,
        to_id BIGINT NOT NULL,
        from_balance BIGINT NOT NULL,
		from_balance_after BIGINT NOT NULL,
        to_balance BIGINT NOT NULL,
		to_balance_after BIGINT NOT NULL,
        amount BIGINT NOT NULL,
        tso BIGINT NOT NULL,
        PRIMARY KEY (tso, from_id, to_id))`)
	}

	if !insert {
		return nil
	}

	var wg sync.WaitGroup

	// TODO: fix the error is Accounts can't be divided by batchSize.
	// Insert batchSize values in one SQL.
	batchSize := 100
	jobCount := c.cfg.Accounts / batchSize

	maxLen := len(remark)
	ch := make(chan int, jobCount)
	for i := 0; i < jobCount; i++ {
		ch <- i * batchSize
	}

	for i := 0; i < c.cfg.Threads; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			args := make([]string, batchSize)
			for {
				startIndex, ok := <-ch
				if !ok {
					return
				}
				start := time.Now()
				for i := 0; i < batchSize; i++ {
					args[i] = fmt.Sprintf("(%d, %d, '%s')", startIndex+i, initBalance, remark[:rand.Intn(maxLen)])
				}
				MustExec(db, fmt.Sprintf("INSERT INTO accounts (id, balance, remark) VALUES %s", strings.Join(args, ",")))
				log.Debugf("insert %d accounts, takes %s", batchSize, time.Now().Sub(start))
			}
		}()
	}

	close(ch)
	wg.Wait()
	return nil
}

func (c *bankCase) startVerify(db *sql.DB) {
	run := func(f func()) {
		ticker := time.NewTicker(c.cfg.Interval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				f()
			}
		}
	}

	go run(func() {
		err := c.verify(db)
		if err != nil {
			log.Errorf("verify error in %s, error: %v", time.Now(), err)
		} else {
			log.Infof("verify success in %s", time.Now())
		}
	})
}

func (c *bankCase) verify(db *sql.DB) error {
	var total int

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

	defer func() {
		_ = tx.Rollback()
	}()

	uuid := fastuuid.MustNewGenerator().Hex128()
	err = tx.QueryRow(fmt.Sprintf("select sum(balance) as total, '%s' from accounts", uuid)).Scan(&total, &uuid)
	if err != nil {
		return err
	}
	check := c.cfg.Accounts * initBalance
	if total != check {
		log.Fatalf("accounts total must %d, but got %d, query uuid is %s", check, total, uuid)
	}
	return nil
}

func (c *bankCase) execute(db *sql.DB) {
	c.startVerify(db)
	log.Info("start to test...")
	defer func() {
		log.Info("test end...")
	}()
	var wg sync.WaitGroup

	run := func(f func()) {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				f()
			}
		}()
	}

	for i := 0; i < c.cfg.Threads; i++ {
		run(func() {
			c.moveMoney(db)
		})
	}
	wg.Wait()
}

func (c *bankCase) moveMoney(db *sql.DB) {
	var (
		from, to int
	)
	for {
		from, to = rand.Intn(c.cfg.Accounts), rand.Intn(c.cfg.Accounts)
		if from == to {
			continue
		}
		break
	}

	if err := c.execTransaction(db, from, to, rand.Intn(999)); err != nil {
		log.Errorf("move money error: %v", err)
		time.Sleep(5 * time.Second)
	}
}

func (c *bankCase) execTransaction(db *sql.DB, from, to int, amount int) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}

	defer func() {
		_ = tx.Rollback()
	}()

	rows, err := tx.Query(fmt.Sprintf("SELECT id, balance FROM accounts WHERE id IN (%d, %d) FOR UPDATE", from, to))
	if err != nil {
		return err
	}
	defer func() {
		_ = 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 err
		}
		switch id {
		case from:
			fromBalance = balance
		case to:
			toBalance = balance
		default:
			return errors.New(fmt.Sprintf("got unexpected account %d", id))
		}
		count++
	}

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

	if count != 2 {
		return errors.New(fmt.Sprintf("select %d(%d) -> %d(%d) invalid count %d", from, fromBalance, to, toBalance, count))
	}

	if fromBalance >= amount {
		tso := time.Now().UnixNano()
		updateSQL := fmt.Sprintf(`
UPDATE accounts
  SET balance = CASE id WHEN %d THEN %d WHEN %d THEN %d END, tso = %d
  WHERE id IN (%d, %d)
`, to, toBalance+amount, from, fromBalance-amount, tso, from, to)
		_, err = tx.Exec(updateSQL)
		if err != nil {
			return err
		}

		if _, err = tx.Exec(fmt.Sprintf(`
INSERT INTO record (from_id, to_id, from_balance, from_balance_after, to_balance, to_balance_after, amount, tso)
    VALUES (%d, %d, %d, %d, %d, %d, %d, %d)`,
			from, to, fromBalance, fromBalance-amount, toBalance, toBalance+amount, amount, tso)); err != nil {
			return err
		}
	}
	return tx.Commit()
}