package main

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

	_ "github.com/go-sql-driver/mysql"
	"github.com/pingcap/errors"
	"github.com/pingcap/test-infra/caselib/ticdc/perf/workload/schema"
)

var (
	tableCount      int
	tableStartIndex int
	tps             int
	rowCount        int64
	thread          int

	dbHost     string
	dbPort     int
	dbUser     string
	dbPassword string
	dbName     string

	total      uint64
	totalError uint64

	workloadType string
)

const (
	bank     = "bank"
	sysbench = "sysbench"
)

func init() {
	flag.StringVar(&workloadType, "workload-type", "sysbench", "bank or sysbench")
	flag.IntVar(&tableCount, "table-count", 1, "table count of the workload")
	flag.IntVar(&tableStartIndex, "table-start-index", 0, "table start index, sbtest<index>")
	flag.IntVar(&tps, "tps", 1000, "tps of the workload")
	flag.IntVar(&thread, "thread", 0, "thread of the workload")
	flag.Int64Var(&rowCount, "row-count", 1000000, "row count of the workload")

	flag.StringVar(&dbHost, "database-host", "127.0.0.1", "database host")
	flag.StringVar(&dbUser, "database-user", "root", "database user")
	flag.StringVar(&dbPassword, "database-password", "", "database password")
	flag.StringVar(&dbName, "database-db-name", "test", "database db name")
	flag.IntVar(&dbPort, "database-port", 4000, "database port")
	flag.Parse()
}

func main() {
	rand.Seed(time.Now().UTC().UnixNano())
	threadCount := getThreadCount()
	fmt.Printf("use thread %d\n", threadCount)
	group := &sync.WaitGroup{}
	group.Add(threadCount)

	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Local", dbUser, dbPassword, dbHost, dbPort, dbName))
	if err != nil {
		panic(err)
	}
	if err = initTables(db); err != nil {
		panic(err)
	}

	rowPerThread := rowCount / int64(threadCount)
	tpsPerThread := int(math.Round(float64(tps) / float64(threadCount)))
	for i := 0; i < threadCount; i++ {
		go func() {
			defer group.Done()
			worker(db, rowPerThread, tpsPerThread)
		}()
	}

	go printTPS()
	group.Wait()
}

func getThreadCount() int {
	cpuNum := runtime.NumCPU()
	if thread > 0 {
		cpuNum = thread
	}
	if tps < cpuNum {
		cpuNum = tps
	}
	return cpuNum
}

// initTables create tables if not exists
func initTables(db *sql.DB) error {
	var sqlFunc func(int) string
	if workloadType == bank {
		sqlFunc = schema.GetBankCreateTableStatement
	} else {
		sqlFunc = schema.GetCreateTableStatement
	}

	for i := 0; i < tableCount; i++ {
		fmt.Printf("try to create table %d\n", i+tableStartIndex)
		if _, err := db.Exec(sqlFunc(i + tableStartIndex)); err != nil {
			return errors.Annotate(err, "create table failed")
		}
	}
	return nil
}

func worker(db *sql.DB, rowPerThread int64, tpsPerThread int) {
	count := rowPerThread
	var t = time.Tick(time.Second)
	printedError := false

	var sqlFunc func(int, int) string
	if workloadType == bank {
		sqlFunc = schema.BuildBankInsertSql
	} else {
		sqlFunc = schema.BuildInsertSql
	}

	for {
		select {
		case <-t:
			var tablesMap = make(map[int]int)
			for q := 0; q < tpsPerThread; q++ {
				n := rand.Int63()
				tableN := int(n)%tableCount + tableStartIndex
				tablesMap[tableN] = tablesMap[tableN] + 1
			}
			for tableN, trCount := range tablesMap {
				insertSql := sqlFunc(tableN, trCount)
				_, err := db.Exec(insertSql)
				if err != nil {
					if !printedError {
						fmt.Println(err)
						printedError = true
					}
					atomic.AddUint64(&totalError, uint64(trCount))
				}
				atomic.AddUint64(&total, uint64(trCount))
				count = count - int64(trCount)
				if count <= 0 {
					return
				}
			}
		}
	}
}

func printTPS() {
	t := time.Tick(time.Second * 5)
	old := uint64(0)
	oldErr := uint64(0)
	for {
		select {
		case <-t:
			temp := atomic.LoadUint64(&total)
			qps := (float64(temp) - float64(old)) / 5.0
			old = temp
			temp = atomic.LoadUint64(&totalError)
			fmt.Printf("total %d, total err %d, qps is %f, err qps %f\n", total, totalError, qps, (float64(temp)-float64(oldErr))/5.0)
			oldErr = temp
		}
	}
}