package main

import (
	"database/sql"
	"fmt"
	"os/exec"
	"regexp"
	"strings"
	"time"

	analyzer "github.com/MyonKeminta/sysbench-analyzer/lib"
	"github.com/juju/errors"
	client "github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/schrodinger/pkg/types"
)

var (
	sysbenchCMD = map[string]string{
		"prepare": "sysbench oltp_read_write --mysql-host=%s --mysql-port=%d --mysql-user=%s --mysql-db=%s " +
			"--tables=%d --table-size=%d --db-driver=mysql --threads=%d --report-interval=%d --percentile=%d prepare",
		"oltp": "sysbench oltp_read_write --mysql-host=%s --mysql-port=%d --mysql-user=%s --mysql-db=%s " +
			"--tables=%d --table-size=%d --db-driver=mysql --threads=%d --report-interval=%d --percentile=%d --time=%d " +
			"--rand-type=%s --rand-pareto-h=%f run",
		"select": "sysbench oltp_point_select --mysql-host=%s --mysql-port=%d --mysql-user=%s --mysql-db=%s " +
			"--tables=%d --table-size=%d --db-driver=mysql --threads=%d --report-interval=%d --percentile=%d --time=%d " +
			"--rand-type=%s --rand-pareto-h=%f run",
		"insert": "sysbench oltp_insert --mysql-host=%s --mysql-port=%d --mysql-user=%s --mysql-db=%s " +
			"--tables=%d --table-size=%d --db-driver=mysql --threads=%d --report-interval=%d --percentile=%d --time=%d " +
			"--rand-type=%s --rand-pareto-h=%f run",
		"update": "sysbench oltp_update_index --mysql-host=%s --mysql-port=%d --mysql-user=%s --mysql-db=%s " +
			"--tables=%d --table-size=%d --db-driver=mysql --threads=%d --report-interval=%d --percentile=%d --time=%d " +
			"--rand-type=%s --rand-pareto-h=%f run",
	}
)

// SysbenchConfig defines the configuration for the sysbench test.
type SysbenchConfig struct {
	Host           string        `toml:"host"`
	Port           int           `toml:"port"`
	User           string        `toml:"user"`
	Password       string        `toml:"password"`
	TableCount     int           `toml:"table_count"`
	TableSize      int           `toml:"table_size"`
	Threads        int           `toml:"threads"`
	MaxTime        int           `toml:"max_time"`
	ReportInterval int           `toml:"report_interval"`
	Percentile     int           `toml:"percentile"`
	DBName         string        `toml:"database"`
	Operation      string        `toml:"operation"`
	WarmupTime     time.Duration `toml:"warm_up_time"`
	Prepare        bool          `toml:"prepare"`
	WarmupData     bool          `toml:"warm_up_data"`
	Threshold      float64       `toml:"threshold"`
	RandType       string        `toml:"rand_type"`
	RandParetoH    float64       `toml:"rand_pareto_h"`
}

// BenchSysbench is a bench case for test sysbench
type BenchSysbench struct {
	cfg *SysbenchConfig
	db  *sql.DB
}

// NewBenchSysbench returns the benchSysbench
func NewBenchSysbench(cfg *SysbenchConfig, db *sql.DB) *BenchSysbench {
	return &BenchSysbench{
		cfg: cfg,
		db:  db,
	}
}

