package main

import (
	"bytes"
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"log"
	"math/rand"
	"os/exec"
	"sync"
	"sync/atomic"
	"time"
)

var (
	oldVersion       = flag.String("o", "v1.0.3", "Old release version")
	newVersion       = flag.String("n", "v2.0.5", "Old release version")
	fromTidbCommitID = flag.String("oldtidb", "", "TiDB Commit ID")
	fromTikvCommitID = flag.String("oldtikv", "", "TiKV Commit ID")
	fromPdCommitID   = flag.String("oldpd", "", "PD Commit ID")
	toTidbCommitID   = flag.String("newtidb", "", "TiDB Commit ID")
	toTikvCommitID   = flag.String("newtikv", "", "TiKV Commit ID")
	toPdCommitID     = flag.String("newpd", "", "PD Commit ID")
	host             = flag.String("h", "127.0.0.1", "TiDB host address")
	pr               = flag.Bool("pr", false, "Wheather to use pr")
)

const (
	totalTable = 600
	port       = 4000
	user       = "root"
	password   = ""
)

func localPrepare(version string) {
	cmd := exec.Command("python", "change_inventory.py", version)
	cmd.Run()
	cmd = exec.Command("rm", "-rf", "downloads")
	cmd.Run()
	cmd = exec.Command("ansible-playbook", "local_prepare.yml")
	var stdout bytes.Buffer
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		log.Println(string(stdout.Bytes()))
		log.Fatal("Local prepare failed!")
	}
	log.Println("Local prepare version:" + version + " succsess")
}

func getTiDBbyCommitID(commitID string) {
	url := ""
	if *pr {
		url = "http://pingcap-dev.hk.ufileos.com/builds/pingcap/tidb/pr/" + commitID + "/centos7/tidb-server.tar.gz"
	} else {
		url = "http://pingcap-dev.hk.ufileos.com/builds/pingcap/tidb/" + commitID + "/centos7/tidb-server.tar.gz"
	}
	cmd := exec.Command("curl", "-SLO", url)
	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}
	cmd = exec.Command("tar", "zxvf", "tidb-server.tar.gz", "-C", "resources")
	if err := cmd.Run(); err != nil {
		log.Fatal("Get TiDB binary from PR failed !")
	}
}

func getTiKVbyCommitID(commitID string) {
	url := ""
	if *pr {
		url = "http://pingcap-dev.hk.ufileos.com/builds/pingcap/tikv/pr/" + commitID + "/unportable_centos7/tikv-server.tar.gz"
	} else {
		url = "http://pingcap-dev.hk.ufileos.com/builds/pingcap/tikv/" + commitID + "/unportable_centos7/tikv-server.tar.gz"
	}
	cmd := exec.Command("curl", "-SLO", url)
	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}
	cmd = exec.Command("tar", "zxvf", "tikv-server.tar.gz", "-C", "resources")
	if err := cmd.Run(); err != nil {
		log.Fatal("Get TiKV binary from PR failed !")
	}
}

func getPDbyCommitID(commitID string) {
	url := ""
	if *pr {
		url = "http://pingcap-dev.hk.ufileos.com/builds/pingcap/pd/pr" + commitID + "/centos7/pd-server.tar.gz"
	} else {
		url = "http://pingcap-dev.hk.ufileos.com/builds/pingcap/pd/" + commitID + "/centos7/pd-server.tar.gz"
	}
	cmd := exec.Command("curl", "-SLO", url)
	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}
	cmd = exec.Command("tar", "zxvf", "pd-server.tar.gz", "-C", "resources")
	if err := cmd.Run(); err != nil {
		log.Fatal("Get PD binary from PR failed !")
	}
}

func bootstrap() {
	cmd := exec.Command("ansible-playbook", "bootstrap.yml")
	var stdout bytes.Buffer
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		log.Println(string(stdout.Bytes()))
		log.Fatal("bootstrap failed")
	}
	log.Println("bootstrap succsess")
}

