package main

import (
	"context"
	"database/sql"
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/kayac/parallel-benchmark/benchmark"
	"github.com/shogo82148/mysql"
)

var mu sync.Mutex
var errCount map[string]int

type logger struct {
}

func (logger) Print(v ...interface{}) {
	s := fmt.Sprint(v...)
	mu.Lock()
	if errCount == nil {
		errCount = map[string]int{}
	}
	errCount[s] = errCount[s] + 1
	mu.Unlock()
}

var value atomic.Value

func main() {
	var p int
	flag.IntVar(&p, "p", 1, "parallelism")
	var d time.Duration
	flag.DurationVar(&d, "d", 10*time.Second, "duration")
	flag.Parse()

	mysql.SetLogger(logger{})

	user := os.Getenv("MYSQL_TEST_USER")
	pass := os.Getenv("MYSQL_TEST_PASS")
	addr := os.Getenv("MYSQL_TEST_ADDR")
	dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?timeout=30s&strict=true", user, pass, addr, "gotest")
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		panic(err)
	}
	db.SetMaxIdleConns(p * runtime.GOMAXPROCS(0))

	go canceler()
	time.Sleep(100 * time.Millisecond)

	benchmark.RunFunc(
		func() int {
			score := 1
			ctx := value.Load().(context.Context)
			if err := db.PingContext(ctx); err != nil {
				logger{}.Print("conn.PingContext: ", err)
				score = 0
			}
			if _, err := db.ExecContext(ctx, "DO ?", 1); err != nil {
				logger{}.Print("conn.ExecContext: ", err)
				score = 0
			}
			var v int
			row := db.QueryRowContext(ctx, "SELECT ?", 1)
			if err := row.Scan(&v); err != nil {
				logger{}.Print("conn.Scan: ", err)
				score = 0
			}

			stmt, err := db.PrepareContext(ctx, "DO ?")
			if err != nil {
				logger{}.Print("conn.PrepareContext: ", err)
				score = 0
			} else {
				defer stmt.Close()
				if _, err := stmt.ExecContext(ctx, 1); err != nil {
					logger{}.Print("stmt.ExecContext: ", err)
					score = 0
				}
			}

			stmt, err = db.PrepareContext(ctx, "SELECT ?")
			if err != nil {
				logger{}.Print("conn.PrepareContext: ", err)
				score = 0
			} else {
				defer stmt.Close()
				row := stmt.QueryRowContext(ctx, 1)
				if err := row.Scan(&v); err != nil {
					logger{}.Print("stmt.Scan: ", err)
					score = 0
				}
			}

			tx, err := db.BeginTx(ctx, nil)
			if err != nil {
				logger{}.Print("conn.BeginContext: ", err)
				score = 0
			} else {
				if _, err := tx.ExecContext(ctx, "DO ?", 1); err != nil {
					logger{}.Print("tx.ExecContext: ", err)
					score = 0
				}
				if err := tx.Commit(); err != nil {
					logger{}.Print("tx.Commit: ", err)
					score = 0
				}
			}

			return score
		},
		d,
		p,
	)

	for err, count := range errCount {
		log.Println(err, count)
	}
}

func canceler() {
	for {
		ctx, cancel := context.WithCancel(context.Background())
		value.Store(ctx)
		time.Sleep(100 * time.Millisecond)
		cancel()
	}
}