package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"math"
	"os"
	"strings"
	"sync"
	"text/tabwriter"
	"time"

	"github.com/go-kit/kit/log/term"
	metrics "github.com/rcrowley/go-metrics"

	"github.com/tendermint/tendermint/libs/log"
	tmrpc "github.com/tendermint/tendermint/rpc/client"
)

var logger = log.NewNopLogger()

type statistics struct {
	TxsThroughput    metrics.Histogram `json:"txs_per_sec"`
	BlocksThroughput metrics.Histogram `json:"blocks_per_sec"`
}

func main() {
	var durationInt, txsRate, connections, txSize int
	var verbose bool
	var outputFormat, broadcastTxMethod string

	flagSet := flag.NewFlagSet("tm-bench", flag.ExitOnError)
	flagSet.IntVar(&connections, "c", 1, "Connections to keep open per endpoint")
	flagSet.IntVar(&durationInt, "T", 10, "Exit after the specified amount of time in seconds")
	flagSet.IntVar(&txsRate, "r", 1000, "Txs per second to send in a connection")
	flagSet.IntVar(&txSize, "s", 250, "The size of a transaction in bytes.")
	flagSet.StringVar(&outputFormat, "output-format", "plain", "Output format: plain or json")
	flagSet.StringVar(&broadcastTxMethod, "broadcast-tx-method", "async", "Broadcast method: async (no guarantees; fastest), sync (ensures tx is checked) or commit (ensures tx is checked and committed; slowest)")
	flagSet.BoolVar(&verbose, "v", false, "Verbose output")

	flagSet.Usage = func() {
		fmt.Println(`Tendermint blockchain benchmarking tool.

Usage:
	tm-bench [-c 1] [-T 10] [-r 1000] [-s 250] [endpoints] [-output-format <plain|json> [-broadcast-tx-method <async|sync|commit>]]

Examples:
	tm-bench localhost:26657`)
		fmt.Println("Flags:")
		flagSet.PrintDefaults()
	}

	flagSet.Parse(os.Args[1:])

	if flagSet.NArg() == 0 {
		flagSet.Usage()
		os.Exit(1)
	}

	if verbose {
		if outputFormat == "json" {
			fmt.Fprintln(os.Stderr, "Verbose mode not supported with json output.")
			os.Exit(1)
		}
		// Color errors red
		colorFn := func(keyvals ...interface{}) term.FgBgColor {
			for i := 1; i < len(keyvals); i += 2 {
				if _, ok := keyvals[i].(error); ok {
					return term.FgBgColor{Fg: term.White, Bg: term.Red}
				}
			}
			return term.FgBgColor{}
		}
		logger = log.NewTMLoggerWithColorFn(log.NewSyncWriter(os.Stdout), colorFn)

		fmt.Printf("Running %ds test @ %s\n", durationInt, flagSet.Arg(0))
	}

	if broadcastTxMethod != "async" &&
		broadcastTxMethod != "sync" &&
		broadcastTxMethod != "commit" {
		fmt.Fprintln(
			os.Stderr,
			"broadcast-tx-method should be either 'sync', 'async' or 'commit'.",
		)
		os.Exit(1)
	}

	var (
		endpoints     = strings.Split(flagSet.Arg(0), ",")
		client        = tmrpc.NewHTTP(endpoints[0], "/websocket")
		initialHeight = latestBlockHeight(client)
	)
	logger.Info("Latest block height", "h", initialHeight)

	transacters := startTransacters(
		endpoints,
		connections,
		txsRate,
		txSize,
		"broadcast_tx_"+broadcastTxMethod,
	)

	// Wait until transacters have begun until we get the start time
	timeStart := time.Now()
	logger.Info("Time last transacter started", "t", timeStart)

	duration := time.Duration(durationInt) * time.Second

	timeEnd := timeStart.Add(duration)
	logger.Info("End time for calculation", "t", timeEnd)

	<-time.After(duration)
	for i, t := range transacters {
		t.Stop()
		numCrashes := countCrashes(t.connsBroken)
		if numCrashes != 0 {
			fmt.Printf("%d connections crashed on transacter #%d\n", numCrashes, i)
		}
	}

	logger.Debug("Time all transacters stopped", "t", time.Now())

	stats, err := calculateStatistics(
		client,
		initialHeight,
		timeStart,
		timeEnd,
		durationInt,
	)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	printStatistics(stats, outputFormat)
}

func latestBlockHeight(client tmrpc.Client) int64 {
	status, err := client.Status()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	return status.SyncInfo.LatestBlockHeight
}

func countCrashes(crashes []bool) int {
	count := 0
	for i := 0; i < len(crashes); i++ {
		if crashes[i] {
			count++
		}
	}
	return count
}