func deploy() {
	cmd := exec.Command("ansible-playbook", "deploy.yml")
	var stdout bytes.Buffer
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		log.Println(string(stdout.Bytes()))
		log.Fatal("deploy failed")
	}
	log.Println("deploy succsess")
}

func startServer() {
	cmd := exec.Command("ansible-playbook", "start.yml")
	var stdout bytes.Buffer
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		log.Println(string(stdout.Bytes()))
		log.Fatal("Start server failed")
	}
	log.Println("Start server succsess")
}

func rollingUpdate() {
	log.Println("Start rolling update")
	cmd := exec.Command("ansible-playbook", "rolling_update.yml")
	var stdout bytes.Buffer
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		log.Println(string(stdout.Bytes()))
		log.Fatal("Rolling update Failed, reading /data1/deploy1/log/ for detail")
	}
	log.Println("Rolling update succsess")
}

func getDB() *sql.DB {
	dbDSN := fmt.Sprintf("%s:%s@tcp(%s:%d)/test", user, password, *host, 4001)
	db, _ := sql.Open("mysql", dbDSN)
	err := db.Ping()
	if err == nil {
		return db
	} else {
		dbDSN := fmt.Sprintf("%s:%s@tcp(%s:%d)/test", user, password, *host, 4000)
		db, _ := sql.Open("mysql", dbDSN)
		return db
	}
}

func mustQuery(query string) *sql.Rows {
	db := getDB()
	retryCount := 0
	rows, err := db.Query(query)
	for err != nil {
		if retryCount >= 10 {
			log.Fatal("try mustQuery too many times, test failed")
		}
		db = getDB()
		rows, err = db.Query(query)
		if err == nil {
			db.Close()
			return rows
		}
		retryCount++
		time.Sleep(time.Duration(7) * time.Second)
	}
	return rows
}

func mustExec(query string, args ...interface{}) {
	retryCount := 0
	db := getDB()
	_, err := db.Exec(query, args...)
	for err != nil {
		if retryCount >= 10 {
			log.Println("A test case failed.")
			log.Println("The case is: " + query)
			log.Println(err)
			break
		}
		db = getDB()
		_, err = db.Exec(query, args...)
		if err == nil {
			db.Close()
			return
		}
		retryCount++
		time.Sleep(time.Duration(7) * time.Second)
	}
	db.Close()
}

func dateImport(ch chan int) {
	log.Println("Importing test data....")
	mustExec("create table test_delete (c1 int, c2 int, primary key(c1))")
	mustExec("create table test_mixed (c1 int, c2 int, primary key(c1))")
	log.Println("Importing test data succsess")

	for _, tbl := range parallel_ddl_tbls {
		dataImportForParallelDDL(tbl)
	}

	ch <- 1
}

func testDelete(ch chan int) {
	//run case
	workerNum := 5
	rowCount := 10000
	batch := rowCount / workerNum
	start := time.Now()
	var wg sync.WaitGroup
	wg.Add(workerNum)
	for i := 0; i < workerNum; i++ {
		go func(i int) {
			defer wg.Done()
			for j := 0; j < batch; j++ {
				k := batch*i + j
				mustExec(fmt.Sprintf("insert into test_delete values (%d, %d)", k, k))
				mustExec(fmt.Sprintf("delete from test_delete where c1 = %d", k))
			}
		}(i)
	}
	wg.Wait()
	end := time.Now()
	log.Printf("[TestSimpleDelete][Time Cost]%v\n", end.Sub(start))
	//check
	cnt := 0
	rows := mustQuery("select * from test_delete")
	defer rows.Close()
	for rows.Next() {
		cnt++
	}
	if cnt != 0 {
		log.Printf("[TestSimpleDelete] Failed!\n")
	} else {
		log.Printf("[TestSimpleDelete] Success!\n")
	}
	ch <- 1
}

