package main

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

	_ "github.com/go-sql-driver/mysql"
)

const CONCURRENCY = 10

var PERIOD = 10 * time.Minute

func help() {
	fmt.Printf("usage: %s prepare\nand then\t%s run\n", os.Args[0], os.Args[0])
}

func main() {
	if len(os.Args) != 2 {
		help()
		return
	}
	switch os.Args[1] {
	case "prepare":
		prepare()
	case "run":
		test()
	default:
		help()
	}
}

func test() {
	db, err := sql.Open("mysql", "root@tcp(127.0.0.1:4000)/test")
	if err != nil {
		panic(err)
	}
	defer db.Close()

	taskCh := make(chan string, 10)
	var wg sync.WaitGroup
	wg.Add(CONCURRENCY)
	for i := 0; i < CONCURRENCY; i++ {
		go queryGoroutine(&wg, db, taskCh)
	}
	go killGoroutine(db)

	wg.Add(1)
	go workloadGen(&wg, taskCh)
	wg.Wait()
}

func queryGoroutine(wg *sync.WaitGroup, db *sql.DB, taskCh <-chan string) {
	defer wg.Done()

	for sql := range taskCh {
		_, err := db.Query(sql)
		if err != nil {
			if !strings.Contains(err.Error(), "Query execution was interrupted") {
				panic(err)
			}
		}
	}
}

func workloadGen(wg *sync.WaitGroup, taskCh chan<- string) {
	defer wg.Done()

	queries := []string{
		"select a from t1",
		"select * from t1 limit 30",
		"select * from t1 order by b",
		"select * from t1 where b < 100",
		"select count(1) from t1, t2 where t1.b = t2.b",
		"select * from t1 where b not in (select 1 from t2 where t1.b = t2.b)",
		"select sum(t1.a) from t1,t2 where t1.b != t2.b",
	}

	done := time.After(PERIOD)
	for {
		idx := rand.Intn(len(queries))
		query := queries[idx]
		select {
		case taskCh <- query:
		case <-done:
			close(taskCh)
			return
		}
	}
}

func killGoroutine(db *sql.DB) {
	for {
		showProcessAndKillQuery(db)
		time.Sleep(100 * time.Millisecond)
	}
}

func showProcessAndKillQuery(db *sql.DB) error {
	infos, err := showProcessList(db)
	if err != nil {
		return err
	}

	rand.Shuffle(len(infos), func(i, j int) { infos[i], infos[j] = infos[j], infos[i] })
	for _, pi := range infos {
		if pi.Command == "Query" {
			if pi.Info == "show processlist" {
				continue
			}

			if pi.Time > 60 {
				panic("query runs run more than 1min" + pi.Info)
			}
			killQuery(db, pi.ID)
		}
	}
	return nil
}

type ProcessInfo struct {
	ID      int
	User    string
	Host    string
	DB      string
	Command string
	Time    int
	State   int
	Info    string
}

func showProcessList(db *sql.DB) ([]ProcessInfo, error) {
	rows, err := db.Query("show processlist")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var infos []ProcessInfo
	for rows.Next() {
		var pi ProcessInfo
		if err := rows.Scan(&pi.ID, &pi.User, &pi.Host, &pi.DB, &pi.Command, &pi.Time, &pi.State, &pi.Info); err != nil {
			return nil, err
		}
		infos = append(infos, pi)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return infos, nil
}

func killQuery(db *sql.DB, id int) error {
	_, err := db.Exec(fmt.Sprintf("kill tidb query %d", id))
	return err
}