// calculateStatistics calculates the tx / second, and blocks / second based
// off of the number the transactions and number of blocks that occurred from
// the start block, and the end time.
func calculateStatistics(
	client tmrpc.Client,
	minHeight int64,
	timeStart, timeStop time.Time,
	duration int,
) (*statistics, error) {
	stats := &statistics{
		BlocksThroughput: metrics.NewHistogram(metrics.NewUniformSample(1000)),
		TxsThroughput:    metrics.NewHistogram(metrics.NewUniformSample(1000)),
	}

	// get blocks between minHeight and last height
	// This returns max(minHeight,(last_height - 20)) to last_height
	info, err := client.BlockchainInfo(minHeight, 0)
	if err != nil {
		return nil, err
	}

	var (
		blockMetas = info.BlockMetas
		lastHeight = info.LastHeight
		diff       = lastHeight - minHeight
		offset     = len(blockMetas)
	)

	for offset < int(diff) {
		// get blocks between minHeight and last height
		info, err := client.BlockchainInfo(minHeight, lastHeight-int64(offset))
		if err != nil {
			return nil, err
		}
		blockMetas = append(blockMetas, info.BlockMetas...)
		offset = len(blockMetas)
	}

	var (
		numBlocksPerSec = make(map[int64]int64)
		numTxsPerSec    = make(map[int64]int64)
	)

	// because during some seconds blocks won't be created...
	for i := int64(0); i < int64(duration); i++ {
		numBlocksPerSec[i] = 0
		numTxsPerSec[i] = 0
	}

	// iterates from max height to min height
	for i, blockMeta := range blockMetas {
		// check if block was created after timeStart
		if blockMeta.Header.Time.Before(timeStart) {
			break
		}

		// check if block was created before timeStop
		if blockMeta.Header.Time.After(timeStop) {
			continue
		}
		sec := secondsSinceTimeStart(timeStart, blockMeta.Header.Time)

		// increase number of blocks for that second
		numBlocksPerSec[sec]++

		// increase number of txs for that second
		numTxsPerSec[sec] += blockMeta.Header.NumTxs
		logger.Debug(fmt.Sprintf("%d txs in block %d, height %d", blockMeta.Header.NumTxs, i, blockMeta.Header.Height))
	}

	for _, n := range numBlocksPerSec {
		stats.BlocksThroughput.Update(n)
	}

	for _, n := range numTxsPerSec {
		stats.TxsThroughput.Update(n)
	}

	return stats, nil
}

func secondsSinceTimeStart(timeStart, timePassed time.Time) int64 {
	return int64(math.Round(timePassed.Sub(timeStart).Seconds()))
}

func startTransacters(
	endpoints []string,
	connections,
	txsRate int,
	txSize int,
	broadcastTxMethod string,
) []*transacter {
	transacters := make([]*transacter, len(endpoints))

	wg := sync.WaitGroup{}
	wg.Add(len(endpoints))
	for i, e := range endpoints {
		t := newTransacter(e, connections, txsRate, txSize, broadcastTxMethod)
		t.SetLogger(logger)
		go func(i int) {
			defer wg.Done()
			if err := t.Start(); err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}
			transacters[i] = t
		}(i)
	}
	wg.Wait()

	return transacters
}

func printStatistics(stats *statistics, outputFormat string) {
	if outputFormat == "json" {
		result, err := json.Marshal(struct {
			TxsThroughput    float64 `json:"txs_per_sec_avg"`
			BlocksThroughput float64 `json:"blocks_per_sec_avg"`
		}{stats.TxsThroughput.Mean(), stats.BlocksThroughput.Mean()})
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		fmt.Println(string(result))
	} else {
		w := tabwriter.NewWriter(os.Stdout, 0, 0, 5, ' ', 0)
		fmt.Fprintln(w, "Stats\tAvg\tStdDev\tMax\tTotal\t")
		fmt.Fprintln(
			w,
			fmt.Sprintf(
				"Txs/sec\t%.0f\t%.0f\t%d\t%d\t",
				stats.TxsThroughput.Mean(),
				stats.TxsThroughput.StdDev(),
				stats.TxsThroughput.Max(),
				stats.TxsThroughput.Sum(),
			),
		)
		fmt.Fprintln(
			w,
			fmt.Sprintf("Blocks/sec\t%.3f\t%.3f\t%d\t%d\t",
				stats.BlocksThroughput.Mean(),
				stats.BlocksThroughput.StdDev(),
				stats.BlocksThroughput.Max(),
				stats.BlocksThroughput.Sum(),
			),
		)
		w.Flush()
	}
}