func testMixed(ch chan int) {
	workerNum := 5
	rowCount := 10000
	batch := rowCount / workerNum

	start := time.Now()

	var wg sync.WaitGroup
	wg.Add(workerNum)
	for i := 0; i < workerNum; i++ {
		go func(i int) {
			defer wg.Done()

			for j := 0; j < batch; j++ {
				k := batch*i + j
				mustExec(fmt.Sprintf("insert into test_mixed values (%d, %d)", k, k))
			}
		}(i)
	}
	wg.Wait()

	end := time.Now()
	fmt.Printf("[TestSimpleMixed][Insert][Time Cost]%v\n", end.Sub(start))

	start = time.Now()

	rowID := int64(rowCount)
	defaultValue := int64(-1)

	wg.Add(workerNum)
	for i := 0; i < workerNum; i++ {
		go func() {
			defer wg.Done()

			for j := 0; j < batch; j++ {
				key := atomic.AddInt64(&rowID, 1)
				mustExec(fmt.Sprintf("insert into test_mixed values (%d, %d)", key, key))
				key = int64(rand.Intn(rowCount))
				mustExec(fmt.Sprintf("update test_mixed set c2 = %d where c1 = %d", defaultValue, key))
				key = int64(rand.Intn(rowCount))
				mustExec(fmt.Sprintf("delete from test_mixed where c1 = %d", key))
			}
		}()
	}
	wg.Wait()

	end = time.Now()
	fmt.Printf("[TestSimpleMixed][Mixed][Time Cost]%v\n", end.Sub(start))
	updateCount := int64(0)
	insertCount := int64(0)
	rows := mustQuery("select * from test_mixed")
	rowNum := 1
	defer rows.Close()
	for rows.Next() {
		var c1, c2 int64
		if err := rows.Scan(&c1, &c2); err != nil {
			log.Fatal(err)
		}
		if c1 == c2 {
			insertCount++
		} else if c2 == defaultValue && c1 < int64(rowCount) {
			updateCount++
		} else {
			log.Fatalf("[TestSimpleMixed fail]invalid row: %d", rowNum)
		}
		rowNum++
	}
	deleteCount := atomic.LoadInt64(&rowID) - insertCount - updateCount
	if insertCount > 0 && updateCount > 0 && deleteCount > 0 {
		log.Println("Mixed Test Success!")
	} else {
		log.Println("Mixed Test Failed!")
	}
	ch <- 1
}

func cleanUp() {
	log.Println("All test passed!")
	mustExec("drop table test_delete")
	mustExec("drop table test_mixed")
}

func main() {
	flag.Parse()
	ch := make(chan int, 10)
	localPrepare(*oldVersion)
	if *fromTidbCommitID != "" {
		getTiDBbyCommitID(*fromTidbCommitID)
	}
	if *fromTikvCommitID != "" {
		getTiKVbyCommitID(*fromTikvCommitID)
	}
	if *fromPdCommitID != "" {
		getPDbyCommitID(*fromPdCommitID)
	}
	bootstrap()
	deploy()
	startServer()
	go dateImport(ch)
	if *oldVersion != *newVersion {
		localPrepare(*newVersion)
	}
	<-ch
	if *toTidbCommitID != "" {
		getTiDBbyCommitID(*toTidbCommitID)
	}
	if *toTikvCommitID != "" {
		getTiKVbyCommitID(*toTikvCommitID)
	}
	if *toPdCommitID != "" {
		getPDbyCommitID(*toPdCommitID)
	}
	go testDelete(ch)
	go testMixed(ch)
	wg := sync.WaitGroup{}
	execParallelDDL(&wg, parallel_ddl_tbls...)
	rollingUpdate()
	for i := 0; i < 2; i++ {
		<-ch
	}
	wg.Wait()
	checkParallelDDLTables()
	cleanUp()
}

func batchInsertRows(db *sql.DB, tbl string, startKey, endKey int) error {
	sql := fmt.Sprintf("insert into %s values", tbl)
	for i := startKey; i < endKey; i++ {
		sql += fmt.Sprintf("(%d, '%d', '%d', %d, %d, %d)", i, i, i, i, i, i)
		if i != endKey-1 {
			sql += ","
		}
	}
	log.Printf("insert startKey %d, endKey %d", startKey, endKey)
	_, err := db.Exec(sql)
	if err != nil {
		log.Printf("sql: %v, err: %v", sql, err)
	}
	return err
}

