package main

import (
	"bytes"
	"context"
	"io"
	"os"
	"time"

	// Register workload
	_ "github.com/pingcap/go-ycsb/pkg/workload"
	// Register TiKV database
	_ "github.com/pingcap/go-ycsb/db/tikv"

	"github.com/juju/errors"
	"github.com/magiconair/properties"
	ycsbClient "github.com/pingcap/go-ycsb/pkg/client"
	"github.com/pingcap/go-ycsb/pkg/measurement"
	"github.com/pingcap/go-ycsb/pkg/prop"
	"github.com/pingcap/go-ycsb/pkg/ycsb"
	client "github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/schrodinger/pkg/types"
)

var (
	casesProperties = map[string]struct {
		Type             string
		ReadProportion   float64
		UpdateProportion float64
		InsertProportion float64
		ScanProportion   float64
	}{
		"raw_get":               {"raw", 1, 0, 0, 0},
		"raw_put":               {"raw", 0, 0, 1, 0},
		"txn_get":               {"txn", 1, 0, 0, 0},
		"txn_put":               {"txn", 0, 0, 1, 0},
		"coprocessor_scan":      {"coprocessor", 0, 0, 0, 1},
		"coprocessor_point_get": {"coprocessor", 1, 0, 0, 0},
	}
)

func getCaseProperties(cfg *YcsbConfig) (*properties.Properties, bool) {
	p := properties.NewProperties()
	p.SetValue("tikv.pd", cfg.PD)
	p.SetValue(prop.OperationCount, cfg.OperationCount)
	p.SetValue(prop.RecordCount, cfg.RecordCount)
	p.SetValue(prop.ThreadCount, cfg.Threads)
	p.SetValue(prop.WarmUpTime, cfg.WarmUpTime)
	p.SetValue(prop.RequestDistribution, cfg.RequestDistribution)

	// set workload
	cp, ok := casesProperties[cfg.CaseName]
	if !ok {
		return nil, false
	}
	p.SetValue("tikv.type", cp.Type)
	p.SetValue(prop.ReadProportion, cp.ReadProportion)
	p.SetValue(prop.UpdateProportion, cp.UpdateProportion)
	p.SetValue(prop.InsertProportion, cp.InsertProportion)
	p.SetValue(prop.ScanProportion, cp.ScanProportion)

	return p, true
}

// YcsbConfig defines the configuration for the ycsb test.
type YcsbConfig struct {
	Threads             int           `toml:"threads"`
	CaseName            string        `toml:"case_name"`
	WarmUpTime          time.Duration `toml:"warm_up_time"`
	Threshold           float64       `toml:"threshold"`
	OperationCount      int           `toml:"operation_count"`
	RecordCount         int           `toml:"record_count"`
	PD                  string        `toml:"pd_address"`
	RequestDistribution string        `toml:"request_distribution"`
}

// BenchYcsb is a bench case for test ycsb
type BenchYcsb struct {
	cfg      *YcsbConfig
	props    *properties.Properties
	workload ycsb.Workload
	db       ycsb.DB
}

// NewBenchYcsb returns the benchYcsb
func NewBenchYcsb(cfg *YcsbConfig) (*BenchYcsb, error) {
	props, ok := getCaseProperties(cfg)
	if !ok {
		return nil, errors.Errorf("case %s is not supported", cfg.CaseName)
	}

	workloadName := props.GetString(prop.Workload, "core")
	workloadCreator := ycsb.GetWorkloadCreator(workloadName)
	workload, err := workloadCreator.Create(props)
	if err != nil {
		return nil, err
	}

	dbName := props.GetString(prop.DB, "tikv")
	dbCreator := ycsb.GetDBCreator(dbName)
	db, err := dbCreator.Create(props)
	if err != nil {
		return nil, err
	}
	db = ycsbClient.DbWrapper{db}

	return &BenchYcsb{
		cfg,
		props,
		workload,
		db,
	}, nil
}

func (s *BenchYcsb) prepareData(ctx context.Context) error {
	s.props.SetValue(prop.DoTransactions, false)

	cli := ycsbClient.NewClient(s.props, s.workload, s.db)
	measurement.InitMeasure(s.props)

	cli.Run(ctx)

	return nil
}

func (s *BenchYcsb) run(ctx context.Context) (*types.TestResult, error) {
	s.props.SetValue(prop.DoTransactions, true)

	cli := ycsbClient.NewClient(s.props, s.workload, s.db)
	measurement.InitMeasure(s.props)

	// capture std output
	old := os.Stdout // keep backup of the real stdout
	r, w, _ := os.Pipe()
	os.Stdout = w

	outC := make(chan string)
	// copy the output in a separate goroutine so printing can't block indefinitely
	go func() {
		var buf bytes.Buffer
		io.Copy(&buf, r)
		outC <- buf.String()
	}()

	cli.Run(ctx)

	// back to normal state
	w.Close()
	os.Stdout = old // restoring the real stdout
	output := <-outC

	// TODO: use analyzer to check drop
	client.Infof("ycsb output: \n %s", output)

	for _, info := range measurement.Info() {
		return &types.TestResult{
			TestName: name(),
			QPS:      ConvertType(info.Get(measurement.QPS)),
			AvgMs:    ConvertType(info.Get(measurement.AVG)) / 1000.0,
			Pct99:    ConvertType(info.Get(measurement.PER99TH)) / 1000.0,
		}, nil
	}
	return nil, errors.Errorf("measurement is none")
}

func (s *BenchYcsb) close() {
	s.db.Close()
	s.workload.Close()
}

// ConvertType converts the value to float64 type.
func ConvertType(value interface{}) float64 {
	switch v := value.(type) {
	case int64:
		return float64(v)
	case float64:
		return v
	case int32:
		return float64(v)
	case int:
		return float64(v)
	default:
		return float64(0)
	}
}