package main

import (
	"bytes"
	"context"
	"database/sql"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"syscall"
	"time"

	_ "github.com/go-sql-driver/mysql"
	client "github.com/pingcap/schrodinger/pkg/test_client"
)

type Op interface {
	Do(db *sql.DB) error
}

type Exec struct {
	SQL  string
	Args []interface{}
}

func (o *Exec) Do(db *sql.DB) error {
	_, err := db.Exec(o.SQL, o.Args...)
	return err
}

type Get struct {
	SQL  string
	Args []interface{}
	N    int
}

func (o *Get) Do(db *sql.DB) error {
	ptrs := make([][]byte, o.N)
	row := make([]interface{}, o.N)
	for i := 0; i < o.N; i++ {
		row[i] = &ptrs[i]
	}
	err := db.QueryRow(o.SQL, o.Args...).Scan(row...)
	if err == sql.ErrNoRows {
		return nil
	}
	return err
}

type Workload struct {
	ICols uint
	IIdxs uint
	SCols uint
	SIdxs uint
	SSize uint
}

func (w *Workload) String() string {
	return fmt.Sprintf("W(%d:%d:%d:%d:%d)", w.ICols, w.IIdxs, w.SCols, w.SIdxs, w.SSize)
}

func (w *Workload) TableName() string {
	return fmt.Sprintf("tbl_i%d_%d_s%d_%d", w.ICols, w.IIdxs, w.SCols, w.SIdxs)
}

func (w *Workload) Schema() string {
	var buf bytes.Buffer
	buf.WriteString("CREATE TABLE ")
	buf.WriteString(w.TableName())
	buf.WriteString(" (id INT UNSIGNED NOT NULL PRIMARY KEY")
	var idxs []string
	for i := uint(0); i < w.ICols; i++ {
		name := w.icol(int(i))
		buf.WriteString(", " + name + " INT")
		if i < w.IIdxs {
			idxs = append(idxs, ", KEY ("+name+")")
		}
	}
	size := strconv.Itoa(int(w.SSize))
	if w.SSize < 20 {
		size = "20"
	}
	for i := uint(0); i < w.SCols; i++ {
		name := w.scol(int(i))
		buf.WriteString(", " + name + " VARCHAR(" + size + ")")
		if i < w.SIdxs {
			idxs = append(idxs, ", KEY ("+name+")")
		}
	}
	for _, idx := range idxs {
		buf.WriteString(idx)
	}
	buf.WriteString(")")
	return buf.String()
}

func (w *Workload) Play(n uint, ins float64, get float64, set float64) <-chan Op {
	ch := make(chan Op)

	if n <= 0 || ins <= 0 || get < 0 || set < 0 {
		close(ch)
		return ch
	}
	p1 := get / (ins + get + set)
	p2 := (get + set) / (ins + get + set)

	go func() {
		i, k := uint(0), 0
		for {
			if i == n {
				break
			}
			r := rand.Float64()
			switch {
			case r < p1 && k > 0:
				ch <- w.get(rand.Intn(k))
			case r < p2 && r >= p1 && k > 0:
				ch <- w.set(rand.Intn(k))
			default:
				ch <- w.ins(k)
				k += 1
			}
			i += 1
		}
		close(ch)
	}()

	return ch
}

func (w *Workload) icol(i int) string {
	return "icol_" + strconv.Itoa(i)
}

func (w *Workload) scol(i int) string {
	return "scol_" + strconv.Itoa(i)
}

func (w *Workload) get(id int) Op {
	return &Get{
		SQL:  "SELECT * FROM " + w.TableName() + " WHERE id=?",
		Args: []interface{}{id},
		N:    1 + int(w.ICols+w.SCols),
	}
}

func (w *Workload) set(id int) Op {
	var buf bytes.Buffer
	buf.WriteString("UPDATE " + w.TableName() + " SET id=?")
	args := make([]interface{}, 0, w.ICols+w.SCols+1)
	args = append(args, id)
	if w.ICols > 0 {
		ik := rand.Intn(int(w.ICols)) + 1
		for i := 0; i < ik; i++ {
			buf.WriteString(", " + w.icol(i) + "=?")
			args = append(args, rand.Intn(1e8))
		}
	}
	if w.SCols > 0 {
		sk := rand.Intn(int(w.SCols)) + 1
		for i := 0; i < sk; i++ {
			buf.WriteString(", " + w.scol(i) + "=?")
			args = append(args, randStr(w.SSize))
		}
	}
	buf.WriteString(" WHERE id=?")
	args = append(args, id)
	return &Exec{SQL: buf.String(), Args: args}
}

