package main

import (
	"database/sql"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"sort"
	"syscall"
	"time"

	"github.com/pingcap/schrodinger-test/tpch-10/query"
	client "github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/schrodinger/pkg/types"
	"golang.org/x/net/context"
)

var (
	backupName string
	dbName     string
)

func init() {
	flag.StringVar(&backupName, "backup-name", "tpch-10g", "the name of backup uesed to restore cluster")
	flag.StringVar(&dbName, "db", "tpch", "the databases to use")
}

func main() {
	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.Infof("[TPCH-10] Got signal [%s] to exist.", sig)
		cancel()
		os.Exit(0)
	}()

	cli := client.New(ctx, client.ModeOnline)

	// Upgrade cluster to the latest master branch.
	if err := upgradeCluster(cli); err != nil {
		client.Fatalf("[TPCH-10] Upgrade cluster FAILED: %v", err)
	}

	if err := cli.RestoreCat(backupName); err != nil {
		client.Fatalf("[TPCH-10] restore cluster FAILED: %v", err)
	}

	// Sleep 10 minutes to wait for the TiDB cluster to be stable.
	time.Sleep(10 * time.Minute)

	// Connect to the cluster.
	db, err := connectCluster(cli)
	if err != nil {
		client.Fatalf("[TPCH-10] Connect to TiDB FAILED: %v", err)
	}
	defer db.Close()

	// Run all the TPC-H queries and report test result.
	caseResult, err := runQueries(db)
	if err != nil {
		client.Errorf("[TPCH-10] Run query failed: %v", err)
	}

	// Report test result.
	err = cli.ReportTestResult(caseResult)
	if err != nil {
		client.Errorf("[TPCH-10] Report case result %v, failed: %v", caseResult, err)
	}
}

func connectCluster(cli *client.Client) (*sql.DB, error) {
	dbAddr := cli.TiDBService()
	dbDSN := fmt.Sprintf("root:@tcp(%s)/%s", dbAddr, dbName)
	db, err := sql.Open("mysql", dbDSN)
	if err != nil {
		return nil, err
	}

	err = db.Ping()
	if err != nil {
		db.Close()
		return nil, err
	}

	return db, nil
}

func upgradeCluster(cli *client.Client) error {
	config := cli.CatConfig()
	newCat := &types.Cat{
		PDConfig:   config.PDConfig,
		TiDBConfig: config.TiDBConfig,
		TiKVConfig: config.TiKVConfig,
	}
	return cli.UpgradeCat(newCat)
}

func runQueries(db *sql.DB) (*types.TestResult, error) {
	caseResult := &types.TestResult{
		TestName:     "tpch_10g_sql_engine_team",
		QueryResults: []*types.QueryResult{},
	}

	for ordinal, query := range query.All {
		if ordinal+1 == 5 || ordinal+1 == 15 {
			continue
		}

		counter := 5
		costTimes := make([]time.Duration, 0, counter)
		for i := 0; i < counter; i++ {
			costTime, err := runQuery(db, query)
			if err != nil {
				return nil, err
			}
			costTimes = append(costTimes, costTime)
		}

		sort.Slice(costTimes, func(i, j int) bool { return costTimes[i] < costTimes[j] })
		costTimes = costTimes[0 : len(costTimes)-1] // remove the lagest.

		var total time.Duration
		for _, costTime := range costTimes {
			total += costTime
		}

		avgCostNs := float64(total.Nanoseconds()) / float64(len(costTimes))
		avgCostDuration := time.Duration(int64(avgCostNs)) * time.Nanosecond

		client.Infof("query: %v, avg time: %vs\n", ordinal+1, avgCostDuration.Seconds())
		caseResult.QueryResults = append(caseResult.QueryResults, &types.QueryResult{
			Query: fmt.Sprintf("query %v", ordinal+1),
			Cost:  avgCostDuration,
		})
	}
	return caseResult, nil
}

func runQuery(db *sql.DB, query string) (costTime time.Duration, err error) {
	startTime := time.Now()
	defer func() {
		costTime = time.Since(startTime)
	}()

	var rows *sql.Rows
	rows, err = db.Query(query)
	if err != nil {
		return
	}
	defer rows.Close()

	// pull result.
	for rows.Next() {
	}

	err = rows.Err()
	return
}