package large

import (
	"database/sql"
	"flag"
	"fmt"
	"time"

	"github.com/VividCortex/gohistogram"
	_ "github.com/go-sql-driver/mysql"
	"github.com/pingcap/errors"
	"github.com/spf13/cobra"
)

func CmdRun(cmd *cobra.Command, args []string) {
	db, err := sql.Open("mysql", Cmd.Conn.getDSN())
	handleErr(err)
	db.SetMaxOpenConns(512)

	err = runTask(&Cmd.Run, db, "")
	handleErr(err)
}

func runTask(conf *RunConfig, db *sql.DB, metricURL string) error {
	workerCount := conf.Small.WorkerCount

	hist := gohistogram.NewHistogram(160)
	report := &Report{Hist: hist}
	statsWorker := newStatsWorker(workerCount)
	go statsWorker.Run(hist)

	if conf.Chaos != nil {
		conf.Chaos.Chaos()
	}

	// notify is used by large txn worker to tell others its finish.
	notify := make(chan struct{})
	if metricURL != "" {
		go memUseMonitor(notify, metricURL, report)
	}

	workers := make([]Worker, workerCount)
	for i := 0; i < workerCount; i++ {
		workers[i] = Worker{ID: i, Count: workerCount, dur: make([]float64, 0, 100)}
		go workers[i].run(conf.Small.TxnMode, conf.Small.Task, db, statsWorker.sampleCh, notify)
	}

	start := time.Now()
	err := largeTxn(&conf.Large, db, notify)
	report.Duration = time.Since(start)
	if err != nil {
		fmt.Println("run large txn fail:", err)
	}

	// The stats worker should wait all workers to finish.
	statsWorker.Wait()

	for i := 0; i < workerCount; i++ {
		report.Succ += workers[i].succ
		report.Fail += workers[i].fail
	}
	report.Report()
	return err
}

func (worker *Worker) run(txnMode TxnMode, task SmallTxner, db *sql.DB, sampleCh chan<- []float64, done <-chan struct{}) {
	err := runInSession(db, txnMode, func(conn *sql.Conn) error {
		for {
			select {
			case <-done:
				sampleCh <- worker.dur
				sampleCh <- nil
				return nil
			default:
			}

			start := time.Now()
			err := task.SmallTxn(worker, conn)
			elapse := time.Since(start)
			if err != nil {
				fmt.Println("run small txn error = ", err)
				worker.fail++
			} else {
				worker.succ++
			}
			if len(worker.dur) == cap(worker.dur) {
				sz := cap(worker.dur)
				sampleCh <- worker.dur
				worker.dur = make([]float64, 0, sz)
			}
			worker.dur = append(worker.dur, elapse.Seconds()*1000)
		}
	})

	if err != nil {
		fmt.Println(err)
	}
}

type BaseRunTask struct {
	SmallTxnMode TxnMode
	LargeUpdate
	SelectRandomPoints
	UpdateRandomPoints
}

func largeTxn(conf *LargeTxnConfig, db *sql.DB, notify chan struct{}) error {
	defer close(notify)

	// Run the large txn in a new session!
	// Set the session mode to optimistic or pessimistic.
	err := runInSession(db, conf.TxnMode, func(conn *sql.Conn) error {
		err := conf.Task.LargeTxn(conn)
		if err != nil {
			return errors.WithStack(err)
		}
		return nil
	})

	return err
}

func (t BaseRunTask) SmallTxn(worker *Worker, conn *sql.Conn) error {
	t.SelectRandomPoints.SmallTxn(worker, conn)
	t.UpdateRandomPoints.SmallTxn(worker, conn)
	return nil
}

func (t BaseRunTask) TxnMode() TxnMode {
	return t.SmallTxnMode
}

var action = flag.String("action", "", "prepare | run | clean")

func RunTest(conf *Config) {
	flag.Parse()

	db, err := sql.Open("mysql", conf.Conn.getDSN())
	handleErr(err)
	defer db.Close()

	if *action == "" || *action == "prepare" {
		err = prepareTask(conf.Prepare.WorkerCount, conf.Prepare.Task, db)
		handleErr(err)
	}

	var metricURL string
	if conf.Conn.Status != 0 {
		metricURL = fmt.Sprintf("http://%s:%d/metrics", conf.Conn.Host, conf.Conn.Status)
	}

	if *action == "" || *action == "run" {
		err = runTask(&conf.Run, db, metricURL)
		handleErr(err)
	}

	if *action == "" || *action == "clean" {
		err = cleanTask(conf.Clean, db)
		handleErr(err)
	}
}

func Run(conf *Config) {
	RunTest(conf)
}