func (s *BenchSysbench) prepareData() error {
	if err := s.createDB(); err != nil {
		return errors.Trace(err)
	}

	cmdStr := fmt.Sprintf(sysbenchCMD["prepare"], s.cfg.Host, s.cfg.Port,
		s.cfg.User, s.cfg.DBName, s.cfg.TableCount, s.cfg.TableSize, s.cfg.Threads,
		s.cfg.ReportInterval, s.cfg.Percentile)
	client.Infof("[%s] start to prepare, exec [%s]", name(), cmdStr)
	if _, err := s.execute(cmdStr); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (s *BenchSysbench) run() (*types.TestResult, error) {
	if s.cfg.Prepare {
		client.Infof("[%s] start to prepare data...", name())
		if err := s.prepareData(); err != nil {
			return nil, errors.Trace(err)
		}
	}

	if s.cfg.WarmupData {
		client.Infof("[%s] start to warm up data", name())
		if err := s.warmupData(); err != nil {
			return nil, errors.Trace(err)
		}
	}

	client.Infof("[%s] warm up, sleep %s", name(), s.cfg.WarmupTime.String())
	time.Sleep(s.cfg.WarmupTime)

	cmdStr := fmt.Sprintf(sysbenchCMD[s.cfg.Operation], s.cfg.Host, s.cfg.Port,
		s.cfg.User, s.cfg.DBName, s.cfg.TableCount, s.cfg.TableSize, s.cfg.Threads,
		s.cfg.ReportInterval, s.cfg.Percentile, s.cfg.MaxTime, s.cfg.RandType, s.cfg.RandParetoH)

	client.Infof("[%s] start to bench, exec [%s]", name(), cmdStr)
	output, err := s.execute(cmdStr)
	if err != nil {
		return nil, errors.Trace(err)
	}

	client.Infof("sysbench output: \n %s", output)

	success, violations := analyzer.CheckQPSDropFromSysbenchOutputText(output, s.cfg.Threshold)
	if !success {
		return nil, errors.Errorf("check from results of sysbench test for QPS dropping, threshold %f%%, \n Abnormal Lines: \n %v", s.cfg.Threshold*100, violations)
	}

	total, err := extractTarget(output, "total")
	if err != nil {
		return nil, err
	}
	totalTime, err := extractTarget(output, "time elapsed")
	if err != nil {
		return nil, err
	}
	avgMs, err := extractTarget(output, "avg")
	if err != nil {
		return nil, err
	}
	percentile99Ms, err := extractTarget(output, "99th percentile")
	if err != nil {
		return nil, err
	}

	return &types.TestResult{
		TestName: name(),
		QPS:      total / totalTime,
		AvgMs:    avgMs,
		Pct99:    percentile99Ms,
	}, nil
}

func (s *BenchSysbench) createDB() error {
	if s.db == nil {
		client.Errorf("[%s] db is nil", name())
		return errors.New("db is nil")
	}

	stmt := fmt.Sprintf("create database IF NOT EXISTS %s", s.cfg.DBName)
	_, err := s.db.Exec(stmt)
	if err != nil {
		client.Errorf("[%s] create database, exec %s failed: %v", name(), stmt, err)
		return errors.Trace(err)
	}
	return nil
}

func (s *BenchSysbench) warmupData() error {
	if s.db == nil {
		client.Errorf("[%s] db is nil", name())
		return errors.New("db is nil")
	}

	for i := 1; i <= s.cfg.TableCount; i++ {
		stmt := fmt.Sprintf("SELECT COUNT(pad) FROM sbtest%d USE INDEX (k_%d);", i, i)
		client.Infof("[%s] exec [%s]", name(), stmt)
		if _, err := s.db.Exec(stmt); err != nil {
			client.Errorf("[%s] warm up data, exec [%s] failed: %v", name(), stmt, err)
			return errors.Trace(err)
		}
	}
	return nil
}

func (s *BenchSysbench) execute(cmdStr string) (string, error) {
	cmd := exec.Command("/bin/sh", "-c", cmdStr)
	buf, err := cmd.CombinedOutput()
	if err != nil {
		client.Errorf("[%s] execute %s: %s \n, output: %s", name(), cmdStr, err, string(buf))
	}
	return string(buf), errors.Trace(err)
}

func extractTarget(output, target string) (float64, error) {
	pattern := fmt.Sprintf("%s:[ ]+[0-9.]+", target)
	re := regexp.MustCompile(pattern)
	s := re.Find([]byte(output))
	if s == nil {
		return 0, fmt.Errorf("failed to extract %s", target)
	}

	seps := strings.Split(string(s), ":")
	if len(seps) != 2 {
		return 0, fmt.Errorf("failed to extract %s", target)
	}

	var res float64
	if n, err := fmt.Sscanf(seps[1], "%f", &res); n != 1 || err != nil {
		return 0, fmt.Errorf("failed to extract %s", target)
	}

	return res, nil
}