func (w *Workload) ins(id int) Op {
	var buf bytes.Buffer
	buf.WriteString("INSERT INTO " + w.TableName() + " (id")
	args := make([]interface{}, 0, w.ICols+w.SCols+1)
	args = append(args, id)
	for i := uint(0); i < w.ICols; i++ {
		buf.WriteString(", " + w.icol(int(i)))
		args = append(args, rand.Intn(1e8))
	}
	for i := uint(0); i < w.SCols; i++ {
		buf.WriteString(", " + w.scol(int(i)))
		args = append(args, randStr(w.SSize))
	}
	buf.WriteString(") VALUES (?")
	for i := uint(0); i < w.ICols+w.SCols; i++ {
		buf.WriteString(", ?")
	}
	buf.WriteString(")")
	return &Exec{SQL: buf.String(), Args: args}
}

func randStr(size uint) string {
	var rs = []rune("0123456abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

	b := make([]rune, size)
	for i := range b {
		b[i] = rs[rand.Intn(len(rs))]
	}
	return string(b)
}

func run(db *sql.DB, w Workload, n uint, p uint, ins float64, get float64, set float64) {

	tag := "[widetbl][" + w.String() + "]"
	client.Infof(tag+" start: n=%d p=%d", n, p)

	_, err := db.Exec("DROP TABLE IF EXISTS " + w.TableName())
	if err != nil {
		client.Fatal(tag+" drop table:", err)
	}

	_, err = db.Exec(w.Schema())
	if err != nil {
		client.Fatal(tag+" create table:", err)
	}

	t0 := time.Now()

	var wg sync.WaitGroup
	ops := w.Play(n, ins, get, set)
	for i := 0; i < int(p); i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			for op := range ops {
				err = op.Do(db)
				if err != nil {
					client.Fatalf(tag+"[#%d] %+v: %+v", id, op, err)
				}
			}
		}(i)
	}
	wg.Wait()

	_, err = db.Exec("DROP TABLE IF EXISTS " + w.TableName())
	if err != nil {
		client.Error(tag+" drop table:", err)
	}

	client.Info(tag+" done in", time.Now().Sub(t0))

}

const (
	modeRandom  = "random"
	modeOneshot = "oneshot"
)

func main() {
	var opts struct {
		mode string
		dsn  string
		n    uint
		p    uint
		ins  float64
		get  float64
		set  float64
	}

	w := Workload{}

	flag.StringVar(&opts.mode, "mode", "random", "running mode")
	flag.StringVar(&opts.dsn, "dsn", "", "dsn to use for debug")
	flag.UintVar(&opts.n, "n", 100, "number of operations to run")
	flag.UintVar(&opts.p, "p", 1, "parallel level")
	flag.Float64Var(&opts.ins, "ins", .6, "insert proportion")
	flag.Float64Var(&opts.get, "get", .2, "select proportion")
	flag.Float64Var(&opts.set, "set", .2, "update proportion")
	flag.UintVar(&w.ICols, "icols", 100, "number of int cols")
	flag.UintVar(&w.SCols, "scols", 100, "number of varchar cols")
	flag.UintVar(&w.IIdxs, "iidxs", 0, "number of indexes on int cols")
	flag.UintVar(&w.SIdxs, "sidxs", 0, "number of indexes on varchar cols")
	flag.UintVar(&w.SSize, "ssize", 20, "size of varchar cols")
	flag.Parse()

	switch opts.mode {
	case modeOneshot, modeRandom:
	default:
		client.Fatal("[widetbl] invalid mode:", opts.mode)
	}
	if w.ICols+w.SCols >= 512 {
		client.Fatal("[widetbl] too many cols:", w.ICols+w.SCols+1)
	}

	ctx, cancel := context.WithCancel(context.Background())

	sc := make(chan os.Signal, 1)
	signal.Notify(sc, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
	go func() {
		sig := <-sc
		client.Info("[widetbl] terminate due to signal:", sig)
		cancel()
		os.Exit(1)
	}()

	if opts.p == 0 {
		opts.p = 1
	}
	if len(opts.dsn) == 0 {
		cli := client.New(ctx, client.ModeOnline)
		opts.dsn = fmt.Sprintf("root:@tcp(%s)/test", cli.TiDBService())
		client.Info("[widetbl] use dsn:", opts.dsn)
	}

	db, err := sql.Open("mysql", opts.dsn)
	if err != nil {
		client.Fatal("[widetbl] open database:", err)
	}

	switch opts.mode {
	case modeOneshot:
		run(db, w, opts.n, opts.p, opts.ins, opts.get, opts.set)
	case modeRandom:
		for {
			cw := Workload{}
			cw.SCols = uint(rand.Intn(int(w.SCols) + 1))
			cw.SIdxs = uint(rand.Intn(int(w.SIdxs) + 1))
			cw.ICols = uint(rand.Intn(int(w.ICols) + 1))
			cw.IIdxs = uint(rand.Intn(int(w.IIdxs) + 1))
			cw.SSize = uint(rand.Intn(int(w.SSize) + 1))
			n := rand.Intn(int(opts.n)) + 1
			p := rand.Intn(int(opts.p)) + 1
			run(db, cw, uint(n), uint(p), opts.ins, opts.get, opts.set)
		}
	}

}

func init() {
	rand.Seed(time.Now().Unix())
}