func dataImportForParallelDDL(tbl string) {
	db := getDB()
	sql := fmt.Sprintf("create table if not exists %s(a int, b varchar(256), c varchar(256), d int, e int, f int)", tbl)
	_, err := db.Exec(sql)
	if err != nil {
		log.Fatalf("create table failed %v", err)
	}
	db.Close()

	totalCnt := 10000
	goCnt := 50
	step := totalCnt / goCnt
	batch := 100
	if step/batch < 1 {
		batch = 1
	}
	wg := sync.WaitGroup{}
	for i := 0; i < goCnt; i++ {
		sKey := i * step
		wg.Add(1)
		go func(num, startKey int) {
			defer wg.Done()
			db := getDB()
			defer db.Close()
			for j := 0; j < step/batch; j++ {
				err = batchInsertRows(db, tbl, startKey, startKey+batch)
				if err != nil {
					log.Printf("insert data err %v", err)
					return
				}
				startKey += batch

			}
			log.Printf("exec sql startKey %d, endKey %d end", startKey, startKey+step)
		}(i, sKey)
	}
	wg.Wait()
}

// buildDDL generates the statements of DDL.
func buildDDL(tbls ...string) [][]string {
	sqlsSli := make([][]string, 0, 5)
	for _, tbl := range tbls {
		sqls := make([]string, 0, 5)
		sql := fmt.Sprintf("alter table %s add index idx0(a)", tbl)
		sqls = append(sqls, sql)
		sql = fmt.Sprintf("alter table %s add index idx1(b)", tbl)
		sqls = append(sqls, sql)
		sql = fmt.Sprintf("alter table %s add index idx2(c)", tbl)
		sqls = append(sqls, sql)
		sql = fmt.Sprintf("alter table %s add index idx3(d)", tbl)
		sqls = append(sqls, sql)
		sql = fmt.Sprintf("alter table %s add index idx4(e)", tbl)
		sqls = append(sqls, sql)
		sqlsSli = append(sqlsSli, sqls)
		log.Printf("tbl %v, sqls %v", tbl, sqls)

		sqls1 := make([]string, 0, 5)
		sql = fmt.Sprintf("alter table %s add column col int", tbl)
		sqls1 = append(sqls1, sql)
		sql = fmt.Sprintf("alter table %s add column col1 int", tbl)
		sqls1 = append(sqls1, sql)
		sql = fmt.Sprintf("alter table %s add column col2 int", tbl)
		sqls1 = append(sqls1, sql)
		sql = fmt.Sprintf("alter table %s add column col3 int", tbl)
		sqls1 = append(sqls1, sql)
		sql = fmt.Sprintf("alter table %s drop column f", tbl)
		sqls1 = append(sqls1, sql)
		sqlsSli = append(sqlsSli, sqls1)
		log.Printf("tbl %v, sqls1 %v", tbl, sqls1)
	}
	log.Printf("sqls len %d ", len(sqlsSli))
	return sqlsSli
}

var parallel_ddl_tbls = []string{"parallel_ddl_tbl_0", "parallel_ddl_tbl_1", "parallel_ddl_tbl_2"}

func execParallelDDL(wg *sync.WaitGroup, tbls ...string) {
	sqlsSli := buildDDL(tbls...)
	for i, sqls := range sqlsSli {
		wg.Add(1)
		go func(num int, sqls []string) {
			defer wg.Done()
			db := getDB()
			defer db.Close()
			log.Printf("start run ddl no.%d", num)
			for _, sql := range sqls {
				_, err := db.Exec(sql)
				if err != nil {
					log.Printf("no.%d, sql: %s, err: %v", num, sql, err)
				} else {
					log.Printf("no.%d, sql: %s", num, sql)
				}
			}
		}(i, sqls)
	}
}

func checkParallelDDLTables() {
	db := getDB()
	defer db.Close()
	for _, tbl := range parallel_ddl_tbls {
		sql := fmt.Sprintf("admin check table %s", tbl)
		_, err := db.Exec(sql)
		if err != nil {
			log.Fatalf("admin check table failed %v", err)
		}
	}
}