package large

import (
	"context"
	"database/sql"
	"fmt"
	"math/rand"
	"time"

	"github.com/pingcap/errors"
)

type InsertIntoSelect struct{}

func (_ InsertIntoSelect) LargeTxn(conn *sql.Conn) error {
	fmt.Println("run large txn InsertIntoSelect")
	_, err := conn.ExecContext(context.Background(), "insert into sbtest2 select * from sbtest1")
	if err != nil {
		return errors.WithStack(err)
	}
	fmt.Println("run large txn InsertIntoSelect finish")
	return nil
}

type TxnMode int

const (
	Optimistic TxnMode = iota
	Pessimistic
)

type LargeDML struct {
	PKRange int
}

// runInSession run the callback function inside a session.
// The transaction mode of the session could be optimistic or pessimistic.
func runInSession(db *sql.DB, txnMode TxnMode, callback func(*sql.Conn) error) error {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	conn, err := db.Conn(ctx)
	if err != nil {
		return errors.WithStack(err)
	}
	cancel()
	defer conn.Close()

	switch txnMode {
	case Optimistic:
		_, err = conn.ExecContext(context.Background(), "set @@tidb_txn_mode = 'optimistic'")
	case Pessimistic:
		_, err = conn.ExecContext(context.Background(), "set @@tidb_txn_mode = 'pessimistic'")
	}
	if err != nil {
		return errors.WithStack(err)
	}

	return callback(conn)
}

func (conf LargeDML) LargeTxn(conn *sql.Conn) error {
	fmt.Println("run large txn DML")
	tx, err := conn.BeginTx(context.Background(), nil)
	if err != nil {
		return errors.WithStack(err)
	}

	for i := 0; i < 50; i++ {
		_, err := tx.Exec("update sbtest2 set k = k + 1 where id in (?, ?, ?)", rand.Intn(conf.PKRange), rand.Intn(conf.PKRange), rand.Intn(conf.PKRange))
		if err != nil {
			return errors.WithStack(err)
		}
		_, err = tx.Exec("delete from sbtest1 limit 1")
		if err != nil {
			return errors.WithStack(err)
		}

		_, err = tx.Exec("update sbtest2 set c = ? where k = ? limit ?", randString(32), rand.Intn(512), conf.PKRange*512/100)
		if err != nil {
			return errors.WithStack(err)
		}

		_, err = tx.Exec("insert ignore into sbtest1 values (?, ?, ?, ?)", conf.PKRange, rand.Intn(1<<11), randString(32), randString(32))
		if err != nil {
			return errors.WithStack(err)
		}

		start := rand.Intn(conf.PKRange)
		end := start + conf.PKRange/200
		_, err = tx.Exec("update sbtest1 set c = ? where id > ? and id < ?", randString(32), start, end)
		if err != nil {
			return errors.WithStack(err)
		}

		fmt.Println("current executed = ", i)
	}
	err = tx.Commit()
	if err != nil {
		return errors.WithStack(err)
	}

	fmt.Println("run large txn DML finish")
	return nil
}

type LargeUpdate struct {
	Table string
}

func (conf LargeUpdate) LargeTxn(conn *sql.Conn) error {
	fmt.Println("run large txn UPDATE")
	_, err := conn.ExecContext(context.Background(), "update sbtest1 set k = k + 1")
	fmt.Println("run large txn UPDATE finish!")
	if err != nil {
		return errors.WithStack(err)
	}
	return nil
}

type LargeDelete struct{}

func (conf LargeDelete) LargeTxn(conn *sql.Conn) error {
	fmt.Println("run large txn DELETE")
	_, err := conn.ExecContext(context.Background(), "delete from sbtest1")
	fmt.Println("run large txn DELETE finish!")
	if err != nil {
		return errors.WithStack(err)
	}
	return nil
}

type LoadData string

func (path LoadData) LargeTxn(conn *sql.Conn) error {
	fmt.Println("run large txn LoadData")
	_, err := conn.ExecContext(context.Background(), fmt.Sprintf("load data local infile '%s' into table sbtest1", string(path)))
	return err
}